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 74d2576bf [Rust] Add constraint-aware structural map and mutation
(#189)
74d2576bf is described below
commit 74d2576bf841542aa92c6be7be3280cd4a73c2c2
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Nov 4 13:58:02 2026 -0500
[Rust] Add constraint-aware structural map and mutation (#189)
Add Rust constraint-aware structural map and mutation using the shared
generic mutation engine. Ordered map callbacks and explicit mutation
handlers observe lexical binding, predicate, and range facts derived
from rewritten values, with definition regions and scope restoration
preserved.
Separate the native semantic base from engine-owned remap storage so
Rust can use the complete native-compatible context layout. Compose the
native factory and legacy mutator with the shared engine, and check live
ownership before specialized in-place descent.
---
docs/ir_guides/working_with_ir_in_rust.rst | 18 ++
.../tirx/object_functor_with_constraint_context.h | 7 +-
.../structural_mutate_with_constraint_context.h | 28 +--
rust/src/tirx.rs | 2 +
.../tirx/structural_map_with_constraint_context.rs | 190 ++++++++++++++++
.../structural_mutate_with_constraint_context.rs | 168 ++++++++++++++
rust/tests/tirx/main.rs | 5 +-
.../tirx/structural_map_with_constraint_context.rs | 242 +++++++++++++++++++++
.../structural_mutate_with_constraint_context.rs | 176 +++++++++++++++
...=> structural_visit_with_constraint_context.rs} | 154 +------------
.../structural_walk_with_constraint_context.rs | 187 ++++++++++++++++
src/tirx/ir/structural_constraints.cc | 3 +-
12 files changed, 1000 insertions(+), 180 deletions(-)
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index d5d260d4d..f2285cf3e 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -669,6 +669,20 @@ replacement before any use of ``i``.
``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.
+ |structural_mutate_with_constraint_context|_ uses
+ ``StructuralMutatorWithConstraintContextView`` for explicit mutation
descent.
+ ``mutator.base().constraints()`` exposes facts derived from rewritten
bindings,
+ predicates, and loop ranges. Call ``default_mutate(self, value)`` to collect
+ those facts while descending; a replacement returned directly skips that
+ node's descent. As with ordinary mutation, moved unique roots may reuse
storage
+ and retained roots copy.
+ |structural_map_with_constraint_context|_ provides ordered mapping with the
+ same lexical facts. Callbacks receive a leading prefix of
+ ``(value, &ConstraintContext, DefRegionKind)`` and return a replacement or
+ ``Unchanged``. ``WalkOrder::PreOrder`` descends into replacement children;
+ ``PostOrder`` invokes the mapper after children are rewritten and the node's
+ constraint scope is restored. Facts for later siblings come from the
+ resulting bindings and expressions.
Typed contexts are neither ``Send`` nor ``Sync``; erased FFI handles retain
the ordinary contract against concurrent access.
@@ -748,6 +762,8 @@ Quick Reference
.. _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
+.. |structural_mutate_with_constraint_context| replace::
``structural_mutate_with_constraint_context``
+.. _structural_mutate_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_mutate_with_constraint_context/fn.structural_mutate_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!``
@@ -758,3 +774,5 @@ Quick Reference
.. _downcast_ref:
../reference/rust/generated/tvm_ffi/object/trait.ObjectRefCore.html#method.downcast_ref
.. |register_global_typed| replace:: ``Function::register_global_typed``
.. _register_global_typed:
../reference/rust/generated/tvm_ffi/function/struct.Function.html#method.register_global_typed
+.. |structural_map_with_constraint_context| replace::
``structural_map_with_constraint_context``
+.. _structural_map_with_constraint_context:
../reference/rust/generated/tvm_tirx/tirx/structural_map_with_constraint_context/fn.structural_map_with_constraint_context.html
diff --git a/include/tvm/tirx/object_functor_with_constraint_context.h
b/include/tvm/tirx/object_functor_with_constraint_context.h
index 1949ff51c..57bb43fb6 100644
--- a/include/tvm/tirx/object_functor_with_constraint_context.h
+++ b/include/tvm/tirx/object_functor_with_constraint_context.h
@@ -87,7 +87,10 @@ class TVM_TIRX_DLL ObjectVisitorWithConstraintContext
* values use structural fallback. Errors propagate without restoring the
context.
*/
class TVM_TIRX_DLL ObjectMutatorWithConstraintContext
- : public StructuralMutatorWithConstraintContextObj {
+ : public
ffi::detail::StructuralMutatorEngineBase<StructuralMutatorWithConstraintContextObj>
{
+ using EngineBase =
+
ffi::detail::StructuralMutatorEngineBase<StructuralMutatorWithConstraintContextObj>;
+
public:
ObjectMutatorWithConstraintContext() :
ObjectMutatorWithConstraintContext(GlobalVTable()) {}
ObjectMutatorWithConstraintContext(const
ObjectMutatorWithConstraintContext&) = delete;
@@ -114,7 +117,7 @@ class TVM_TIRX_DLL ObjectMutatorWithConstraintContext
const ffi::Object*, ObjectMutatorWithConstraintContext*,
ffi::InplaceMode)>;
explicit ObjectMutatorWithConstraintContext(const VTable* table)
- : StructuralMutatorWithConstraintContextObj(StructuralVTable()),
native_vtable_(table) {}
+ : EngineBase(StructuralVTable()), native_vtable_(table) {}
static void InitVTable(VTable* table) {}
diff --git a/include/tvm/tirx/structural_mutate_with_constraint_context.h
b/include/tvm/tirx/structural_mutate_with_constraint_context.h
index 9d79feef8..b5fa84c8b 100644
--- a/include/tvm/tirx/structural_mutate_with_constraint_context.h
+++ b/include/tvm/tirx/structural_mutate_with_constraint_context.h
@@ -33,17 +33,14 @@ namespace tirx {
/*!
* \brief Mutates values using their structural mutation behavior while
maintaining
* the constraint context relevant to each callback.
- * Constraints use the resulting expressions and bindings.
+ * Constraints use the resulting expressions and bindings. This semantic base
contains
+ * only context state; concrete structural engines own the identity-remap
environment.
* \sa ffi::StructuralMutator, sym::ConstraintContext
*/
-class StructuralMutatorWithConstraintContextObj : public
ffi::StructuralMapEngineBase {
+class StructuralMutatorWithConstraintContextObj : public
ffi::StructuralMutatorObj {
public:
using MutatorObjType = StructuralMutatorWithConstraintContextObj;
using StateTupleType = std::tuple<const sym::ConstraintContext&>;
- /*!
- * \brief Construct a mutator using constraint-aware default descent.
- */
- StructuralMutatorWithConstraintContextObj();
/*!
* \brief Active effective premises. Consumers must not append during
traversal.
* Retain this handle and constraints->LocalSize() for a stable position, or
use
@@ -95,7 +92,7 @@ class StructuralMutatorWithConstraintContextObj : public
ffi::StructuralMapEngin
protected:
explicit StructuralMutatorWithConstraintContextObj(const
ffi::StructuralMutatorVTable* vtable)
- : ffi::StructuralMapEngineBase(vtable) {}
+ : ffi::StructuralMutatorObj(vtable) {}
StateTupleType StateTuple() const noexcept { return {constraints}; }
@@ -206,23 +203,6 @@ inline constexpr
ffi::reflection::AttrKey<FStructuralMutateWithConstraintContext
"__s_maybe_inplace_mutate_with_constraint_context__"};
} // namespace type_attr
-inline
StructuralMutatorWithConstraintContextObj::StructuralMutatorWithConstraintContextObj()
- : StructuralMutatorWithConstraintContextObj([] {
- static const ffi::StructuralMutatorVTable table{
- [](ffi::StructuralMutatorObj* self, ffi::AnyView value) noexcept {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
- static_cast<StructuralMutatorWithConstraintContextObj*>(self)
- ->DefaultMutateExpected(value,
ffi::InplaceMode::kDisallow));
- },
- [](ffi::StructuralMutatorObj* self, ffi::AnyView value) noexcept {
- return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
- static_cast<StructuralMutatorWithConstraintContextObj*>(self)
- ->DefaultMutateExpected(value,
ffi::InplaceMode::kAllow));
- },
- DispatchVarRemapGet, DispatchVarRemapSet};
- return &table;
- }()) {}
-
inline ffi::Expected<ffi::UnchangedOr<ffi::Any>>
StructuralMutatorWithConstraintContextObj::DefaultMutateExpected(ffi::AnyView
value,
ffi::InplaceMode mode) noexcept {
diff --git a/rust/src/tirx.rs b/rust/src/tirx.rs
index 6cadcebfc..52f1ba824 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_map_with_constraint_context;
+pub mod structural_mutate_with_constraint_context;
pub mod structural_visit_with_constraint_context;
pub mod structural_walk_with_constraint_context;
pub mod tensor;
diff --git a/rust/src/tirx/structural_map_with_constraint_context.rs
b/rust/src/tirx/structural_map_with_constraint_context.rs
new file mode 100644
index 000000000..2d5e721ec
--- /dev/null
+++ b/rust/src/tirx/structural_map_with_constraint_context.rs
@@ -0,0 +1,190 @@
+/*
+ * 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.
+ */
+
+//! Ordered constraint-aware mapping through the shared mutation engine.
+//!
+//! This adapter owns callback order only. The mutation engine owns dispatch,
+//! remaps and copy-on-write; native hooks own lexical facts and scope
restoration.
+
+use std::borrow::BorrowMut;
+
+use tvm_ffi::{
+ Any, AnyView, DefRegionKind, Dispatcher, IntoDispatcher, MutationView,
Object, Result,
+ WalkOrder,
+};
+
+use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
+
+use super::structural_mutate_with_constraint_context::{
+ structural_mutate_with_constraint_context,
StructuralMutatorWithConstraintContextView,
+};
+use crate::sym::ConstraintContext;
+
+struct MapMutator<'a, D, const PRE: bool> {
+ mapper: &'a mut D,
+}
+
+impl<'v, 'c, 'r, 'a, D, const PRE: bool>
+ Dispatcher<
+ (
+ MutationView<'v>,
+ &'c mut StructuralMutatorWithConstraintContextView<'r,
MapMutator<'a, D, PRE>>,
+ ),
+ Option<Result<Any>>,
+ > for &mut MapMutator<'a, D, PRE>
+where
+ for<'s, 'b, 'ctx> &'s mut D:
+ Dispatcher<(AnyView<'b>, &'ctx ConstraintContext, DefRegionKind),
Option<Result<Any>>>,
+{
+ #[inline(always)]
+ fn dispatch_tuple(
+ self,
+ (value, mutator): (
+ MutationView<'v>,
+ &'c mut StructuralMutatorWithConstraintContextView<'r,
MapMutator<'a, D, PRE>>,
+ ),
+ ) -> Option<Result<Any>> {
+ Some(self.map(value, mutator))
+ }
+}
+
+impl<D, const PRE: bool> MapMutator<'_, D, PRE>
+where
+ for<'s, 'v, 'c> &'s mut D:
+ Dispatcher<(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
Option<Result<Any>>>,
+{
+ #[inline(always)]
+ fn map(
+ &mut self,
+ value: MutationView<'_>,
+ mutator: &mut StructuralMutatorWithConstraintContextView<'_, Self>,
+ ) -> Result<Any> {
+ let kind = mutator.def_region_kind();
+ if PRE {
+ let mode = value.inplace_mode();
+ // The mapper cannot access this runtime. Its context borrow ends
+ // before descent, which may replace the interior-mutable field.
+ let result = unsafe {
+ self.mapper.dispatch_tuple((
+ value.as_any_view(),
+ mutator.base().constraints_for_observer(),
+ kind,
+ ))
+ };
+ let Some(result) = result else {
+ return mutator.default_mutate(self, value).map(Into::into);
+ };
+ let mut mapped = result?;
+ if mapped.type_index() == TVMFFITypeIndex::kTVMFFIUnchanged as i32
{
+ return mutator.default_mutate(self, value).map(Into::into);
+ }
+ let same_object = value
+ .downcast_ref::<Object>()
+ .zip(mapped.downcast_ref::<Object>())
+ .is_some_and(|(original, replacement)| std::ptr::eq(original,
replacement));
+ if same_object {
+ // An identity callback must not make a unique root appear
shared.
+ drop(mapped);
+ return mutator.default_mutate(self, value).map(Into::into);
+ }
+ // An owned replacement can lend a capability, restricted by the
+ // original path's permission. The shared engine checks live
aliases.
+ let replacement = MutationView::from(&mut
mapped).with_inplace_mode(mode);
+ let result = mutator.default_mutate(self, replacement)?;
+ Ok(result.value_or(mapped))
+ } else {
+ // SAFETY: this private adapter runs only inside the shared engine
+ // callback. The engine owns the root; recursive callers keep each
+ // input field owned until this synchronous callback returns.
Default
+ // descent can mutate fields but cannot destroy or retype that
input.
+ // This snapshot holds no typed node reference across descent. Only
+ // reborrow it afterward when Unchanged preserves the input, and
never
+ // let the reborrow escape this callback.
+ let original: AnyView<'_> = unsafe {
std::mem::transmute(value.as_any_view()) };
+ let mapped = mutator.default_mutate(self, value)?;
+ let unchanged = mapped.is_unchanged();
+ let mapped: Any = mapped.into();
+ let current = if unchanged { original } else {
AnyView::from(&mapped) };
+ // Scoped native hooks have restored the incoming context here.
+ let result = unsafe {
+ self.mapper.dispatch_tuple((
+ current,
+ mutator.base().constraints_for_observer(),
+ kind,
+ ))
+ };
+ match result {
+ Some(result) => {
+ let result = result?;
+ Ok(if result.type_index() ==
TVMFFITypeIndex::kTVMFFIUnchanged as i32 {
+ mapped
+ } else {
+ result
+ })
+ }
+ None => Ok(mapped),
+ }
+ }
+ }
+}
+
+/// Map every occurrence with the constraint facts active at that position.
+///
+/// Pre-order callbacks run before descending into the replacement's children;
+/// post-order callbacks see the rewritten value after its scope is restored.
+/// Facts for later siblings use the resulting bindings and expressions.
+/// `Unchanged` keeps the current value and still descends in pre-order.
+/// The root is consumed: moved unique roots may reuse storage, retained roots
copy.
+/// Errors stop traversal without rolling back completed in-place changes.
+///
+/// Accepts the shared [`IntoDispatcher`] forms: owned or borrowed dispatchers,
+/// typed callbacks, and nested tuples. Callbacks take a leading prefix of
+/// `(value, &ConstraintContext, DefRegionKind)` and return a replacement,
+/// `Unchanged`, or a `Result` thereof. Observers must not append to the active
+/// context or panic across native callbacks. Clone a context to retain its
live
+/// handle, or use `snapshot()` / `to_array()` for a fixed prefix.
+pub fn structural_map_with_constraint_context<R, H, M, D, O>(
+ root: R,
+ mapper: H,
+ order: WalkOrder,
+) -> Result<Any>
+where
+ R: Into<Any>,
+ O: BorrowMut<D>,
+ for<'v, 'c> H: IntoDispatcher<
+ (AnyView<'v>, &'c ConstraintContext, DefRegionKind),
+ Option<Result<Any>>,
+ M,
+ Target = D,
+ Output = O,
+ >,
+ for<'s, 'v, 'c> &'s mut D:
+ Dispatcher<(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
Option<Result<Any>>>,
+{
+ let mut storage = mapper.into_dispatcher();
+ let mapper = storage.borrow_mut();
+ match order {
+ WalkOrder::PreOrder => {
+ structural_mutate_with_constraint_context(root, MapMutator::<D,
true> { mapper })
+ }
+ WalkOrder::PostOrder => {
+ structural_mutate_with_constraint_context(root, MapMutator::<D,
false> { mapper })
+ }
+ }
+}
diff --git a/rust/src/tirx/structural_mutate_with_constraint_context.rs
b/rust/src/tirx/structural_mutate_with_constraint_context.rs
new file mode 100644
index 000000000..88ddb50e1
--- /dev/null
+++ b/rust/src/tirx/structural_mutate_with_constraint_context.rs
@@ -0,0 +1,168 @@
+/*
+ * 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 structural mutation through the shared Rust engine.
+//!
+//! The native semantic base carries only lexical context state. Rust owns its
+//! dispatcher and remaps in `StructuralMutateEngine<Base>`; native constraint
+//! hooks recurse through that engine's vtable and collect facts from results.
+
+use std::borrow::BorrowMut;
+use std::cell::UnsafeCell;
+
+use tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock;
+use tvm_ffi::extra::structural_mutate::{StructuralMutateEngine,
StructuralMutatorBase};
+use tvm_ffi::reflection::TypeAttrColumn;
+use tvm_ffi::object::ObjectRef;
+use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex};
+use tvm_ffi::{
+ Any, Dispatcher, InplaceMode, IntoDispatcher, MutationView, ObjectArc,
ObjectCore,
+ ObjectRefCore, Result, StructuralMutatorObj, StructuralMutatorView,
+};
+
+use super::StructuralMutatorWithConstraintContextObj;
+use crate::sym::ConstraintContext;
+
+/// Complete native-compatible semantic base, without an owning object
conversion.
+///
+/// Native hooks may replace the current context during descent through a
shared
+/// view. All native-mutated state is interior mutable. This representation
+/// requires the matching native header/build, including the private scope
flag.
+#[repr(C)]
+pub struct StructuralMutatorWithConstraintContextBase {
+ base: StructuralMutatorObj,
+ constraints: UnsafeCell<ConstraintContext>,
+ needs_child_context: UnsafeCell<bool>,
+}
+
+impl StructuralMutatorWithConstraintContextBase {
+ /// Clone the current live context without retaining a borrow of its field.
+ /// Use `snapshot()` on the returned handle to freeze its current prefix.
+ #[inline]
+ pub fn constraints(&self) -> ConstraintContext {
+ unsafe { (&*self.constraints.get()).clone() }
+ }
+
+ /// Borrow only while a mapper cannot recurse through this mutator.
+ /// End this borrow before default descent can replace the context field.
+ #[inline]
+ pub(super) unsafe fn constraints_for_observer(&self) -> &ConstraintContext
{
+ &*self.constraints.get()
+ }
+}
+
+// SAFETY: repr(C) matches the complete native state-only semantic base. The
+// shared engine installs its vtable, owns remaps and publishes the
dispatcher's
+// actual recursive reborrow. Native hooks borrow this runtime synchronously
and
+// cannot retain it. No shared reference freezes the replaceable context/flag.
+// Each default call transfers exactly one owning replacement/Unchanged/Error.
+unsafe impl StructuralMutatorBase for
StructuralMutatorWithConstraintContextBase {
+ fn new() -> Result<Self> {
+ Ok(Self {
+ base: StructuralMutatorObj::default(),
+ constraints: UnsafeCell::new(ConstraintContext::new()?),
+ needs_child_context: UnsafeCell::new(false),
+ })
+ }
+
+ fn type_index() -> i32 {
+ StructuralMutatorWithConstraintContextObj::type_index()
+ }
+
+ #[inline(always)]
+ unsafe fn default_mutate_raw(runtime: *mut Self, value: MutationView<'_>)
-> TVMFFIAny {
+ static COPYING: ColdInitLazyLock<Option<TypeAttrColumn>> =
+ ColdInitLazyLock::new(||
TypeAttrColumn::new("__s_mutate_with_constraint_context__"));
+ static INPLACE: ColdInitLazyLock<Option<TypeAttrColumn>> =
ColdInitLazyLock::new(|| {
+
TypeAttrColumn::new("__s_maybe_inplace_mutate_with_constraint_context__")
+ });
+ let borrowed = value.as_any_view();
+ let lookup = |column: Option<TypeAttrColumn>| {
+ column
+ .and_then(|column|
column.get_raw_ffi_any(borrowed.type_index()))
+ .filter(|hook| hook.type_index != TVMFFITypeIndex::kTVMFFINone
as i32)
+ };
+ let hook = if value.inplace_mode() == InplaceMode::Allow {
+ lookup(*INPLACE).filter(|_| {
+ // Permission alone is insufficient after a callback retains an
+ // alias. The supported owning count includes this temporary;
+ // release it before invoking a hook or entering native
descent.
+ borrowed.try_as::<ObjectRef>().is_some_and(|owner| {
+ let unique =
ObjectArc::strong_count(ObjectRef::data(&owner)) == 2;
+ drop(owner);
+ unique
+ })
+ })
+ } else {
+ None
+ }
+ .or_else(|| lookup(*COPYING));
+ if let Some(hook) = hook {
+ // Copy the borrowed registry token; the shared helper acquires
+ // Function ownership before packed reentry and handles error
context.
+ StructuralMutatorObj::call_structural_mutate_hook_raw(
+ runtime.cast(),
+ *borrowed.as_raw_ffi_any(),
+ hook,
+ )
+ } else {
+ // Prefix-only fallback uses installed remap/callback thunks, never
+ // an ordinary Rust engine's dispatcher offset.
+ StructuralMutatorObj::default_mutate_raw(runtime.cast(), value)
+ }
+ }
+}
+
+/// Scoped constraint mutator view. Pass the handler's actual `self` to
recursive
+/// operations and use `view.base().constraints()` to inspect resulting
premises.
+/// Callbacks follow the shared engine's non-escaping, non-unwinding contract.
+pub type StructuralMutatorWithConstraintContextView<'a, D> =
+ StructuralMutatorView<'a, D, StructuralMutatorWithConstraintContextBase>;
+
+/// Mutate using native lexical binding, predicate and range semantics.
+///
+/// The first matching handler owns descent. Use `view.default_mutate(self,
value)`
+/// for constraint-aware default descent and `view.mutate(self, child)` to
dispatch
+/// recursively. A moved unique root may reuse storage; retained roots copy.
+/// Specialized hooks obey the ordinary synchronous mutation ABI and
registration
+/// trust contract. They must not retain the runtime or append through observer
+/// context handles during descent. Errors terminate the traversal.
+pub fn structural_mutate_with_constraint_context<R, H, M, D, O>(
+ root: R,
+ dispatcher: H,
+) -> Result<Any>
+where
+ R: Into<Any>,
+ O: BorrowMut<D>,
+ for<'v, 'c, 'r> H: IntoDispatcher<
+ (MutationView<'v>, &'c mut
StructuralMutatorWithConstraintContextView<'r, D>),
+ Option<Result<Any>>,
+ M,
+ Target = D,
+ Output = O,
+ >,
+ for<'s, 'v, 'c, 'r> &'s mut D: Dispatcher<
+ (MutationView<'v>, &'c mut
StructuralMutatorWithConstraintContextView<'r, D>),
+ Option<Result<Any>>,
+ >,
+{
+
StructuralMutateEngine::<StructuralMutatorWithConstraintContextBase>::mutate_root(
+ root, dispatcher,
+ )
+}
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index 67d2ad347..0a250f061 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -28,7 +28,10 @@ mod function;
mod index_map;
mod layout;
mod op;
-mod structural;
+mod structural_map_with_constraint_context;
+mod structural_mutate_with_constraint_context;
+mod structural_visit_with_constraint_context;
+mod structural_walk_with_constraint_context;
mod tensor;
mod thread_scope;
mod tile;
diff --git a/rust/tests/tirx/structural_map_with_constraint_context.rs
b/rust/tests/tirx/structural_map_with_constraint_context.rs
new file mode 100644
index 000000000..b4d8ed280
--- /dev/null
+++ b/rust/tests/tirx/structural_map_with_constraint_context.rs
@@ -0,0 +1,242 @@
+/*
+ * 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.
+ */
+
+//! Ordered mapping, result-derived facts and mutation permissions.
+
+use tvm_tirx::ir::*;
+use tvm_tirx::sym::{
+ BindingConstraintObj, ConstraintContext, PredicateConstraintObj,
RangeConstraintObj,
+};
+use
tvm_tirx::tirx::structural_map_with_constraint_context::structural_map_with_constraint_context
as map;
+use tvm_tirx::tvm_ffi::{
+ Any, AnyView, DefRegionKind, Error, ObjectRefCore, Result, Unchanged,
WalkOrder, RUNTIME_ERROR,
+};
+
+use super::{int32, op};
+use crate::common::serial;
+
+fn eval(value: i64) -> Stmt {
+ Evaluate::new(int32(value), UNKNOWN_LOC).into()
+}
+
+#[test]
+fn ordered_mapping_collects_result_facts_and_restores_scopes() {
+ let _serial = serial();
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ for moved in [false, true] {
+ let x = Var::new("map_x", PrimType::int(32), UNKNOWN_LOC);
+ let i = Var::new("map_i", PrimType::int(32), UNKNOWN_LOC);
+ let body = For::new(
+ &i,
+ int32(0),
+ int32(4),
+ int32(1),
+ ForKind::Default,
+ eval(11),
+ None,
+ UNKNOWN_LOC,
+ )
+ .unwrap();
+ let branch =
+ If::new(op("prim._OpLT", &x, int32(4)), body, Some(&eval(12)),
UNKNOWN_LOC)
+ .unwrap();
+ let root = SeqStmt::new(
+ [
+ Stmt::from(Bind::new(&x, int32(4), UNKNOWN_LOC).unwrap()),
+ branch.into(),
+ eval(13),
+ ],
+ UNKNOWN_LOC,
+ );
+ let pointer = &*root as *const SeqStmtObj;
+ let retained = (!moved).then(|| root.clone());
+ let mut facts = Vec::new();
+ let mut definitions = Vec::new();
+ let mut loop_contexts = Vec::new();
+ let result = map(
+ root,
+ (
+ |node: &IntImmObj| -> Any {
+ if node.value.to_i64() == Some(4) {
+ int32(8).into()
+ } else {
+ Unchanged.into()
+ }
+ },
+ |_: &VarObj, _: &ConstraintContext, kind: DefRegionKind| {
+ definitions.push(kind);
+ Unchanged
+ },
+ (
+ |node: &EvaluateObj, context: &ConstraintContext| {
+ facts.push((
+ node.value
+ .downcast_ref::<IntImmObj>()
+ .unwrap()
+ .value
+ .to_i64()
+ .unwrap(),
+ context.to_array().iter().collect::<Vec<_>>(),
+ ));
+ Unchanged
+ },
+ |_: &ForObj, context: &ConstraintContext| {
+ loop_contexts.push(context.to_array().len());
+ Unchanged
+ },
+ ),
+ ),
+ order,
+ )
+ .unwrap();
+ let result = SeqStmt::try_from(result).unwrap();
+ assert_eq!(std::ptr::eq(pointer, &*result), moved);
+ assert_eq!(
+ facts.iter().map(|(marker, facts)| (*marker,
facts.len())).collect::<Vec<_>>(),
+ [(11, 5), (12, 2), (13, 1)]
+ );
+ assert_eq!(loop_contexts, [2]);
+ assert!(definitions.contains(&DefRegionKind::Simple));
+ let binding = result.seq[0].downcast_ref::<BindObj>().unwrap();
+ let branch = result.seq[1].downcast_ref::<IfObj>().unwrap();
+ let body =
branch.then_case.seq[0].downcast_ref::<ForObj>().unwrap();
+ let fact =
facts[0].1[0].downcast_ref::<BindingConstraintObj>().unwrap();
+ assert!(fact.var.same_as(&binding.var));
+ assert!(fact.value.same_as(&binding.value));
+
assert_eq!(fact.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+ let fact =
facts[0].1[1].downcast_ref::<PredicateConstraintObj>().unwrap();
+ assert!(fact.condition.same_as(&branch.condition));
+ let fact =
facts[0].1[2].downcast_ref::<RangeConstraintObj>().unwrap();
+ assert!(fact.var.same_as(&body.var));
+ assert_eq!(
+
fact.range.begin.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
+ Some(0)
+ );
+ assert_eq!(
+
fact.range.extent.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
+ Some(8)
+ );
+ assert!(facts[0].1[0].same_as(&facts[2].1[0]));
+ if let Some(retained) = retained {
+ let binding =
retained.seq[0].downcast_ref::<BindObj>().unwrap();
+ assert_eq!(
+
binding.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
+ Some(4)
+ );
+ }
+ }
+ }
+}
+
+#[test]
+fn replacement_order_identity_aliases_and_errors() {
+ let _serial = serial();
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ let mut seen = Vec::new();
+ let result = map(
+ eval(11),
+ (
+ |_: &EvaluateObj| eval(21),
+ |node: &IntImmObj| {
+ let value = node.value.to_i64().unwrap();
+ seen.push(value);
+ int32(value + 1)
+ },
+ ),
+ order,
+ )
+ .unwrap();
+ let result = result.downcast_ref::<EvaluateObj>().unwrap();
+ let expected = match order {
+ WalkOrder::PreOrder => 22,
+ WalkOrder::PostOrder => 21,
+ };
+ assert_eq!(
+ result.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
+ Some(expected)
+ );
+ assert_eq!(seen, [if matches!(order, WalkOrder::PreOrder) { 21 } else
{ 11 }]);
+
+ let mut calls = 0;
+ let error = map(
+ SeqStmt::new([eval(11), eval(12)], UNKNOWN_LOC),
+ |_: &EvaluateObj| -> Result<Any> {
+ calls += 1;
+ Err(Error::new(RUNTIME_ERROR, "map stop", ""))
+ },
+ order,
+ )
+ .err()
+ .unwrap();
+ assert!(error.to_string().contains("map stop"));
+ assert_eq!(calls, 1);
+ }
+ // Returning the same owning handle must not prevent reuse; retaining that
+ // handle separately must force copying even after callback entry.
+ for order in [WalkOrder::PreOrder, WalkOrder::PostOrder] {
+ for retain_alias in [false, true] {
+ let i = Var::new("map_alias", PrimType::int(32), UNKNOWN_LOC);
+ let root = For::new(
+ &i,
+ int32(0),
+ int32(4),
+ int32(1),
+ ForKind::Default,
+ eval(11),
+ None,
+ UNKNOWN_LOC,
+ )
+ .unwrap();
+ let pointer = &*root as *const ForObj;
+ let mut saved: Option<Any> = None;
+ let result = map(
+ root,
+ (
+ |node: &ForObj| -> Any {
+ let value: Any = AnyView::from(node).into();
+ if retain_alias {
+ saved = Some(value.clone());
+ }
+ value
+ },
+ |node: &IntImmObj| -> Any {
+ if node.value.to_i64() == Some(4) {
+ int32(8).into()
+ } else {
+ Unchanged.into()
+ }
+ },
+ ),
+ order,
+ )
+ .unwrap();
+ let result = result.downcast_ref::<ForObj>().unwrap();
+ let post = matches!(order, WalkOrder::PostOrder);
+ assert_eq!(std::ptr::eq(pointer, result), !retain_alias || post);
+
assert_eq!(result.end.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+ if let Some(saved) = saved {
+ let saved = saved.downcast_ref::<ForObj>().unwrap();
+ assert_eq!(
+
saved.end.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
+ Some(if post { 8 } else { 4 })
+ );
+ }
+ }
+ }
+}
diff --git a/rust/tests/tirx/structural_mutate_with_constraint_context.rs
b/rust/tests/tirx/structural_mutate_with_constraint_context.rs
new file mode 100644
index 000000000..88f674c18
--- /dev/null
+++ b/rust/tests/tirx/structural_mutate_with_constraint_context.rs
@@ -0,0 +1,176 @@
+/*
+ * 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.
+ */
+
+//! Result-derived constraint facts and shared-engine mutation permissions.
+
+use tvm_tirx::ir::*;
+use tvm_tirx::sym::{BindingConstraintObj, Constraint, PredicateConstraintObj,
RangeConstraintObj};
+use tvm_tirx::tirx::structural_mutate_with_constraint_context::{
+ structural_mutate_with_constraint_context as mutate,
+ StructuralMutatorWithConstraintContextView as View,
+};
+use tvm_tirx::tvm_ffi::{dispatch, Any, DefRegionKind, MutationView,
ObjectRefCore, Result, Unchanged};
+
+use super::{int32, op};
+use crate::common::serial;
+
+fn eval(value: i64) -> Stmt {
+ Evaluate::new(int32(value), UNKNOWN_LOC).into()
+}
+
+#[derive(Default)]
+struct Rewrite {
+ facts: Vec<(i64, Vec<Constraint>)>,
+ definitions: Vec<DefRegionKind>,
+ restored: bool,
+}
+
+#[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>), Result<Any>)]
+impl Rewrite {
+ fn mutate_integer(
+ &mut self,
+ node: MutationView<'_, IntImm>,
+ _: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ if node.value.to_i64().unwrap() == 4 {
+ Ok(int32(8).into())
+ } else {
+ Ok(Unchanged.into())
+ }
+ }
+ fn mutate_var(
+ &mut self,
+ node: MutationView<'_, Var>,
+ view: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ self.definitions.push(view.def_region_kind());
+ view.default_mutate(self, node).map(Into::into)
+ }
+ fn mutate_evaluate(
+ &mut self,
+ node: MutationView<'_, Evaluate>,
+ view: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ let marker =
node.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64().unwrap();
+ self.facts.push((marker,
view.base().constraints().to_array().iter().collect()));
+ view.default_mutate(self, node).map(Into::into)
+ }
+ fn mutate_for(
+ &mut self,
+ node: MutationView<'_, For>,
+ view: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ let before = view.base().constraints().to_array().len();
+ let result = view.default_mutate(self, node);
+ self.restored = view.base().constraints().to_array().len() == before;
+ result.map(Into::into)
+ }
+}
+
+#[test]
+fn result_facts_and_scope_restore_for_retained_and_moved_roots() {
+ let _serial = serial();
+ for moved in [false, true] {
+ let x = Var::new("mutated_x", PrimType::int(32), UNKNOWN_LOC);
+ let i = Var::new("mutated_i", PrimType::int(32), UNKNOWN_LOC);
+ let body = For::new(
+ &i,
+ int32(0),
+ int32(4),
+ int32(1),
+ ForKind::Default,
+ eval(11),
+ None,
+ UNKNOWN_LOC,
+ )
+ .unwrap();
+ let branch =
+ If::new(op("prim._OpLT", &x, int32(4)), body, Some(&eval(12)),
UNKNOWN_LOC).unwrap();
+ let root = SeqStmt::new(
+ [Stmt::from(Bind::new(&x, int32(4), UNKNOWN_LOC).unwrap()),
branch.into(), eval(13)],
+ UNKNOWN_LOC,
+ );
+ let retained = (!moved).then(|| root.clone());
+ let mut seen = Rewrite::default();
+ let result = SeqStmt::try_from(mutate(root, &mut
seen).unwrap()).unwrap();
+ assert_eq!(
+ seen.facts.iter().map(|(n, facts)| (*n,
facts.len())).collect::<Vec<_>>(),
+ [(11, 5), (12, 2), (13, 1)]
+ );
+ assert!(seen.restored);
+ assert!(seen.definitions.contains(&DefRegionKind::Simple));
+ let result_binding = result.seq[0].downcast_ref::<BindObj>().unwrap();
+ let result_branch = result.seq[1].downcast_ref::<IfObj>().unwrap();
+ let result_loop =
result_branch.then_case.seq[0].downcast_ref::<ForObj>().unwrap();
+ let binding =
seen.facts[0].1[0].downcast_ref::<BindingConstraintObj>().unwrap();
+ assert!(binding.var.same_as(&result_binding.var));
+ assert!(binding.value.same_as(&result_binding.value));
+
assert_eq!(binding.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+ let guard =
seen.facts[0].1[1].downcast_ref::<PredicateConstraintObj>().unwrap();
+ assert!(guard.condition.same_as(&result_branch.condition));
+ let range =
seen.facts[0].1[2].downcast_ref::<RangeConstraintObj>().unwrap();
+ assert!(range.var.same_as(&result_loop.var));
+
assert_eq!(range.range.begin.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(0));
+
assert_eq!(range.range.extent.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+
assert_eq!(result_loop.end.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+ assert!(seen.facts[0].1[0].same_as(&seen.facts[2].1[0]));
+ if let Some(retained) = retained {
+ let original = retained.seq[0].downcast_ref::<BindObj>().unwrap();
+
assert_eq!(original.value.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(4));
+ }
+ }
+}
+
+#[test]
+fn alias_acquired_in_callback_forces_copying_default_descent() {
+ let _serial = serial();
+ #[derive(Default)]
+ struct Alias {
+ saved: Option<Any>,
+ }
+ #[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>), Result<Any>)]
+ impl Alias {
+ fn mutate_for(
+ &mut self,
+ node: MutationView<'_, For>,
+ view: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ self.saved = Some(node.as_any_view().into());
+ view.default_mutate(self, node).map(Into::into)
+ }
+ fn mutate_integer(
+ &mut self,
+ node: MutationView<'_, IntImm>,
+ _: &mut View<'_, Self>,
+ ) -> Result<Any> {
+ Ok(if node.value.to_i64() == Some(4) { int32(8).into() } else {
Unchanged.into() })
+ }
+ }
+ let mut alias = Alias::default();
+ let i = Var::new("aliased_i", PrimType::int(32), UNKNOWN_LOC);
+ let root =
+ For::new(&i, int32(0), int32(4), int32(1), ForKind::Default, eval(11),
None, UNKNOWN_LOC)
+ .unwrap();
+ let result = For::try_from(mutate(root, &mut alias).unwrap()).unwrap();
+ let saved = alias.saved.unwrap();
+ let saved = saved.downcast_ref::<ForObj>().unwrap();
+ assert_eq!(saved.end.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(4));
+ assert_eq!(result.end.downcast_ref::<IntImmObj>().unwrap().value.to_i64(),
Some(8));
+ assert!(!std::ptr::eq(saved, &*result));
+}
diff --git a/rust/tests/tirx/structural.rs
b/rust/tests/tirx/structural_visit_with_constraint_context.rs
similarity index 62%
rename from rust/tests/tirx/structural.rs
rename to rust/tests/tirx/structural_visit_with_constraint_context.rs
index 554e5836a..7279bee99 100644
--- a/rust/tests/tirx/structural.rs
+++ b/rust/tests/tirx/structural_visit_with_constraint_context.rs
@@ -20,17 +20,13 @@
//! 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::sym::{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,
+ RUNTIME_ERROR,
};
use super::int32;
@@ -82,97 +78,6 @@ fn variable(name: &str) -> Var {
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();
@@ -284,61 +189,6 @@ fn interrupts_restore_native_scope_and_errors_terminate() {
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]
diff --git a/rust/tests/tirx/structural_walk_with_constraint_context.rs
b/rust/tests/tirx/structural_walk_with_constraint_context.rs
new file mode 100644
index 000000000..20f71c35c
--- /dev/null
+++ b/rust/tests/tirx/structural_walk_with_constraint_context.rs
@@ -0,0 +1,187 @@
+/*
+ * 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.
+ */
+
+//! Typed constraint-aware structural walk callbacks and analyzer scopes.
+
+use tvm_tirx::ir::*;
+use tvm_tirx::sym::{
+ Analyzer, BindingConstraintObj, ConstraintContext, PredicateConstraintObj,
RangeConstraintObj,
+};
+use
tvm_tirx::tirx::structural_walk_with_constraint_context::structural_walk_with_constraint_context
as walk;
+use tvm_tirx::tvm_ffi::{
+ AnyView, DefRegionKind, Error, ObjectRefCore, Result, WalkOrder,
WalkResult, RUNTIME_ERROR,
+};
+use super::int32;
+use crate::common::serial;
+
+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)
+}
+
+#[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));
+ }
+ 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);
+ // 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"));
+ }
+}
diff --git a/src/tirx/ir/structural_constraints.cc
b/src/tirx/ir/structural_constraints.cc
index 97d3998d0..d03c1dbae 100644
--- a/src/tirx/ir/structural_constraints.cc
+++ b/src/tirx/ir/structural_constraints.cc
@@ -1028,7 +1028,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("tirx.StructuralMutatorWithConstraintContext",
[] {
return StructuralMutatorWithConstraintContext(
-
ffi::make_object<StructuralMutatorWithConstraintContextObj>());
+ ffi::make_object<
+
ffi::StructuralMutateEngine<StructuralMutatorWithConstraintContextObj>>());
})
.def("tirx.StructuralVisitorWithConstraintContextDefaultVisit",
[](StructuralVisitorWithConstraintContext visitor, ffi::AnyView
value) {
