This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 7e49a148b [Rust] Add constraint-aware structural visit and walk (#187)
7e49a148b is described below
commit 7e49a148b8f7ce00c0e0d5629d1b4c4d46d1c3f9
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Nov 4 11:58:52 2026 -0500
[Rust] Add constraint-aware structural visit and walk (#187)
Rust constraint-aware visitors and walkers inspect the binding, range
and guard facts active at each TIRx node. A complete native-compatible
base supplies constraint state and sparse native default hooks to the
shared TVM-FFI visitor engine. Visit and walk reuse `IntoDispatcher` for
owned or borrowed dispatchers, typed callbacks and nested callback
tuples; the walk adapter adds context observation and pre/post control.
The actual `ConstraintContext` types mirror native parent and
captured-prefix semantics. Clones remain live, snapshots capture
prefixes, and typed analyzer synchronization consumes those contexts.
Interior mutability covers native updates without retaining references
across recursive descent. Complete private layouts require matching
native headers and builds.
The dependency is merged TVM-FFI upstream
`3e906bc685ea45624bc80e8d7b7d86497d492af3`, including shared dispatcher
conversion and the safe borrowed type-attribute lookup. The C++ context
constructor parameter is renamed to `parent_prefix_size` to match the
Rust `with_parent_prefix_size` API.
---
3rdparty/tvm-ffi | 2 +-
docs/ir_guides/working_with_ir_in_rust.rst | 25 +-
include/tvm/sym/constraint.h | 19 +-
rust/src/sym.rs | 1 +
rust/src/sym/analyzer.rs | 19 +-
rust/src/sym/constraint.rs | 174 +++++++--
rust/src/tirx.rs | 2 +
.../structural_visit_with_constraint_context.rs | 195 ++++++++++
.../structural_walk_with_constraint_context.rs | 162 ++++++++
rust/tests/sym/analyzer.rs | 2 +-
rust/tests/sym/main.rs | 7 +-
rust/tests/tirx/main.rs | 1 +
rust/tests/tirx/structural.rs | 419 +++++++++++++++++++++
13 files changed, 976 insertions(+), 52 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index afa8dce89..3e906bc68 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit afa8dce89705e844457899b42e3d0f69da8c5c03
+Subproject commit 3e906bc685ea45624bc80e8d7b7d86497d492af3
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index fd80dd7b8..d5d260d4d 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -661,11 +661,16 @@ replacement before any use of ``i``.
.. note::
- TIRx also has traversals that track the facts valid at each node, such
- as loop ranges and branch conditions. In Rust they are
- |StructuralVisitorWithConstraintContext|_ and
- |StructuralMutatorWithConstraintContext|_; ``constraints()`` returns the
- current facts.
+ |structural_visit_with_constraint_context|_ tracks facts such as loop
+ ranges and branch conditions through the same Rust visitor engine. Its
+ |StructuralVisitorWithConstraintContextView|_ supports explicit descent;
+ ``visitor.base().constraints()`` returns the ordinary owning
+ ``ConstraintContext`` handle to the current facts. Use ``snapshot()`` to
retain the visible prefix or
+ ``to_array()`` for an independent sequence of fact handles.
+ |structural_walk_with_constraint_context|_ supplies the same context to
+ pre- or post-order observers and handles descent through the visitor engine.
+ Typed contexts are neither ``Send`` nor ``Sync``; erased FFI handles retain
+ the ordinary contract against concurrent access.
With fields, typed callbacks, and the structural utilities, a Rust function
works with the same IR that Python builds and prints, so an analysis or a
@@ -739,10 +744,12 @@ Quick Reference
.. _StructuralMutatorView:
../reference/rust/generated/tvm_ffi/extra/structural_mutate/struct.StructuralMutatorView.html
.. |ObjectIdentity| replace:: ``ObjectIdentity``
.. _ObjectIdentity:
../reference/rust/generated/tvm_ffi/object/struct.ObjectIdentity.html
-.. |StructuralVisitorWithConstraintContext| replace::
``StructuralVisitorWithConstraintContext``
-.. _StructuralVisitorWithConstraintContext:
../reference/rust/generated/tvm_tirx/tirx/struct.StructuralVisitorWithConstraintContext.html
-.. |StructuralMutatorWithConstraintContext| replace::
``StructuralMutatorWithConstraintContext``
-.. _StructuralMutatorWithConstraintContext:
../reference/rust/generated/tvm_tirx/tirx/struct.StructuralMutatorWithConstraintContext.html
+.. |structural_visit_with_constraint_context| replace::
``structural_visit_with_constraint_context``
+.. _structural_visit_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_visit_with_constraint_context/fn.structural_visit_with_constraint_context.html
+.. |structural_walk_with_constraint_context| replace::
``structural_walk_with_constraint_context``
+.. _structural_walk_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_walk_with_constraint_context/fn.structural_walk_with_constraint_context.html
+.. |StructuralVisitorWithConstraintContextView| replace::
``StructuralVisitorWithConstraintContextView``
+.. _StructuralVisitorWithConstraintContextView:
../reference/rust/generated/tvm_tirx/tirx/structural_visit_with_constraint_context/type.StructuralVisitorWithConstraintContextView.html
.. |match_any!| replace:: ``match_any!``
.. _match_any!: ../reference/rust/generated/tvm_ffi/macro.match_any.html
.. |Var::new| replace:: ``Var::new``
diff --git a/include/tvm/sym/constraint.h b/include/tvm/sym/constraint.h
index 98257ab24..01b23e36e 100644
--- a/include/tvm/sym/constraint.h
+++ b/include/tvm/sym/constraint.h
@@ -195,14 +195,13 @@ class ConstraintContextObj : public ffi::Object {
/*!
* \brief Capture an inherited prefix and start an empty local sequence.
* \param parent Non-null parent context.
- * \param parent_local_prefix Number of parent-local entries to capture.
+ * \param parent_prefix_size Number of parent-local entries to capture.
* \throws ffi::Error If parent is null or the prefix is outside [0,
parent->LocalSize()].
*/
- ConstraintContextObj(ffi::ObjectPtr<ConstraintContextObj> parent, int64_t
parent_local_prefix)
- : parent_(std::move(parent)), parent_prefix_size_(parent_local_prefix) {
- TVM_FFI_CHECK(
- parent_ && 0 <= parent_local_prefix && parent_local_prefix <=
parent_->LocalSize(),
- ValueError)
+ ConstraintContextObj(ffi::ObjectPtr<ConstraintContextObj> parent, int64_t
parent_prefix_size)
+ : parent_(std::move(parent)), parent_prefix_size_(parent_prefix_size) {
+ TVM_FFI_CHECK(parent_ && 0 <= parent_prefix_size && parent_prefix_size <=
parent_->LocalSize(),
+ ValueError)
<< "Constraint context parent local prefix is out of bounds";
}
// Context nodes have stable identity; use Snapshot to capture their current
position.
@@ -309,16 +308,16 @@ class ConstraintContext : public ffi::ObjectRef {
/*! \brief Construct an empty root. */
ConstraintContext() :
ffi::ObjectRef(ffi::make_object<ConstraintContextObj>()) {}
/*!
- * \brief Construct a child capturing [0, parent_local_prefix) of its
parent's local entries.
+ * \brief Construct a child capturing [0, parent_prefix_size) of its
parent's local entries.
* Later parent appends do not enter this child's effective constraints.
* \param parent Context supplying inherited constraints.
- * \param parent_local_prefix Number of parent-local entries to capture,
between zero
+ * \param parent_prefix_size Number of parent-local entries to capture,
between zero
* and parent->LocalSize(), inclusive.
* \throws ffi::Error If the prefix is outside that range.
*/
- ConstraintContext(ConstraintContext parent, int64_t parent_local_prefix)
+ ConstraintContext(ConstraintContext parent, int64_t parent_prefix_size)
: ffi::ObjectRef(ffi::make_object<ConstraintContextObj>(
- ffi::GetObjectPtr<ConstraintContextObj>(parent.get()),
parent_local_prefix)) {}
+ ffi::GetObjectPtr<ConstraintContextObj>(parent.get()),
parent_prefix_size)) {}
/*! \brief Adopt a reference to a non-null context object. */
explicit ConstraintContext(ffi::ObjectPtr<ConstraintContextObj> node)
: ffi::ObjectRef(std::move(node)) {}
diff --git a/rust/src/sym.rs b/rust/src/sym.rs
index 80e0f31e3..5d9cede82 100644
--- a/rust/src/sym.rs
+++ b/rust/src/sym.rs
@@ -27,6 +27,7 @@
// An object registered under `sym` that no file has a block for yet is
generated here, to be moved
// to the file of the Python module that defines it.
// tvm-ffi-stubgen(prefix): sym
+// tvm-ffi-stubgen(skip): sym.ConstraintContext
pub mod analyzer;
pub mod constraint;
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index 936e71315..7c69e88cb 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -39,6 +39,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::Range
// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::sym::int_set::IntSet
+// tvm-ffi-stubgen(import-object): crate::sym::constraint::ConstraintContext
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
@@ -50,6 +51,7 @@ use crate::ir::expr::IntoExpr;
use crate::ir::expr::Range;
use crate::ir::expr::Var;
use crate::ir::support;
+use crate::sym::constraint::ConstraintContext;
use crate::sym::int_set::IntSet;
use std::ops::Deref;
use tvm_ffi::Error;
@@ -115,15 +117,26 @@ impl Analyzer {
/// Synchronize with a `sym.ConstraintContext` captured at its current
position.
///
- /// The context is an opaque FFI object until typed constraint bindings
are available.
/// Switching contexts removes facts that are no longer present. Mutating
the context
- /// after this call takes effect only after another synchronization.
- pub fn sync_constraint_context(&self, context:
&tvm_ffi::object::ObjectRef) -> Result<()> {
+ /// after this call takes effect only after another synchronization.
Constraint and
+ /// expression nodes are shared and must not be mutated through aliases.
Reentrant
+ /// synchronization is rejected; discard the analyzer if premise ingestion
fails.
+ pub fn sync_constraint_context(&self, context: &ConstraintContext) ->
Result<()> {
tvm_ffi::cached_global_func!("sym.AnalyzerSyncConstraintContext")
.call_packed(&[tvm_ffi::AnyView::from(self),
tvm_ffi::AnyView::from(context)])?;
Ok(())
}
+ /// Snapshot the accepted premises in an independently extensible context.
+ ///
+ /// Fact and expression nodes remain shared. Appending to this context
changes the
+ /// analyzer only after [`Self::sync_constraint_context`] is called again.
+ pub fn get_constraint_context(&self) -> Result<ConstraintContext> {
+ tvm_ffi::cached_global_func!("sym.AnalyzerGetConstraintContext")
+ .call_tuple((self,))?
+ .try_into()
+ }
+
/// The constant bounds of `expr` (`sym.AnalyzerConstIntBound`).
pub fn const_int_bound(&self, expr: impl IntoExpr) ->
Result<ConstIntBound> {
tvm_ffi::cached_global_func!("sym.AnalyzerConstIntBound")
diff --git a/rust/src/sym/constraint.rs b/rust/src/sym/constraint.rs
index fe3a49b83..b85615be8 100644
--- a/rust/src/sym/constraint.rs
+++ b/rust/src/sym/constraint.rs
@@ -26,6 +26,20 @@
//! boolean `condition` its collector vouches for. A [`ConstraintContext`]
holds the constraints in
//! effect at a point of a traversal, ancestors' first; the structural visitor
and mutator of
//! [`tirx::structural`](crate::tirx::structural) keep one for each value they
reach.
+//!
+//! The context representation mirrors `include/tvm/sym/constraint.h`,
including
+//! its private parent, captured prefix and append-only array. It requires
matching
+//! native headers/builds: stub generation cannot certify these private fields.
+//! Context handles use ordinary FFI ownership. Their native-mutated array
makes
+//! typed handles neither Send nor Sync. Erasing them into ObjectRef or a
packed Function
+//! does not enforce that confinement; callers retain the ordinary FFI
obligation
+//! to prevent concurrent context access. Constraint and IR nodes remain
shared.
+
+// tvm-ffi-stubgen(ty-map): sym.ConstraintContext ->
crate::sym::constraint::ConstraintContext
+
+use std::cell::UnsafeCell;
+use std::mem::{align_of, offset_of, size_of};
+use tvm_ffi::{AnyView, Array, Error, ObjectCore, Result, VALUE_ERROR};
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::Expr;
@@ -151,30 +165,6 @@ impl BindingConstraint {
tvm_ffi::impl_object_upcast!(BindingConstraint => Constraint);
// tvm-ffi-stubgen(end)
-// tvm-ffi-stubgen(begin): object/sym.ConstraintContext
-/// Opaque: bytes [24, 48) of [24, 48) are not accounted for by reflected
fields. Fields are read through the C ABI getters.
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "sym.ConstraintContext"]
-#[type_final]
-pub struct ConstraintContextObj {
- base: Object,
-}
-
-#[repr(C)]
-#[derive(tvm_ffi::derive::ObjectRef, Clone)]
-pub struct ConstraintContext {
- base: ObjectArc<ConstraintContextObj>,
-}
-
-impl Deref for ConstraintContext {
- type Target = ConstraintContextObj;
- fn deref(&self) -> &ConstraintContextObj {
- &self.base
- }
-}
-// tvm-ffi-stubgen(end)
-
// tvm-ffi-stubgen(begin): object/sym.PredicateConstraint
/// Complete: reflected fields fill [24, 40) exactly (alignment padding [33,
40)).
#[repr(C)]
@@ -284,3 +274,139 @@ impl RangeConstraint {
tvm_ffi::impl_object_upcast!(RangeConstraint => Constraint);
// tvm-ffi-stubgen(end)
+
+/// Native-compatible stable-identity constraint context.
+///
+/// Only the local array changes after construction. Its handle is replaced by
+/// native append, so shared context inspection must use interior mutability.
+/// Parent links and their captured local prefix are immutable and acyclic.
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "sym.ConstraintContext"]
+#[type_final]
+pub struct ConstraintContextObj {
+ base: Object,
+ parent: Option<ObjectArc<ConstraintContextObj>>,
+ parent_prefix_size: i64,
+ constraints: UnsafeCell<Array<Constraint>>,
+}
+
+/// An owning reference to a live context. Clones share identity and later
appends.
+/// Use [`ConstraintContextObj::snapshot`] to capture the current prefix.
+#[repr(C)]
+#[derive(tvm_ffi::derive::ObjectRef, Clone)]
+pub struct ConstraintContext {
+ base: ObjectArc<ConstraintContextObj>,
+}
+
+impl Deref for ConstraintContext {
+ type Target = ConstraintContextObj;
+ fn deref(&self) -> &Self::Target {
+ &self.base
+ }
+}
+
+impl ConstraintContext {
+ /// Construct an empty root using the complete native-compatible layout.
+ pub fn new() -> Result<Self> {
+ Self::allocate(None, 0)
+ }
+
+ /// Capture a prefix of the parent's local entries and start an empty
child.
+ /// Earlier ancestor prefixes are inherited unchanged; later appends are
excluded.
+ pub fn with_parent_prefix_size(parent: Self, parent_prefix_size: i64) ->
Result<Self> {
+ if parent_prefix_size < 0 || parent_prefix_size > parent.local_size() {
+ return Err(Error::new(
+ VALUE_ERROR,
+ "Constraint context parent local prefix is out of bounds",
+ "",
+ ));
+ }
+ Self::allocate(Some(parent.base), parent_prefix_size)
+ }
+
+ fn allocate(parent: Option<ObjectArc<ConstraintContextObj>>, prefix: i64)
-> Result<Self> {
+ Ok(Self {
+ base: ObjectArc::new(ConstraintContextObj {
+ base: Object::new(),
+ parent,
+ parent_prefix_size: prefix,
+ // Start with an empty Rust array. Native Append owns all
growth;
+ // never install a nonempty Rust allocation in this mutable
field.
+ constraints: UnsafeCell::new(Array::default()),
+ }),
+ })
+ }
+}
+
+impl ConstraintContextObj {
+ /// Append through the existing native operation, retaining its array COW
behavior.
+ /// A running native visitor reserves its active context for semantic
hooks;
+ /// that visitor contract forbids observer appends during traversal.
Standalone
+ /// contexts and [`Self::for_each`] permit append.
+ /// Premise nodes and referenced expressions must not be mutated through
aliases.
+ pub fn append(&self, constraint: &Constraint) -> Result<()> {
+ tvm_ffi::cached_global_func!("sym.ConstraintContextAppend")
+ .call_packed(&[AnyView::from(self), AnyView::from(constraint)])?;
+ Ok(())
+ }
+
+ /// Number of local facts, excluding captured ancestors.
+ pub fn local_size(&self) -> i64 {
+ // No callback or mutation occurs during this short borrow.
+ unsafe { (&*self.constraints.get()).len() as i64 }
+ }
+
+ /// Capture the visible prefix in an independently extensible child.
+ pub fn snapshot(&self) -> Result<ConstraintContext> {
+ let parent = AnyView::from(self)
+ .try_as::<ConstraintContext>()
+ .expect("context has its registered native type");
+ ConstraintContext::with_parent_prefix_size(parent, self.local_size())
+ }
+
+ fn local_array(&self) -> Array<Constraint> {
+ // Hold an owning clone before any callback; native append then uses
COW.
+ unsafe { (&*self.constraints.get()).clone() }
+ }
+
+ /// Visit a fixed inherited/local prefix. A callback may append without
entering
+ /// this iteration; each borrowed fact is valid only during its callback.
+ pub fn for_each(&self, mut callback: impl FnMut(&Constraint)) {
+ let retained = AnyView::from(self)
+ .try_as::<ConstraintContext>()
+ .expect("context has its registered native type");
+ retained.for_each_prefix(retained.local_size(), &mut callback);
+ }
+
+ fn for_each_prefix(&self, prefix: i64, callback: &mut impl
FnMut(&Constraint)) {
+ let local = self.local_array();
+ if let Some(parent) = &self.parent {
+ parent.for_each_prefix(self.parent_prefix_size, callback);
+ }
+ for constraint in local.iter().take(prefix as usize) {
+ callback(&constraint);
+ }
+ }
+
+ /// Return an independent sequence in inherited/local order. Later context
+ /// appends do not change it; fact nodes and expressions remain shared.
+ pub fn to_array(&self) -> Array<Constraint> {
+ if self.parent.is_none() {
+ return self.local_array();
+ }
+ let mut facts = Vec::new();
+ self.for_each(|fact| facts.push(fact.clone()));
+ Array::new(facts)
+ }
+}
+
+// These verify the Rust mirror against the matching x86_64 header layout, not
+// an arbitrary loaded binary's hidden fields. Stub generation cannot certify
them.
+const _: () = {
+ assert!(size_of::<ConstraintContextObj>() == 48);
+ assert!(align_of::<ConstraintContextObj>() == 8);
+ assert!(offset_of!(ConstraintContextObj, parent) == 24);
+ assert!(offset_of!(ConstraintContextObj, parent_prefix_size) == 32);
+ assert!(offset_of!(ConstraintContextObj, constraints) == 40);
+};
diff --git a/rust/src/tirx.rs b/rust/src/tirx.rs
index e1ddf66e4..6cadcebfc 100644
--- a/rust/src/tirx.rs
+++ b/rust/src/tirx.rs
@@ -36,6 +36,8 @@ pub mod layout;
pub mod op;
pub mod op_attr;
pub mod structural;
+pub mod structural_visit_with_constraint_context;
+pub mod structural_walk_with_constraint_context;
pub mod tensor;
pub mod thread_scope;
pub mod tile;
diff --git a/rust/src/tirx/structural_visit_with_constraint_context.rs
b/rust/src/tirx/structural_visit_with_constraint_context.rs
new file mode 100644
index 000000000..6bdb01fc5
--- /dev/null
+++ b/rust/src/tirx/structural_visit_with_constraint_context.rs
@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//! Structural visiting with the lexical constraint context.
+//!
+//! The shared Rust structural engine owns dispatch and recursion. This module
+//! supplies its complete native-compatible base and sparse constraint-aware
+//! default hooks; ordinary default traversal is reused when a hook is absent.
+//! Native semantic hooks create no analyzer and do not snapshot mutable IR
reads.
+//!
+//! Hook and global-function registration follows the ordinary visitor/binding
+//! trust contract: registered hooks must honor the synchronous ABI, may not
retain
+//! or concurrently use the visitor, and registration must not race with
traversal.
+
+use std::borrow::BorrowMut;
+use std::cell::UnsafeCell;
+
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
+use tvm_ffi::extra::structural_visit::StructuralVisitorBase;
+use tvm_ffi::reflection::TypeAttrColumn;
+use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex};
+use tvm_ffi::{
+ AnyView, Dispatcher, IntoDispatcher, ObjectCore, Result,
StructuralVisitEngine,
+ StructuralVisitorObj, StructuralVisitorView, VisitInterrupt,
+};
+
+use super::StructuralVisitorWithConstraintContextObj;
+use crate::sym::{ConstraintContext, ConstraintContextObj};
+
+/// Complete Rust representation of the native constraint-aware visitor base.
+///
+/// This is distinct from the generated opaque
+/// [`StructuralVisitorWithConstraintContextObj`]. The engine allocates this
+/// complete representation and installs its own callback; no generated opaque
+/// object is allocated. The base has no owning object conversion.
+///
+/// Native semantic hooks can replace the active context and update the private
+/// scope flag through a shared [`StructuralVisitorView::base`] reference.
Those
+/// fields use interior mutability; context access clones the current handle.
+#[repr(C)]
+pub struct StructuralVisitorWithConstraintContextBase {
+ base: StructuralVisitorObj,
+ constraints: UnsafeCell<ConstraintContext>,
+ needs_child_context: UnsafeCell<bool>,
+}
+
+impl StructuralVisitorWithConstraintContextBase {
+ /// Read the current live context. This ordinary owning handle may
+ /// outlive the visitor. Use [`ConstraintContextObj::snapshot`] to retain
the
+ /// current prefix or [`ConstraintContextObj::to_array`] for a detached
sequence.
+ pub fn constraints(&self) -> ConstraintContext {
+ // No reference to the replaceable field survives this short clone.
+ unsafe { (&*self.constraints.get()).clone() }
+ }
+
+ /// Borrow only while an observer has no access to this visitor's descent.
+ /// The caller must end the borrow before any operation can replace the
context.
+ #[inline]
+ pub(super) unsafe fn constraints_for_observer(&self) -> &ConstraintContext
{
+ &*self.constraints.get()
+ }
+}
+
+// SAFETY: repr(C) mirrors the complete base in the matching native header.
+// Stub generation cannot certify its private tail; the binding requires the
+// matching native build. All fields begin with the same state as the native
+// constructor. Native-mutated fields use UnsafeCell. The type identity is the
+// registered specialized type, while the shared engine installs its vtable and
+// owns the dispatcher. Default hooks receive that live runtime synchronously
and
+// return one owning None/interrupt/error token under the ordinary hook
contract.
+unsafe impl StructuralVisitorBase for
StructuralVisitorWithConstraintContextBase {
+ fn new() -> Result<Self> {
+ Ok(Self {
+ base: StructuralVisitorObj::default(),
+ constraints: UnsafeCell::new(ConstraintContext::new()?),
+ needs_child_context: UnsafeCell::new(false),
+ })
+ }
+
+ fn type_index() -> i32 {
+ StructuralVisitorWithConstraintContextObj::type_index()
+ }
+
+ #[inline(always)]
+ unsafe fn default_visit_raw(runtime: *mut Self, value: AnyView<'_>) ->
TVMFFIAny {
+ static HOOKS: ColdInitLazyLock<Option<TypeAttrColumn>> =
+ ColdInitLazyLock::new(||
TypeAttrColumn::new("__s_visit_with_constraint_context__"));
+ // Copy the borrowed token before reentry. Packed dispatch acquires an
+ // owning Function before invoking a callback that could register
hooks.
+ let hook = HOOKS
+ .and_then(|column| column.get_raw_ffi_any(value.type_index()))
+ .filter(|hook| hook.type_index != TVMFFITypeIndex::kTVMFFINone as
i32);
+ if let Some(hook) = hook {
+ StructuralVisitorObj::call_structural_visit_hook_raw(
+ runtime.cast(),
+ *value.as_raw_ffi_any(),
+ hook,
+ )
+ } else {
+ // Prefix-only default descent calls the installed vtable for each
+ // child; it must not project an ordinary engine's dispatcher
offset.
+ StructuralVisitorObj::default_visit_raw(runtime.cast(), value)
+ }
+ }
+}
+
+/// Scoped access to the shared Rust engine with a constraint-aware base.
+///
+/// Pass the handler's `self` to recursive operations so the engine publishes
the
+/// actual exclusive dispatcher reborrow. Inspect premises via `view.base()`.
+///
+/// ```compile_fail
+/// use
tvm_tirx::tirx::structural_visit_with_constraint_context::StructuralVisitorWithConstraintContextView;
+/// fn escape<'a, D>(v: StructuralVisitorWithConstraintContextView<'a, D>)
+/// -> StructuralVisitorWithConstraintContextView<'static, D> { v }
+/// ```
+///
+/// ```compile_fail
+/// use
tvm_tirx::tirx::structural_visit_with_constraint_context::StructuralVisitorWithConstraintContextView;
+/// use tvm_tirx::tvm_ffi::Any;
+/// fn retain<D>(v: &StructuralVisitorWithConstraintContextView<'_, D>) -> Any
{
+/// v.base().into()
+/// }
+/// ```
+pub type StructuralVisitorWithConstraintContextView<'a, D> =
+ StructuralVisitorView<'a, D, StructuralVisitorWithConstraintContextBase>;
+
+/// Visit with native binding, range and guard semantics through the shared
Rust
+/// engine. A matched dispatcher handler owns descent. Unmatched values use the
+/// base's specialized hook or ordinary structural fallback.
+///
+/// ```
+/// use tvm_tirx::tirx::structural_visit_with_constraint_context::{
+/// structural_visit_with_constraint_context as visit,
+/// StructuralVisitorWithConstraintContextView,
+/// };
+/// use tvm_tirx::tvm_ffi::{dispatch, AnyView, Result, VisitInterrupt};
+///
+/// #[derive(Default)]
+/// struct Count(usize);
+/// #[dispatch(visit, (AnyView<'_>, &mut
StructuralVisitorWithConstraintContextView<'_, Self>),
Result<Option<VisitInterrupt>>)]
+/// impl Count {
+/// fn visit_number(&mut self, _value: &i64, visitor: &mut
StructuralVisitorWithConstraintContextView<'_, Self>)
+/// -> Result<Option<VisitInterrupt>> {
+/// assert_eq!(visitor.base().constraints().local_size(), 0);
+/// self.0 += 1;
+/// Ok(None)
+/// }
+/// }
+/// let mut count = Count::default();
+/// visit(&7_i64, &mut count)?;
+/// assert_eq!(count.0, 1);
+/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// ```
+///
+/// Callbacks must return failures through `Result` and must not panic: a panic
+/// reaching the non-unwinding native callback boundary aborts the process.
+/// Hooks and reflected metadata retain the ordinary visitor's trust
obligations.
+pub fn structural_visit_with_constraint_context<'input, R, H, M, D, O>(
+ root: R,
+ visitor: H,
+) -> Result<Option<VisitInterrupt>>
+where
+ R: Into<AnyView<'input>>,
+ O: BorrowMut<D>,
+ for<'v, 'c, 'r> H: IntoDispatcher<
+ (AnyView<'v>, &'c mut StructuralVisitorWithConstraintContextView<'r,
D>),
+ Option<Result<Option<VisitInterrupt>>>,
+ M,
+ Target = D,
+ Output = O,
+ >,
+ for<'s, 'v, 'c, 'r> &'s mut D: Dispatcher<
+ (AnyView<'v>, &'c mut StructuralVisitorWithConstraintContextView<'r,
D>),
+ Option<Result<Option<VisitInterrupt>>>,
+ >,
+{
+
StructuralVisitEngine::<StructuralVisitorWithConstraintContextBase>::visit_root(root,
visitor)
+}
diff --git a/rust/src/tirx/structural_walk_with_constraint_context.rs
b/rust/src/tirx/structural_walk_with_constraint_context.rs
new file mode 100644
index 000000000..e265196f3
--- /dev/null
+++ b/rust/src/tirx/structural_walk_with_constraint_context.rs
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//! Constraint-aware observation through the shared Rust visitor engine.
+//!
+//! This adapter chooses pre/post callback order and translates walk control;
+//! the visit engine and native constraint hooks own descent, scopes and
reborrows.
+
+use std::borrow::BorrowMut;
+use tvm_ffi::{
+ Any, AnyView, DefRegionKind, Dispatcher, IntoDispatcher, Result,
VisitInterrupt, WalkOrder,
+ WalkResult,
+};
+
+use crate::sym::ConstraintContext;
+use super::structural_visit_with_constraint_context::{
+ structural_visit_with_constraint_context,
StructuralVisitorWithConstraintContextView,
+};
+
+struct WalkVisitor<'a, D, const PRE: bool> {
+ observer: &'a mut D,
+}
+
+impl<'v, 'c, 'r, 'a, D, const PRE: bool>
+ Dispatcher<
+ (
+ AnyView<'v>,
+ &'c mut StructuralVisitorWithConstraintContextView<'r,
WalkVisitor<'a, D, PRE>>,
+ ),
+ Option<Result<Option<VisitInterrupt>>>,
+ > for &mut WalkVisitor<'a, D, PRE>
+where
+ for<'s, 'b, 'ctx> &'s mut D: Dispatcher<
+ (AnyView<'b>, &'ctx ConstraintContext, DefRegionKind),
+ Option<Result<WalkResult>>,
+ >,
+{
+ #[inline(always)]
+ fn dispatch_tuple(
+ self,
+ (value, visitor): (
+ AnyView<'v>,
+ &'c mut StructuralVisitorWithConstraintContextView<'r,
WalkVisitor<'a, D, PRE>>,
+ ),
+ ) -> Option<Result<Option<VisitInterrupt>>> {
+ Some(self.visit(value, visitor))
+ }
+}
+
+impl<D, const PRE: bool> WalkVisitor<'_, D, PRE>
+where
+ for<'s, 'v, 'c> &'s mut D:
+ Dispatcher<(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
Option<Result<WalkResult>>>,
+{
+ #[inline(always)]
+ fn visit(
+ &mut self,
+ value: AnyView<'_>,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let kind = visitor.def_region_kind();
+ if !PRE {
+ if let Some(interrupt) = visitor.default_visit(self, value)? {
+ core::hint::cold_path();
+ return Ok(Some(interrupt));
+ }
+ }
+ // Post callbacks observe the restored context. The observer cannot
access
+ // this runtime's visitor; a nested walk creates its own runtime.
Appending
+ // mutates context contents, not this handle. End the field borrow
before
+ // any default descent can replace it; observers may clone it to
retain it.
+ let outcome = unsafe {
+ (&mut *self.observer).dispatch_tuple((
+ value,
+ visitor.base().constraints_for_observer(),
+ kind,
+ ))
+ };
+ let outcome = match outcome {
+ Some(result) => result?,
+ None => WalkResult::Advance,
+ };
+ match outcome {
+ WalkResult::Advance if PRE => visitor.default_visit(self, value),
+ WalkResult::Advance => Ok(None),
+ WalkResult::Skip => {
+ core::hint::cold_path();
+ Ok(None)
+ }
+ WalkResult::Interrupt => {
+ core::hint::cold_path();
+ Ok(Some(VisitInterrupt::with(Any::new())))
+ }
+ WalkResult::InterruptWith(payload) => {
+ core::hint::cold_path();
+ Ok(Some(VisitInterrupt::with(payload)))
+ }
+ }
+ }
+}
+
+/// Walk every occurrence with its active constraint context and optional
definition
+/// region. Pre-order Skip prunes descent; post-order Skip completes the
callback.
+/// Interrupt and errors stop traversal. The native hooks own fact timing and
scope
+/// restoration; sequence/lazy-call hooks omit intermediate Array callbacks,
as in C++.
+///
+/// As specified by the native visitor, observers must not append to its active
+/// context during traversal. Standalone contexts and `for_each` allow append.
+/// Observers must not panic across native callbacks.
+/// Cloned contexts remain live; capture a snapshot or array for a fixed
prefix.
+/// Accepts the shared [`IntoDispatcher`] forms: owned or borrowed dispatchers,
+/// typed value or node callbacks, and nested callback tuples. Callbacks
receive
+/// a leading prefix of `(value, &ConstraintContext, DefRegionKind)` and
return an
+/// `IntoWalkResult`. Named dispatchers use that full argument tuple and return
+/// `Option<Result<WalkResult>>`; unmatched values advance normally.
+/// The local adapter adds context observation to the shared visitor engine.
+pub fn structural_walk_with_constraint_context<'input, R, H, M, D, O>(
+ root: R,
+ walker: H,
+ order: WalkOrder,
+) -> Result<Option<VisitInterrupt>>
+where
+ R: Into<AnyView<'input>>,
+ O: BorrowMut<D>,
+ for<'v, 'c> H: IntoDispatcher<
+ (AnyView<'v>, &'c ConstraintContext, DefRegionKind),
+ Option<Result<WalkResult>>,
+ M,
+ Target = D,
+ Output = O,
+ >,
+ for<'s, 'v, 'c> &'s mut D:
+ Dispatcher<(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
Option<Result<WalkResult>>>,
+{
+ let mut storage = walker.into_dispatcher();
+ let observer = storage.borrow_mut();
+ let root = root.into();
+ match order {
+ WalkOrder::PreOrder => {
+ structural_visit_with_constraint_context(root, WalkVisitor::<D,
true> { observer })
+ }
+ WalkOrder::PostOrder => {
+ structural_visit_with_constraint_context(root, WalkVisitor::<D,
false> { observer })
+ }
+ }
+}
diff --git a/rust/tests/sym/analyzer.rs b/rust/tests/sym/analyzer.rs
index 686fdedc6..9ac71542e 100644
--- a/rust/tests/sym/analyzer.rs
+++ b/rust/tests/sym/analyzer.rs
@@ -117,7 +117,7 @@ fn contexts_restore_parent_facts() {
let analyzer = Analyzer::new();
let x = var("x");
let x_mod_3 = op("prim._OpFloorMod", &x, int32(3));
- let parent = context([]);
+ let parent = analyzer.get_constraint_context().unwrap();
let child = context([predicate_constraint(op("prim._OpEQ", &x_mod_3,
int32(0)))]);
analyzer.sync_constraint_context(&child).unwrap();
assert_eq!(analyzer.modular_set(&x).unwrap().coeff, 3);
diff --git a/rust/tests/sym/main.rs b/rust/tests/sym/main.rs
index 0b26f5d06..7611356c2 100644
--- a/rust/tests/sym/main.rs
+++ b/rust/tests/sym/main.rs
@@ -28,7 +28,7 @@ mod int_set;
mod iter_affine_map;
use tvm_tirx::ir::{UNKNOWN_LOC, dtype, Expr, IntImm, Location, PrimType,
Range, Var};
-use tvm_tirx::sym::Analyzer;
+use tvm_tirx::sym::{Analyzer, ConstraintContext};
use tvm_tirx::tvm_ffi::String as FfiString;
use tvm_tirx::tvm_ffi::object::ObjectRef;
@@ -54,9 +54,8 @@ fn range(begin: i64, extent: i64) -> Range {
Range::from_fields(int32(begin), int32(extent), Location::unknown())
}
-// Constraint data is passed opaquely until the Rust bindings expose typed
contexts.
-fn context(constraints: impl IntoIterator<Item = ObjectRef>) -> ObjectRef {
- let context: ObjectRef = global("sym.ConstraintContext")
+fn context(constraints: impl IntoIterator<Item = ObjectRef>) ->
ConstraintContext {
+ let context: ConstraintContext = global("sym.ConstraintContext")
.call_tuple((None::<ObjectRef>, 0_i64))
.unwrap()
.try_into()
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index 7a58fe011..67d2ad347 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -28,6 +28,7 @@ mod function;
mod index_map;
mod layout;
mod op;
+mod structural;
mod tensor;
mod thread_scope;
mod tile;
diff --git a/rust/tests/tirx/structural.rs b/rust/tests/tirx/structural.rs
new file mode 100644
index 000000000..554e5836a
--- /dev/null
+++ b/rust/tests/tirx/structural.rs
@@ -0,0 +1,419 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+//! Constraint-aware Rust-engine descent through an owned or borrowed
dispatcher.
+
+use tvm_tirx::ir::*;
+use tvm_tirx::sym::{
+ Analyzer, BindingConstraintObj, Constraint, ConstraintContext,
PredicateConstraintObj,
+ RangeConstraintObj,
+};
+use tvm_tirx::tirx::structural_visit_with_constraint_context::{
+ structural_visit_with_constraint_context as visit,
StructuralVisitorWithConstraintContextView,
+};
+use
tvm_tirx::tirx::structural_walk_with_constraint_context::structural_walk_with_constraint_context
as walk;
+use tvm_tirx::tvm_ffi::{
+ dispatch, AnyView, Array, DefRegionKind, Error, Map, ObjectRefCore,
Result, VisitInterrupt,
+ WalkOrder, WalkResult, RUNTIME_ERROR,
+};
+
+use super::int32;
+use crate::common::serial;
+
+#[derive(Default)]
+struct Observe {
+ facts: Vec<(i64, Vec<Constraint>)>,
+ definitions: Vec<(String, DefRegionKind)>,
+ live: Option<ConstraintContext>,
+ snapshot: Option<ConstraintContext>,
+ array: Option<Array<Constraint>>,
+}
+
+#[dispatch(visit, (AnyView<'_>, &mut
StructuralVisitorWithConstraintContextView<'_, Self>),
Result<Option<VisitInterrupt>>)]
+impl Observe {
+ fn visit_evaluate(
+ &mut self,
+ node: &EvaluateObj,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ if let Some(value) = node.value.downcast_ref::<IntImmObj>() {
+ let context = visitor.base().constraints();
+ if self.live.is_none() {
+ self.snapshot = Some(context.snapshot()?);
+ self.array = Some(context.to_array());
+ self.live = Some(context.clone());
+ }
+ self.facts.push((value.value.to_i64().unwrap(),
context.to_array().iter().collect()));
+ }
+ visitor.default_visit(self, node)
+ }
+ fn visit_var(
+ &mut self,
+ node: &VarObj,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ self.definitions.push((node.name.as_str().into(),
visitor.def_region_kind()));
+ visitor.default_visit(self, node)
+ }
+}
+
+fn eval(value: i64) -> Stmt {
+ Evaluate::new(int32(value), UNKNOWN_LOC).into()
+}
+fn variable(name: &str) -> Var {
+ Var::new(name, PrimType::int(32), UNKNOWN_LOC)
+}
+fn fact_counts(seen: &Observe) -> Vec<(i64, usize)> {
+ seen.facts.iter().map(|(n, facts)| (*n, facts.len())).collect()
+}
+#[test]
+fn typed_walk_callbacks_collect_facts_and_restore_analyzer_scopes() {
+ let _serial = serial();
+ let (x, i) = (variable("typed_x"), variable("typed_i"));
+ let bound = int32(4);
+ let condition = Var::new("typed_condition", PrimType::bool(), UNKNOWN_LOC);
+ let body =
+ For::new(&i, int32(2), int32(8), int32(2), ForKind::Default, eval(11),
None, UNKNOWN_LOC)
+ .unwrap();
+ let branch = If::new(&condition, &body, Some(&eval(12)),
UNKNOWN_LOC).unwrap();
+ let empty_loops =
+ [(0, 0, 1), (4, 0, -1)].into_iter().enumerate().map(|(n, (begin, end,
step))| {
+ Stmt::from(
+ For::new(
+ &i,
+ int32(begin),
+ int32(end),
+ int32(step),
+ ForKind::Default,
+ eval(14 + n as i64),
+ None,
+ UNKNOWN_LOC,
+ )
+ .unwrap(),
+ )
+ });
+ let root = SeqStmt::new(
+ [Stmt::from(Bind::new(&x, &bound, UNKNOWN_LOC).unwrap()),
branch.into(), eval(13)]
+ .into_iter()
+ .chain(empty_loops),
+ UNKNOWN_LOC,
+ );
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let analyzer = Analyzer::new();
+ let mut facts = Vec::new();
+ let mut loop_bounds = Vec::new();
+ let mut definitions = Vec::new();
+ walk(
+ &root,
+ (
+ |node: &VarObj, _: &ConstraintContext, kind: DefRegionKind| {
+ definitions.push((node.name.as_str().to_owned(), kind));
+ WalkResult::Advance
+ },
+ (
+ |node: &EvaluateObj, context: &ConstraintContext| ->
Result<WalkResult> {
+ let marker =
&node.value.downcast_ref::<IntImmObj>().unwrap().value;
+ facts.push((
+ marker.to_i64().unwrap(),
+ context.to_array().iter().collect::<Vec<_>>(),
+ ));
+ analyzer.sync_constraint_context(context)?;
+ loop_bounds.push(if
analyzer.const_int_bound_is_bound(&i)? {
+ let bounds = analyzer.const_int_bound(&i)?;
+ Some((bounds.min_value, bounds.max_value))
+ } else {
+ None
+ });
+ Ok(WalkResult::Advance)
+ },
+ |_: AnyView<'_>| WalkResult::Advance,
+ ),
+ ),
+ order,
+ )
+ .unwrap();
+ assert_eq!(
+ facts.iter().map(|(marker, facts)| (*marker,
facts.len())).collect::<Vec<_>>(),
+ [(11, 5), (12, 2), (13, 1), (14, 1), (15, 1)]
+ );
+ assert_eq!(loop_bounds, [Some((2, 7)), None, None, None, None]);
+ let binding =
facts[0].1[0].downcast_ref::<BindingConstraintObj>().unwrap();
+ assert!(binding.var.same_as(&x));
+ assert!(binding.value.same_as(&bound));
+ assert!(facts[0].1[0].same_as(&facts[2].1[0]));
+ let positive =
facts[0].1[1].downcast_ref::<PredicateConstraintObj>().unwrap();
+ assert!(positive.condition.same_as(&condition));
+ assert!(!positive.assume_always);
+ let negative =
facts[1].1[1].downcast_ref::<PredicateConstraintObj>().unwrap();
+
assert!(negative.condition.downcast_ref::<prim::NotObj>().unwrap().a.same_as(&condition));
+ assert!(!negative.assume_always);
+ let range =
facts[0].1[2].downcast_ref::<RangeConstraintObj>().unwrap();
+ assert!(range.var.same_as(&i));
+
assert_eq!(range.range.begin.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(2));
+
assert_eq!(range.range.extent.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(6));
+ assert!(definitions
+ .iter()
+ .any(|(name, kind)| name == "typed_i" && *kind ==
DefRegionKind::Simple));
+ }
+}
+
+#[test]
+fn live_context_and_snapshot_have_distinct_prefix_semantics() {
+ let _serial = serial();
+ let (x, y) = (variable("snapshot_x"), variable("snapshot_y"));
+ let root = SeqStmt::new(
+ [
+ Stmt::from(Bind::new(&x, int32(1), UNKNOWN_LOC).unwrap()),
+ eval(10),
+ Bind::new(&y, int32(2), UNKNOWN_LOC).unwrap().into(),
+ eval(20),
+ ],
+ UNKNOWN_LOC,
+ );
+ let mut seen = Observe::default();
+ visit(&root, &mut seen).unwrap();
+ let live = seen.live.unwrap();
+ assert_eq!(live.local_size(), 2);
+ assert_eq!(live.to_array().len(), 2);
+ assert_eq!(seen.snapshot.unwrap().to_array().len(), 1);
+ assert_eq!(seen.array.unwrap().len(), 1);
+ assert_eq!(seen.facts[0].1.len(), 1);
+ assert_eq!(seen.facts[1].1.len(), 2);
+}
+
+#[test]
+fn explicit_reentry_uses_borrowed_state_and_restores_definition_regions() {
+ let _serial = serial();
+ struct Reenter<'a> {
+ seen: &'a mut Vec<(i64, DefRegionKind)>,
+ depth: usize,
+ }
+ #[dispatch(visit, (AnyView<'_>, &mut
StructuralVisitorWithConstraintContextView<'_, Self>),
Result<Option<VisitInterrupt>>)]
+ impl<'a> Reenter<'a> {
+ fn visit_integer(
+ &mut self,
+ value: i64,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ self.seen.push((value, visitor.def_region_kind()));
+ if self.depth == 0 {
+ self.depth += 1;
+ visitor.with_def_region_kind(DefRegionKind::Pattern, |visitor|
{
+ visitor.with_def_region_kind(DefRegionKind::Simple,
|visitor| {
+ visitor.visit(self, &(value + 1))
+ })
+ })?;
+ self.depth -= 1;
+ self.seen.push((value + 2, visitor.def_region_kind()));
+ }
+ Ok(None)
+ }
+ }
+ let mut observed = Vec::new();
+ visit(&1_i64, Reenter { seen: &mut observed, depth: 0 }).unwrap();
+ assert_eq!(
+ observed,
+ [(1, DefRegionKind::None), (2, DefRegionKind::Pattern), (3,
DefRegionKind::None)]
+ );
+}
+
+#[test]
+fn interrupts_restore_native_scope_and_errors_terminate() {
+ let _serial = serial();
+ struct Halt {
+ fail: bool,
+ seen: usize,
+ restored: bool,
+ }
+ #[dispatch(visit, (AnyView<'_>, &mut
StructuralVisitorWithConstraintContextView<'_, Self>),
Result<Option<VisitInterrupt>>)]
+ impl Halt {
+ fn visit_for(
+ &mut self,
+ node: &ForObj,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ // A shared base reference remains valid across native context/flag
+ // mutation because every such field has interior mutability.
+ let base = visitor.base();
+ let result = visitor.default_visit(self, node);
+ if result.is_ok() {
+ self.restored = base.constraints().to_array().is_empty();
+ }
+ result
+ }
+ fn visit_evaluate(
+ &mut self,
+ _: &EvaluateObj,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ self.seen += 1;
+ assert_eq!(visitor.base().constraints().to_array().len(), 3);
+ if self.fail {
+ Err(Error::new(RUNTIME_ERROR, "constraint stop", ""))
+ } else {
+ Ok(Some(VisitInterrupt::with(17_i64)))
+ }
+ }
+ }
+ let i = variable("halt_i");
+ let body =
+ For::new(&i, int32(0), int32(4), int32(1), ForKind::Default, eval(11),
None, UNKNOWN_LOC)
+ .unwrap();
+ let root = SeqStmt::new([Stmt::from(body), eval(12)], UNKNOWN_LOC);
+ let mut state = Halt { fail: false, seen: 0, restored: false };
+ let interrupt = visit(&root, &mut state).unwrap().unwrap();
+ assert_eq!(i64::try_from(interrupt.value.clone()).unwrap(), 17);
+ assert_eq!(state.seen, 1);
+ assert!(state.restored);
+ state.fail = true;
+ assert!(visit(&root, &mut
state).err().unwrap().to_string().contains("constraint stop"));
+ assert_eq!(state.seen, 2);
+
+ // Public walk order and Skip: pre-order prunes the leaf, post-order
observes it first.
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut trace = Vec::new();
+ walk(
+ &root,
+ |value: AnyView<'_>| {
+ if let Some(node) = value.downcast_ref::<EvaluateObj>() {
+ let marker =
+
node.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64().unwrap();
+ trace.push(("evaluate", marker));
+ return WalkResult::Skip;
+ }
+ if let Some(node) = value.downcast_ref::<IntImmObj>() {
+ let marker = node.value.to_i64().unwrap();
+ if [11, 12].contains(&marker) {
+ trace.push(("leaf", marker));
+ }
+ }
+ WalkResult::Advance
+ },
+ order,
+ )
+ .unwrap();
+ let expected = match order {
+ WalkOrder::PreOrder => vec![("evaluate", 11), ("evaluate", 12)],
+ WalkOrder::PostOrder => {
+ vec![("leaf", 11), ("evaluate", 11), ("leaf", 12),
("evaluate", 12)]
+ }
+ };
+ assert_eq!(trace, expected);
+ let mut calls = 0;
+ let interrupted = walk(
+ &root,
+ |_: &EvaluateObj| {
+ calls += 1;
+ WalkResult::interrupt_with(19_i64)
+ },
+ order,
+ )
+ .unwrap()
+ .unwrap();
+ assert_eq!(i64::try_from(interrupted.value.clone()).unwrap(), 19);
+ assert_eq!(calls, 1);
+ let failure = walk(
+ &root,
+ |_: &EvaluateObj| -> Result<WalkResult> {
+ Err(Error::new(RUNTIME_ERROR, "walk stop", ""))
+ },
+ order,
+ )
+ .err()
+ .unwrap();
+ assert!(failure.to_string().contains("walk stop"));
+ }
+}
+
+#[test]
+fn lazy_call_arms_gain_predicates_but_eager_select_arms_do_not() {
+ let _serial = serial();
+ #[derive(Default)]
+ struct Leaves(Vec<(i64, Vec<Constraint>)>);
+ #[dispatch(visit, (AnyView<'_>, &mut
StructuralVisitorWithConstraintContextView<'_, Self>),
Result<Option<VisitInterrupt>>)]
+ impl Leaves {
+ fn visit_integer(
+ &mut self,
+ node: &IntImmObj,
+ visitor: &mut StructuralVisitorWithConstraintContextView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ let value = node.value.to_i64().unwrap();
+ if [21, 22, 31, 32].contains(&value) {
+ self.0.push((value,
visitor.base().constraints().to_array().iter().collect()));
+ }
+ Ok(None)
+ }
+ }
+ let condition = Var::new("lazy_condition", PrimType::bool(), UNKNOWN_LOC);
+ let lazy = prim::op::if_then_else(&condition, int32(21),
int32(22)).unwrap();
+ let eager = prim::Select::new(&condition, int32(31), int32(32),
UNKNOWN_LOC).unwrap();
+ let root = Array::new(vec![lazy, eager.into()]);
+ let mut seen = Leaves::default();
+ visit(&root, &mut seen).unwrap();
+ assert_eq!(
+ seen.0.iter().map(|(n, facts)| (*n, facts.len())).collect::<Vec<_>>(),
+ [(21, 1), (22, 1), (31, 0), (32, 0)]
+ );
+ let positive =
seen.0[0].1[0].downcast_ref::<PredicateConstraintObj>().unwrap();
+ assert!(positive.condition.same_as(&condition));
+ let negative =
seen.0[1].1[0].downcast_ref::<PredicateConstraintObj>().unwrap();
+
assert!(negative.condition.downcast_ref::<prim::NotObj>().unwrap().a.same_as(&condition));
+}
+
+#[test]
+fn native_assume_and_thread_launch_hooks_preserve_fact_timing() {
+ let _serial = serial();
+ let condition = Var::new("assumption", PrimType::bool(), UNKNOWN_LOC);
+ let assume = Call::new(
+ None,
+ Op::get("tirx.assume_always").unwrap(),
+ [&condition],
+ None,
+ Vec::<Type>::new(),
+ UNKNOWN_LOC,
+ )
+ .unwrap();
+ let thread = variable("thread_index");
+ let launch = RegionStmt::new(
+ Op::get("tirx.launch_thread").unwrap(),
+ [Expr::from(StringImm::new("threadIdx.x", UNKNOWN_LOC)), int32(8)],
+ [&thread],
+ DictAttrs::new(Map::new()),
+ eval(42),
+ Vec::<Var>::new(),
+ UNKNOWN_LOC,
+ )
+ .unwrap();
+ let root = SeqStmt::new(
+ [eval(40), Evaluate::new(assume, UNKNOWN_LOC).into(), eval(41),
launch.into(), eval(43)],
+ UNKNOWN_LOC,
+ );
+ let mut seen = Observe::default();
+ visit(&root, &mut seen).unwrap();
+ assert_eq!(fact_counts(&seen), [(40, 0), (41, 1), (42, 2), (43, 1)]);
+ let predicate =
seen.facts[1].1[0].downcast_ref::<PredicateConstraintObj>().unwrap();
+ assert!(predicate.assume_always);
+ assert!(predicate.condition.same_as(&condition));
+ let range =
seen.facts[2].1[1].downcast_ref::<RangeConstraintObj>().unwrap();
+ assert!(range.var.same_as(&thread));
+ assert!(seen
+ .definitions
+ .iter()
+ .any(|(name, kind)| name == "thread_index" && *kind ==
DefRegionKind::Simple));
+}
