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 24173d5ce [Rust] Update TVM FFI and use Result<UnchangedOr<Any>> map
and mutate results (#194)
24173d5ce is described below
commit 24173d5ce4662fb12a2b78686704b39af4c27828
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Nov 4 16:22:34 2026 -0500
[Rust] Update TVM FFI and use Result<UnchangedOr<Any>> map and mutate
results (#194)
This PR updates TVM FFI to include apache/tvm-ffi#925, which makes
`Result<UnchangedOr<Any>>` the result of Rust structural map and mutate
dispatchers, and migrates TIRx to it.
- The constraint-aware map and mutate adapters require dispatchers
producing `Option<Result<UnchangedOr<Any>>>`. The map adapter returns
recursion results directly and checks for the unchanged marker with
`is_unchanged()`.
- The constraint-aware mutation test, the IR guide and extension guide
examples, and their docs use `Result<UnchangedOr<Any>>` handlers.
Handlers return `default_mutate` directly and keep a value with
`Unchanged.into()`.
- The TVM FFI update also brings in apache/tvm-ffi#920, #922, and #924.
---
3rdparty/tvm-ffi | 2 +-
docs/advanced/defining_ir_extensions.rst | 6 +--
docs/ir_guides/working_with_ir_in_rust.rst | 6 ++-
examples/extension_guides/rust/src/lib.rs | 10 ++--
examples/ir_guides/rust/src/lib.rs | 14 +++---
.../tirx/structural_map_with_constraint_context.rs | 53 +++++++++++-----------
.../structural_mutate_with_constraint_context.rs | 6 +--
.../structural_mutate_with_constraint_context.rs | 28 ++++++------
8 files changed, 64 insertions(+), 61 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 3e906bc68..2e7dcacb6 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 3e906bc685ea45624bc80e8d7b7d86497d492af3
+Subproject commit 2e7dcacb6f21dcc3cc9603493f06f3760d9fa7e8
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index 7af329cac..d728efd77 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -292,9 +292,9 @@ The rewrites are:
impl Lowering {
fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> { ... }
}
- #[dispatch(map, (AnyView<'_>,), Result<Any>)]
+ #[dispatch(map, (AnyView<'_>,), Result<UnchangedOr<Any>>)]
impl Lowering {
- fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
+ fn map_sequence(&mut self, seq: SeqStmt) ->
Result<UnchangedOr<Any>> { ... }
:append: }
.. group-tab:: Python
@@ -309,7 +309,7 @@ and ``parts`` reads the real and imaginary parts of an
operand, either a
validates a call to one of the operators, then applies the three rules;
``map_bind`` splits a complex binding, and ``map_sequence`` flattens the
sequence that a split binding leaves inside another. In Rust,
-``#[dispatch(map, (AnyView<'_>,), Result<Any>)]`` sends each node to the
+``#[dispatch(map, (AnyView<'_>,), Result<UnchangedOr<Any>>)]`` sends each node
to the
``map_`` method of its type. In Python, the next excerpt pairs each method with
the type of node it takes.
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 98bc11fdb..cd8db1c46 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -645,8 +645,10 @@ The output is:
``mutate_for`` records the replacement for its loop variable, then
``mutator.default_mutate(self, value)`` rewrites the loop's children
through the same struct, so ``mutate_var`` sees the replacement before any
-use of ``i``. ``default_mutate`` returns ``UnchangedOr<Any>``, and
-``.map(Into::into)`` turns it into the handler's ``Any`` result.
+use of ``i``. Mutate handlers return ``Result<UnchangedOr<Any>>``, the
+same type that ``default_mutate`` returns, so ``mutate_for`` returns its
+result directly. ``mutate_var`` returns the new index with ``.into()``, or
+``Unchanged`` to keep the variable.
A mutate handler receives its node as a ``MutationView``, here
``MutationView<'_, For>``. We read the loop's fields through it, as with
diff --git a/examples/extension_guides/rust/src/lib.rs
b/examples/extension_guides/rust/src/lib.rs
index 94c5e8f3d..295ce1b2a 100644
--- a/examples/extension_guides/rust/src/lib.rs
+++ b/examples/extension_guides/rust/src/lib.rs
@@ -31,7 +31,7 @@ use tvm_tirx::tirx::transform::FunctionPass;
use tvm_tirx::tirx::Function;
use tvm_tirx::tvm_ffi::{
self, def_object_ref, dispatch, rs_object, Any, AnyView, Error,
ObjectRefCast, ObjectRefCore,
- Result, Unchanged, WalkOrder,
+ Result, Unchanged, UnchangedOr, WalkOrder,
};
// begin type
@@ -192,14 +192,14 @@ impl Lowering {
}
}
-#[dispatch(map, (AnyView<'_>,), Result<Any>)]
+#[dispatch(map, (AnyView<'_>,), Result<UnchangedOr<Any>>)]
impl Lowering {
- fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> {
+ fn map_sequence(&mut self, seq: SeqStmt) -> Result<UnchangedOr<Any>> {
Ok(SeqStmt::new(seq.seq.iter(), seq.loc.clone()).into())
}
// begin lowering
- fn map_call(&mut self, call: Call) -> Result<Any> {
+ fn map_call(&mut self, call: Call) -> Result<UnchangedOr<Any>> {
if is_complex_op(&call.op)? {
validate_call(&call)?;
}
@@ -227,7 +227,7 @@ impl Lowering {
Ok(call.into())
}
- fn map_bind(&mut self, bind: Bind) -> Result<Any> {
+ fn map_bind(&mut self, bind: Bind) -> Result<UnchangedOr<Any>> {
if let Some(ty) = bind.var.ty.downcast_ref::<ComplexTypeObj>() {
let (re, im) = self.parts(&bind.value)?;
let name = bind.var.name.as_str();
diff --git a/examples/ir_guides/rust/src/lib.rs
b/examples/ir_guides/rust/src/lib.rs
index a7f6df938..fed3f2868 100644
--- a/examples/ir_guides/rust/src/lib.rs
+++ b/examples/ir_guides/rust/src/lib.rs
@@ -37,7 +37,7 @@ use tvm_tirx::tirx::Function;
use tvm_tirx::tvm_ffi::{
self, dispatch, match_any, Any, AnyView, Array, DefRegionKind,
MutationView, ObjectIdentity,
ObjectRefCast, ObjectRefCore, Result, StructuralMutatorView,
StructuralVisitorView, Unchanged,
- VisitInterrupt, WalkOrder, WalkResult,
+ UnchangedOr, VisitInterrupt, WalkOrder, WalkResult,
};
/// A Python list of strings.
@@ -283,27 +283,27 @@ struct ReverseLoops {
reversed_index: HashMap<ObjectIdentity, Expr>,
}
-#[dispatch(mutate, (MutationView<'_>, &mut StructuralMutatorView<'_, Self>),
Result<Any>)]
+#[dispatch(mutate, (MutationView<'_>, &mut StructuralMutatorView<'_, Self>),
Result<UnchangedOr<Any>>)]
impl ReverseLoops {
fn mutate_for(
&mut self,
value: MutationView<'_, For>,
mutator: &mut StructuralMutatorView<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
let index = (&value.end - 1 - &value.var).try_into()?;
self.reversed_index.insert(ObjectIdentity::of(&value.var), index);
- mutator.default_mutate(self, value).map(Into::into)
+ mutator.default_mutate(self, value)
}
fn mutate_var(
&mut self,
var: Var,
mutator: &mut StructuralMutatorView<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
let is_use = mutator.def_region_kind() == DefRegionKind::None;
Ok(match self.reversed_index.get(&ObjectIdentity::of(&var)) {
Some(index) if is_use => index.clone().into(),
- _ => var.into(),
+ _ => Unchanged.into(),
})
}
}
@@ -335,7 +335,7 @@ fn remove_redundant_guards(func: Function) ->
Result<Function> {
let analyzer = Analyzer::new();
structural_map_with_constraint_context(
func,
- |branch: If, constraints: &ConstraintContext| -> Result<Any> {
+ |branch: If, constraints: &ConstraintContext| ->
Result<UnchangedOr<Any>> {
analyzer.sync_constraint_context(constraints)?;
Ok(if analyzer.can_prove(branch.condition.clone(),
ProofStrength::Default)? {
branch.then_case.clone().into()
diff --git a/rust/src/tirx/structural_map_with_constraint_context.rs
b/rust/src/tirx/structural_map_with_constraint_context.rs
index 2d5e721ec..4f24d6640 100644
--- a/rust/src/tirx/structural_map_with_constraint_context.rs
+++ b/rust/src/tirx/structural_map_with_constraint_context.rs
@@ -26,11 +26,9 @@ use std::borrow::BorrowMut;
use tvm_ffi::{
Any, AnyView, DefRegionKind, Dispatcher, IntoDispatcher, MutationView,
Object, Result,
- WalkOrder,
+ UnchangedOr, WalkOrder,
};
-use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
-
use super::structural_mutate_with_constraint_context::{
structural_mutate_with_constraint_context,
StructuralMutatorWithConstraintContextView,
};
@@ -46,11 +44,13 @@ impl<'v, 'c, 'r, 'a, D, const PRE: bool>
MutationView<'v>,
&'c mut StructuralMutatorWithConstraintContextView<'r,
MapMutator<'a, D, PRE>>,
),
- Option<Result<Any>>,
+ Option<Result<UnchangedOr<Any>>>,
> for &mut MapMutator<'a, D, PRE>
where
- for<'s, 'b, 'ctx> &'s mut D:
- Dispatcher<(AnyView<'b>, &'ctx ConstraintContext, DefRegionKind),
Option<Result<Any>>>,
+ for<'s, 'b, 'ctx> &'s mut D: Dispatcher<
+ (AnyView<'b>, &'ctx ConstraintContext, DefRegionKind),
+ Option<Result<UnchangedOr<Any>>>,
+ >,
{
#[inline(always)]
fn dispatch_tuple(
@@ -59,22 +59,24 @@ where
MutationView<'v>,
&'c mut StructuralMutatorWithConstraintContextView<'r,
MapMutator<'a, D, PRE>>,
),
- ) -> Option<Result<Any>> {
+ ) -> Option<Result<UnchangedOr<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>>>,
+ for<'s, 'v, 'c> &'s mut D: Dispatcher<
+ (AnyView<'v>, &'c ConstraintContext, DefRegionKind),
+ Option<Result<UnchangedOr<Any>>>,
+ >,
{
#[inline(always)]
fn map(
&mut self,
value: MutationView<'_>,
mutator: &mut StructuralMutatorWithConstraintContextView<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
let kind = mutator.def_region_kind();
if PRE {
let mode = value.inplace_mode();
@@ -88,12 +90,13 @@ where
))
};
let Some(result) = result else {
- return mutator.default_mutate(self, value).map(Into::into);
+ return mutator.default_mutate(self, value);
};
- let mut mapped = result?;
- if mapped.type_index() == TVMFFITypeIndex::kTVMFFIUnchanged as i32
{
- return mutator.default_mutate(self, value).map(Into::into);
+ let mapped = result?;
+ if mapped.is_unchanged() {
+ return mutator.default_mutate(self, value);
}
+ let mut mapped: Any = mapped.into();
let same_object = value
.downcast_ref::<Object>()
.zip(mapped.downcast_ref::<Object>())
@@ -101,13 +104,13 @@ where
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);
+ return mutator.default_mutate(self, value);
}
// 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))
+ Ok(UnchangedOr::changed(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
@@ -118,9 +121,7 @@ where
// 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) };
+ let current = if mapped.is_unchanged() { original } else {
AnyView::from(&mapped) };
// Scoped native hooks have restored the incoming context here.
let result = unsafe {
self.mapper.dispatch_tuple((
@@ -132,11 +133,7 @@ where
match result {
Some(result) => {
let result = result?;
- Ok(if result.type_index() ==
TVMFFITypeIndex::kTVMFFIUnchanged as i32 {
- mapped
- } else {
- result
- })
+ Ok(if result.is_unchanged() { mapped } else { result })
}
None => Ok(mapped),
}
@@ -169,13 +166,15 @@ where
O: BorrowMut<D>,
for<'v, 'c> H: IntoDispatcher<
(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
- Option<Result<Any>>,
+ Option<Result<UnchangedOr<Any>>>,
M,
Target = D,
Output = O,
>,
- for<'s, 'v, 'c> &'s mut D:
- Dispatcher<(AnyView<'v>, &'c ConstraintContext, DefRegionKind),
Option<Result<Any>>>,
+ for<'s, 'v, 'c> &'s mut D: Dispatcher<
+ (AnyView<'v>, &'c ConstraintContext, DefRegionKind),
+ Option<Result<UnchangedOr<Any>>>,
+ >,
{
let mut storage = mapper.into_dispatcher();
let mapper = storage.borrow_mut();
diff --git a/rust/src/tirx/structural_mutate_with_constraint_context.rs
b/rust/src/tirx/structural_mutate_with_constraint_context.rs
index 88ddb50e1..a921d43cc 100644
--- a/rust/src/tirx/structural_mutate_with_constraint_context.rs
+++ b/rust/src/tirx/structural_mutate_with_constraint_context.rs
@@ -33,7 +33,7 @@ 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,
+ ObjectRefCore, Result, StructuralMutatorObj, StructuralMutatorView,
UnchangedOr,
};
use super::StructuralMutatorWithConstraintContextObj;
@@ -152,14 +152,14 @@ where
O: BorrowMut<D>,
for<'v, 'c, 'r> H: IntoDispatcher<
(MutationView<'v>, &'c mut
StructuralMutatorWithConstraintContextView<'r, D>),
- Option<Result<Any>>,
+ Option<Result<UnchangedOr<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>>,
+ Option<Result<UnchangedOr<Any>>>,
>,
{
StructuralMutateEngine::<StructuralMutatorWithConstraintContextBase>::mutate_root(
diff --git a/rust/tests/tirx/structural_mutate_with_constraint_context.rs
b/rust/tests/tirx/structural_mutate_with_constraint_context.rs
index 88f674c18..38158e6f7 100644
--- a/rust/tests/tirx/structural_mutate_with_constraint_context.rs
+++ b/rust/tests/tirx/structural_mutate_with_constraint_context.rs
@@ -25,7 +25,9 @@ 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 tvm_tirx::tvm_ffi::{
+ dispatch, Any, DefRegionKind, MutationView, ObjectRefCore, Result,
Unchanged, UnchangedOr,
+};
use super::{int32, op};
use crate::common::serial;
@@ -41,13 +43,13 @@ struct Rewrite {
restored: bool,
}
-#[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>), Result<Any>)]
+#[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>),
Result<UnchangedOr<Any>>)]
impl Rewrite {
fn mutate_integer(
&mut self,
node: MutationView<'_, IntImm>,
_: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
if node.value.to_i64().unwrap() == 4 {
Ok(int32(8).into())
} else {
@@ -58,28 +60,28 @@ impl Rewrite {
&mut self,
node: MutationView<'_, Var>,
view: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
self.definitions.push(view.def_region_kind());
- view.default_mutate(self, node).map(Into::into)
+ view.default_mutate(self, node)
}
fn mutate_evaluate(
&mut self,
node: MutationView<'_, Evaluate>,
view: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<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)
+ view.default_mutate(self, node)
}
fn mutate_for(
&mut self,
node: MutationView<'_, For>,
view: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<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)
+ result
}
}
@@ -144,21 +146,21 @@ fn
alias_acquired_in_callback_forces_copying_default_descent() {
struct Alias {
saved: Option<Any>,
}
- #[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>), Result<Any>)]
+ #[dispatch(mutate, (MutationView<'_>, &mut View<'_, Self>),
Result<UnchangedOr<Any>>)]
impl Alias {
fn mutate_for(
&mut self,
node: MutationView<'_, For>,
view: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
self.saved = Some(node.as_any_view().into());
- view.default_mutate(self, node).map(Into::into)
+ view.default_mutate(self, node)
}
fn mutate_integer(
&mut self,
node: MutationView<'_, IntImm>,
_: &mut View<'_, Self>,
- ) -> Result<Any> {
+ ) -> Result<UnchangedOr<Any>> {
Ok(if node.value.to_i64() == Some(4) { int32(8).into() } else {
Unchanged.into() })
}
}
