This is an automated email from the ASF dual-hosted git repository.
Seven-Streams 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 2025a90d4 [Rust] Update FFI and migrate structural dispatch (#170)
2025a90d4 is described below
commit 2025a90d459514fcb649705ecd73777824504bf2
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Nov 3 09:25:51 2026 -0500
[Rust] Update FFI and migrate structural dispatch (#170)
Update the TVM-FFI submodule to the current Rust dispatch APIs. Migrate
structural traversal examples and the IR type test to the shared
dispatcher, keep traversal state on each dispatcher, and use scoped
visitor/mutator views for recursive descent. Align example dependencies
and guide snippets with the pinned FFI source.
---
3rdparty/tvm-ffi | 2 +-
docs/advanced/defining_ir_extensions.rst | 6 +-
docs/ir_guides/working_with_ir_in_rust.rst | 62 +++++++++++----------
examples/extension/Cargo.toml | 2 +
examples/extension/src/lib.rs | 22 +++++---
examples/ir_guides/rust/Cargo.toml | 2 +
examples/ir_guides/rust/src/lib.rs | 79 +++++++++++++++------------
examples/ir_guides/working_with_ir_in_rust.py | 2 +
rust/tests/ir/type.rs | 17 ++----
9 files changed, 104 insertions(+), 90 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 3380316d0..b0a048e4d 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 3380316d02864ff51513a5edd353b6d1f2e86d40
+Subproject commit b0a048e4d510febff28f1ecea102716516c5b952
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index 50458ea45..b9f2a4709 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -180,7 +180,7 @@ once per run and keeps them. ``Lowering`` implements the
rewrites with them:
}
#[dispatch(map)]
impl Lowering {
- fn map_sequence(&mut self, seq: SeqStmt) -> SeqStmt { ... }
+ fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
:append: }
``ops`` holds the four operators, ``bindings`` records each complex variable
@@ -189,7 +189,9 @@ parts of an operand, either a ``complex(re, im)`` call or a
variable already
split. ``#[dispatch(map)]`` sends each node to the ``map_`` method of its type,
and leaves other nodes unchanged: ``map_call`` applies the three rules,
``map_bind`` splits a complex binding, and ``map_sequence`` flattens the
-sequence that a split binding leaves inside another.
+sequence that a split binding leaves inside another. The value-only callback
+passes each node to ``lowering.dispatch`` and returns ``Unchanged`` when no
+handler matches; these rewrites do not need a definition-region argument.
The pass is built from an ordinary named function:
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 65fa4a086..ba8013e21 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -317,7 +317,10 @@ the default behavior.
``#[dispatch(walk)]`` on an ``impl`` block turns its ``walk_*`` methods
into callbacks for |structural_walk|_, which visits every node. Each
method borrows the nodes of its parameter type, and the methods are tried
-in source order. Let's count the variables and additions:
+in source order. Walk handlers share the ``Result<WalkResult>`` return type;
+one handler declares the trailing ``DefRegionKind`` argument used by the
+walker, and handlers that do not need it may omit it. Let's count the
+variables and additions:
.. literalinclude:: ../../examples/ir_guides/rust/src/lib.rs
:language: rust
@@ -518,17 +521,16 @@ are structurally equal.
Visitors and Mutators in Rust
-----------------------------
-Stateful traversals use the same four structural utilities. The state
-lives either on a |dispatch|_ struct, as in Dispatching on Node Types, or
-in |VisitCallbacks|_ and |MutateCallbacks|_, which attach it to a tuple of
-callbacks. ``#[dispatch(map)]``, ``#[dispatch(visit)]``, and
-``#[dispatch(mutate)]`` work like ``#[dispatch(walk)]`` for ``map_*``,
-``visit_*``, and ``mutate_*`` methods.
+Stateful traversals use the same four structural utilities. State lives
+on a |dispatch|_ struct, as in Dispatching on Node Types.
+``#[dispatch(map)]`` and ``#[dispatch(visit)]`` group ``map_*`` and
+``visit_*`` methods. ``#[dispatch(mutate, MutationView)]`` groups
+``mutate_*`` methods and carries mutation permission through matching.
**A visit or mutate callback takes control of its node's children.**
-Let's record the loops around each write. With ``VisitCallbacks``, each
-callback receives the node and a ``VisitContext``, which holds the shared
-state:
+Let's record the loops around each write. Each handler receives the node
+and can take a ``StructuralVisitorView<'_, Self>`` for recursive traversal;
+the loop stack and collected writes remain on ``self``:
.. literalinclude:: ../../examples/ir_guides/rust/src/lib.rs
:language: rust
@@ -541,20 +543,20 @@ The output is:
['B inside [i]']
-``visit_for`` pushes the loop variable, lets ``visit_children`` visit the
-loop's children with their usual traversal, then pops the variable. It
-returns the result of ``visit_children``, so that an interrupt from a child
-reaches the caller. The context also provides ``visit`` for a selected
-child and ``def_region_kind`` for the current node.
+``visit_for`` pushes the loop variable, lets ``default_visit(self, loop_)``
+visit the loop's children with their usual traversal, then pops the variable.
+It returns the traversal result, so an interrupt or error from a child
+reaches the caller. The view also provides ``visit(self, child)`` for a
+selected child and ``def_region_kind`` for the current node.
-For the mutation we use the struct form instead; either form works for
-both visits and mutations. A ``#[dispatch(mutate)]`` method that takes a
-trailing ``&mut Mutator`` gets the mutation counterpart of
-``VisitContext``, with ``def_region_kind()`` and child-mutation methods. A
+For mutation, a trailing ``&mut StructuralMutatorView<'_, Self>`` provides
+``def_region_kind()`` and child-mutation methods. A
``MutationView<'_, For>`` carries the current loop and the engine's
-permission to reuse it. ``as_any_view()`` borrows it for inspection; after
-that borrow ends, ``default_maybe_inplace_mutate`` consumes the handle to
-recurse with the same permission. Reuse still requires unique ownership,
+permission to reuse it. Its typed fields can be borrowed directly; after
+that borrow ends, ``default_mutate`` consumes the handle to recurse with
+the same permission. The returned ``UnchangedOr<Any>`` becomes the common
+``Any`` handler result without discarding an unchanged marker. Reuse still
+requires unique ownership,
so an alias held by Python remains unchanged. Let's reverse the loop
by replacing each use of ``i`` with ``n - 1 - i``. As before,
``.try_into()?`` extracts the built expression or returns its error:
@@ -578,7 +580,7 @@ The output is:
B[n - 1 - i] = A[n - 1 - i] + T.float32(1.0)
``mutate_for`` records the replacement, then
-``default_maybe_inplace_mutate`` rewrites the loop's children through the
+``default_mutate`` rewrites the loop's children through the
same struct, so ``mutate_var`` sees the
replacement before any use of ``i``.
@@ -626,10 +628,10 @@ Quick Reference
- Receives the definition region kind
* - ``Function::get_global("ffi.StructuralEqual")``
- Compares by structure
- * - |dispatch|_\ ``(walk)``, ``(map)``, ``(visit)``, ``(mutate)``
+ * - |dispatch|_\ ``(walk)``, ``(map)``, ``(visit)``, ``(mutate,
MutationView)``
- Turns the methods of a struct into typed callbacks
- * - |VisitCallbacks|_, |MutateCallbacks|_
- - Attach state to a tuple of visit or mutate callbacks
+ * - |StructuralVisitorView|_, |StructuralMutatorView|_
+ - Recurse with the same dispatcher and current definition region
* - |static_init_block!|_, |register_global_typed|_
- Registers Rust functions as global functions
@@ -661,10 +663,10 @@ Quick Reference
.. _make_const:
../reference/rust/generated/tvm_tirx/ir/prim/op/fn.make_const.html
.. |dispatch| replace:: ``#[dispatch]``
.. _dispatch: ../reference/rust/generated/tvm_ffi/attr.dispatch.html
-.. |VisitCallbacks| replace:: ``VisitCallbacks``
-.. _VisitCallbacks:
../reference/rust/generated/tvm_ffi/extra/structural_visit/struct.VisitCallbacks.html
-.. |MutateCallbacks| replace:: ``MutateCallbacks``
-.. _MutateCallbacks:
../reference/rust/generated/tvm_ffi/extra/structural_mutate/struct.MutateCallbacks.html
+.. |StructuralVisitorView| replace:: ``StructuralVisitorView``
+.. _StructuralVisitorView:
../reference/rust/generated/tvm_ffi/extra/structural_visit/struct.StructuralVisitorView.html
+.. |StructuralMutatorView| replace:: ``StructuralMutatorView``
+.. _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``
diff --git a/examples/extension/Cargo.toml b/examples/extension/Cargo.toml
index 61beb674c..12ef76e38 100644
--- a/examples/extension/Cargo.toml
+++ b/examples/extension/Cargo.toml
@@ -28,3 +28,5 @@ crate-type = ["cdylib"]
[dependencies]
# Use this checkout so the Rust API and native libraries share the same FFI
revision.
tvm-tirx = { path = "../../rust" }
+# Procedural macros resolve the FFI crate directly; use the same pinned
sources.
+tvm-ffi = { path = "../../3rdparty/tvm-ffi/rust/tvm-ffi" }
diff --git a/examples/extension/src/lib.rs b/examples/extension/src/lib.rs
index c2368002c..f630aae1f 100644
--- a/examples/extension/src/lib.rs
+++ b/examples/extension/src/lib.rs
@@ -25,13 +25,13 @@ use tvm_tirx::ir::sig::{self, PrimExpr};
use tvm_tirx::ir::transform::{Pass, PassContext, PassInfo};
use tvm_tirx::ir::{
dtype, register_type, Bind, Call, CallEffectKind, CallObj, Expr, IRModule,
Op, OpDef, PrimType,
- SeqStmt, Stmt, Type, TypeObj, Var,
+ SeqStmt, Type, TypeObj, Var,
};
use tvm_tirx::tirx::transform::FunctionPass;
use tvm_tirx::tirx::Function;
use tvm_tirx::tvm_ffi::{
- self, def_object_ref, dispatch, rs_class, Error, ObjectRefCast,
ObjectRefCore, Result,
- WalkOrder,
+ self, def_object_ref, dispatch, rs_class, Any, AnyView, Error,
ObjectRefCast, ObjectRefCore,
+ Result, Unchanged, WalkOrder,
};
// begin type
@@ -210,12 +210,12 @@ impl Lowering {
#[dispatch(map)]
impl Lowering {
- fn map_sequence(&mut self, seq: SeqStmt) -> SeqStmt {
- SeqStmt::new(seq.seq.iter(), seq.loc.clone())
+ fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> {
+ Ok(SeqStmt::new(seq.seq.iter(), seq.loc.clone()).into())
}
// begin lowering
- fn map_call(&mut self, call: Call) -> Result<Expr> {
+ fn map_call(&mut self, call: Call) -> Result<Any> {
let ops = &self.ops;
if ops.contains(&call.op) {
validate_call(&call)?;
@@ -238,12 +238,12 @@ impl Lowering {
}
if call.op.same_as(&ops.real) || call.op.same_as(&ops.imag) {
let (re, im) = self.parts(&call.args.get(0)?)?;
- return Ok(if call.op.same_as(&ops.real) { re } else { im });
+ return Ok(if call.op.same_as(&ops.real) { re.into() } else {
im.into() });
}
Ok(call.into())
}
- fn map_bind(&mut self, bind: Bind) -> Result<Stmt> {
+ fn map_bind(&mut self, bind: Bind) -> Result<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();
@@ -298,7 +298,11 @@ fn lower_function(
_ctx: PassContext,
) -> Result<Option<Function>> {
let mut lowering = Lowering { ops: ComplexOps::get()?, bindings:
Vec::new() };
- let rewritten = tvm_ffi::structural_map(func, &mut lowering,
WalkOrder::PostOrder)?;
+ let rewritten = tvm_ffi::structural_map(
+ func,
+ |value: AnyView<'_>| lowering.dispatch(value).unwrap_or_else(||
Ok(Unchanged.into())),
+ WalkOrder::PostOrder,
+ )?;
let ops = &lowering.ops;
let checked = tvm_ffi::structural_map(
rewritten,
diff --git a/examples/ir_guides/rust/Cargo.toml
b/examples/ir_guides/rust/Cargo.toml
index 12fae240a..61a6ba338 100644
--- a/examples/ir_guides/rust/Cargo.toml
+++ b/examples/ir_guides/rust/Cargo.toml
@@ -29,3 +29,5 @@ crate-type = ["cdylib"]
[dependencies]
# Use the bindings from this checkout to match the current IR node fields.
tvm-tirx = { path = "../../../rust" }
+# Procedural macros resolve the FFI crate directly; use the same pinned
sources.
+tvm-ffi = { path = "../../../3rdparty/tvm-ffi/rust/tvm-ffi" }
diff --git a/examples/ir_guides/rust/src/lib.rs
b/examples/ir_guides/rust/src/lib.rs
index c906ab4c6..65d901473 100644
--- a/examples/ir_guides/rust/src/lib.rs
+++ b/examples/ir_guides/rust/src/lib.rs
@@ -31,8 +31,8 @@ use tvm_tirx::ir::{
use tvm_tirx::script::script;
use tvm_tirx::tirx::Function;
use tvm_tirx::tvm_ffi::{
- self, dispatch, match_any, Any, AnyView, Array, DefRegionKind,
MutationView, Mutator,
- ObjectIdentity, ObjectRefCast, ObjectRefCore, Result, VisitCallbacks,
VisitContext,
+ self, dispatch, match_any, Any, AnyView, Array, DefRegionKind,
MutationView, ObjectIdentity,
+ ObjectRefCast, ObjectRefCore, Result, StructuralMutatorView,
StructuralVisitorView,
VisitInterrupt, WalkOrder, WalkResult,
};
@@ -179,14 +179,14 @@ struct Counter {
#[dispatch(walk)]
impl Counter {
- fn walk_var(&mut self, _var: &VarObj) -> WalkResult {
+ fn walk_var(&mut self, _var: &VarObj, _kind: DefRegionKind) ->
Result<WalkResult> {
self.vars += 1;
- WalkResult::Advance
+ Ok(WalkResult::Advance)
}
- fn walk_add(&mut self, _add: &AddObj) -> WalkResult {
+ fn walk_add(&mut self, _add: &AddObj) -> Result<WalkResult> {
self.adds += 1;
- WalkResult::Advance
+ Ok(WalkResult::Advance)
}
}
// [counter.end]
@@ -207,29 +207,32 @@ struct LoopNest {
writes: Vec<String>,
}
-fn visit_for(
- loop_: &ForObj,
- visitor: &mut VisitContext<'_, LoopNest>,
-) -> Result<Option<VisitInterrupt>> {
- visitor.state_mut().loops.push(loop_.var.name.as_str().to_owned());
- let result = visitor.visit_children()?;
- visitor.state_mut().loops.pop();
- Ok(result)
-}
+#[dispatch(visit)]
+impl LoopNest {
+ fn visit_for(
+ &mut self,
+ loop_: &ForObj,
+ visitor: &mut StructuralVisitorView<'_, Self>,
+ ) -> Result<Option<VisitInterrupt>> {
+ self.loops.push(loop_.var.name.as_str().to_owned());
+ let result = visitor.default_visit(self, loop_);
+ self.loops.pop();
+ result
+ }
-fn visit_store(store: &TensorStoreObj, visitor: &mut VisitContext<'_,
LoopNest>) -> Result<()> {
- let loops = visitor.state().loops.join(", ");
- let write = format!("{} inside [{}]", script(&store.dest)?.as_str(),
loops);
- visitor.state_mut().writes.push(write);
- Ok(())
+ fn visit_store(&mut self, store: &TensorStoreObj) ->
Result<Option<VisitInterrupt>> {
+ let loops = self.loops.join(", ");
+ let write = format!("{} inside [{}]", script(&store.dest)?.as_str(),
loops);
+ self.writes.push(write);
+ Ok(None)
+ }
}
fn loops_around_writes(func: Function) -> Result<Strings> {
- let mut visitor = VisitCallbacks::new(LoopNest::default(), (visit_for,
visit_store));
+ let mut visitor = LoopNest::default();
tvm_ffi::structural_visit(&func, &mut visitor)?;
- let state = visitor.into_state();
- assert!(state.loops.is_empty());
- Ok(state.writes.iter().map(tvm_ffi::String::from).collect())
+ assert!(visitor.loops.is_empty());
+ Ok(visitor.writes.iter().map(tvm_ffi::String::from).collect())
}
// [visit.end]
@@ -239,22 +242,28 @@ struct ReverseLoops {
reversed_index: HashMap<ObjectIdentity, Expr>,
}
-#[dispatch(mutate)]
+#[dispatch(mutate, MutationView)]
impl ReverseLoops {
- fn mutate_for(&mut self, value: MutationView<'_, For>, mutator: &mut
Mutator) -> Result<Any> {
- let view = value.as_any_view();
- let loop_ = view.downcast_ref::<ForObj>().unwrap();
- let index = (&loop_.end - 1 - &loop_.var).try_into()?;
- self.reversed_index.insert(ObjectIdentity::of(&loop_.var), index);
- mutator.default_maybe_inplace_mutate(self, value)
+ fn mutate_for(
+ &mut self,
+ value: MutationView<'_, For>,
+ mutator: &mut StructuralMutatorView<'_, Self>,
+ ) -> Result<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)
}
- fn mutate_var(&mut self, var: Var, mutator: &mut Mutator) -> Expr {
+ fn mutate_var(
+ &mut self,
+ var: Var,
+ mutator: &mut StructuralMutatorView<'_, Self>,
+ ) -> Result<Any> {
let is_use = mutator.def_region_kind() == DefRegionKind::None;
- match self.reversed_index.get(&ObjectIdentity::of(&var)) {
- Some(index) if is_use => index.clone(),
+ Ok(match self.reversed_index.get(&ObjectIdentity::of(&var)) {
+ Some(index) if is_use => index.clone().into(),
_ => var.into(),
- }
+ })
}
}
diff --git a/examples/ir_guides/working_with_ir_in_rust.py
b/examples/ir_guides/working_with_ir_in_rust.py
index f034249b1..6a0a341bc 100644
--- a/examples/ir_guides/working_with_ir_in_rust.py
+++ b/examples/ir_guides/working_with_ir_in_rust.py
@@ -87,7 +87,9 @@ assert
list(tirx_ir_guide_examples.rust.compare_specialized(add_one, n)) == [Fal
print(list(tirx_ir_guide_examples.rust.loops_around_writes(add_one)))
assert list(tirx_ir_guide_examples.rust.loops_around_writes(add_one)) == ["B
inside [i]"]
+original_script = add_one.script()
add_one_reversed = tirx_ir_guide_examples.rust.reverse_loops(add_one)
print(add_one_reversed.script())
assert "B[n - 1 - i] = A[n - 1 - i] + T.float32(1.0)" in
add_one_reversed.script()
assert add_one_reversed.body.seq[0].var.same_as(add_one.body.seq[0].var)
+assert add_one.script() == original_script # the Python-held input remains
unchanged
diff --git a/rust/tests/ir/type.rs b/rust/tests/ir/type.rs
index adc3f0036..9e0621511 100644
--- a/rust/tests/ir/type.rs
+++ b/rust/tests/ir/type.rs
@@ -27,10 +27,10 @@ use tvm_tirx::ir::{
UNKNOWN_LOC, Location, dtype, Evaluate, FuncType, MissingType, NoneType,
PrimType, PtrType,
TensorRegionType, TupleType, Type, TypeObj, Var, VarObj,
};
-use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class, structural_map};
+use tvm_tirx::tvm_ffi::{def_object_ref, rs_class, structural_map};
use tvm_tirx::tvm_ffi::{
- Any, AnyView, Array, DLDataType, DLDataTypeExt, DefRegionKind,
FieldGetter, MapDispatch,
- ObjectCore, ObjectRefCore, Result, ShallowCopy, String as FfiString,
WalkOrder,
+ Any, AnyView, Array, DLDataType, DLDataTypeExt, FieldGetter, ObjectCore,
ObjectRefCore, Result,
+ ShallowCopy, String as FfiString, WalkOrder,
};
use super::{scope};
@@ -259,22 +259,13 @@ fn
rust_type_cpp_deep_copy_goes_through_the_shallow_copy_and_the_setters() {
assert_eq!(copy.note.as_str(), "seven");
}
-/// Adds one to every integer leaf.
-struct IncrementIntegers;
-
-impl MapDispatch for IncrementIntegers {
- fn dispatch_map(&mut self, value: AnyView<'_>, _: DefRegionKind) ->
Option<Result<Any>> {
- value.try_as::<i64>().map(|integer| Ok(Any::from(integer + 1)))
- }
-}
-
#[test]
fn rust_type_rust_engines_descend_through_the_fields() {
let _serial = serial();
let value = int_float_pair(7, "seven");
// Shared with `value`, so the engine copies the node through
`__ffi_shallow_copy__`.
let mapped =
- structural_map(value.clone(), &mut IncrementIntegers,
WalkOrder::PostOrder).unwrap();
+ structural_map(value.clone(), |value: i64| value + 1,
WalkOrder::PostOrder).unwrap();
let mapped = PairType::try_from(mapped).unwrap();
assert!(!mapped.same_as(&value));
assert_eq!(mapped.tag, 8);
