This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 c8d0ae636 [DOCS][RUST] Polish the extension tutorial and example (#176)
c8d0ae636 is described below
commit c8d0ae636c265bb5e03a427b6f2792f41629c172
Author: Linzhang Li <[email protected]>
AuthorDate: Tue Nov 3 13:09:06 2026 -0500
[DOCS][RUST] Polish the extension tutorial and example (#176)
---
AGENTS.md | 2 +-
docs/advanced/defining_ir_extensions.rst | 85 +++++++++++++++-----------------
docs/advanced/extending_ir_in_rust.rst | 31 ++++++------
docs/advanced/using_ir_extensions.rst | 2 +-
docs/reference/rust/index.rst | 8 +--
examples/extension/Cargo.toml | 3 +-
examples/extension/src/lib.rs | 65 +++++++++---------------
rust/Cargo.toml | 2 +-
rust/src/lib.rs | 5 +-
tests/python/ir/test_op_rust.py | 2 +-
10 files changed, 94 insertions(+), 111 deletions(-)
diff --git a/AGENTS.md b/AGENTS.md
index 7e2523ebe..2d0a31d25 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -34,7 +34,7 @@ harnesses on top. The Python package is `tvm_tirx`; the
distribution is
include/tvm/ C++ public headers (ir, sym, tirx, target, runtime,
backends)
src/ C++ compiler and runtime implementations
python/tvm_tirx/ Python package and reflected C++ bindings
-rust/ Rust bindings (stubgen)
+rust/ Rust bindings (stubgen) and extension tools
tests/cpp/ GoogleTest C++ tests
tests/python/ pytest tests organized by owning namespace and backend
tests/lint/ ASF header and file-type checks
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index b9f2a4709..ba65df853 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -27,8 +27,9 @@ the validation and traversal helpers they use.
Defining a Type
---------------
-``ComplexType`` has one field, ``elem``, the scalar type shared by the real and
-imaginary parts:
+To work with complex numbers directly in the IR, we first introduce a type for
+them. ``ComplexType`` has one field, ``dtype``, the scalar type shared by the
+real and imaginary parts, as in ``TensorType``:
.. literalinclude:: ../../examples/extension/src/lib.rs
:language: rust
@@ -50,18 +51,14 @@ reference. Its arguments are:
as the result of type inference.
- ``ComplexTypeObj``: the object it refers to. A ``ComplexTypeObj`` converts
into a ``ComplexType`` with ``into()``, and a ``ComplexType`` dereferences to
- it, so ``ty.elem`` reads the field.
-
-``infer_complex`` in the next section uses both conversions: it builds a
-``ComplexTypeObj``, turns it into a ``ComplexType``, and returns it as a
-``Type``.
+ it, so ``ty.dtype`` reads the field.
The ``static_init_block!`` runs ``register_type`` when the library loads, as a
``TVM_FFI_STATIC_INIT_BLOCK`` does for C++, and Python then sees the type as
``my_ext.Complex``.
A ``ComplexType`` records the scalar type of the two parts, and a
-``complex(re, im)`` call supplies their values as its operands. ``elem`` can be
+``complex(re, im)`` call supplies their values as its operands. ``dtype`` can
be
any scalar float type, such as ``float32``, ``float64`` or ``bfloat16``; this
tutorial uses ``float32``.
@@ -71,7 +68,7 @@ Defining Operators
To support operations on the new type, we need to define new operators. In our
case, complex arithmetic needs four: one builds a complex value from two
scalars, one multiplies complex values, and two read the parts back as scalars
-so that results can be stored in ordinary ``float32`` buffers. Like the type,
+so that results can be stored in ordinary ``float32`` tensors. Like the type,
they are registered under the ``my_ext`` namespace:
.. list-table::
@@ -94,7 +91,7 @@ they are registered under the ``my_ext`` namespace:
- Extract the imaginary part
- ``T.float32``
-The constructor is registered as an ``OpDef`` chain in a
+The ``complex`` operator is registered as an ``OpDef`` chain in a
``static_init_block!``. Each step updates the operator registry, and the chain
completes when its statement ends:
@@ -113,14 +110,18 @@ checks the arguments and computes the result type:
:end-before: // end inference
``set_validator`` repeats these checks for a call whose result type is given
-explicitly rather than inferred. ``Pure`` marks the operator as computing a
-value without writing memory.
-
-``infer_complex`` requires both parts to be scalar floats of the same type, and
-converts neither. The other three operators follow the same pattern. ``cmul``
-requires two complex arguments of the same type and returns it; ``real`` and
-``imag`` require one and return its ``elem``. Passing an ordinary scalar to
``cmul`` therefore fails
-with ``TypeError: my_ext.cmul expects a complex value``.
+explicitly rather than inferred. ``Pure`` marks the operator as a pure
+function, which does not interact with any external state.
+
+``infer_complex`` computes the type of a ``complex(re, im)`` call: a
+``ComplexType`` whose ``dtype`` is the type of the two parts. It requires both
+parts to be scalar floats of the same type, and converts neither. It then uses
+both conversions of the reference: it builds a ``ComplexTypeObj``, turns it
into
+a ``ComplexType``, and returns it as a ``Type``. The other three operators
+follow the same pattern. ``cmul`` requires two complex arguments of the same
+type and returns it; ``real`` and ``imag`` require one and return its
+``dtype``. Passing an ordinary scalar to ``cmul`` therefore fails with
+``TypeError: my_ext.cmul expects a complex value``.
Writing a Pass
--------------
@@ -153,26 +154,18 @@ that pair. This lets the rules apply to the step-by-step
function from
R[i] = product_re
Im[i] = product_im
-Only scalar arithmetic remains. The pass looks up the four operators by name,
-as C++ does with ``Op::Get``:
-
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin operators
- :end-before: // end operators
-
-The operators are registered when the library loads, and ``Op::get`` returns
-the registry's single object for a name. ``call.op.same_as(&ops.cmul)``
-therefore tests whether a call's ``op`` is ``my_ext.cmul`` by comparing
-pointers. A lookup goes through the registry, so the pass fetches the operators
-once per run and keeps them. ``Lowering`` implements the rewrites with them:
+Only scalar arithmetic remains. The pass looks up the operators by name, as C++
+does with ``Op::Get``. The operators are registered when the library loads, and
+``Op::get`` returns the registry's single object for a name.
+``call.op.same_as(&Op::get("my_ext.cmul")?)`` therefore tests whether a call's
+``op`` is ``my_ext.cmul`` by comparing pointers. ``Lowering`` implements the
+rewrites:
.. literalinclude:: ../../examples/extension/src/lib.rs
:language: rust
:start-after: // begin lowering
:end-before: // end lowering
:prepend: struct Lowering {
- ops: ComplexOps,
bindings: Vec<(Var, Var, Var)>,
}
impl Lowering {
@@ -183,15 +176,15 @@ once per run and keeps them. ``Lowering`` implements the
rewrites with them:
fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
:append: }
-``ops`` holds the four operators, ``bindings`` records each complex variable
-with its two scalar replacements, and ``parts`` reads the real and imaginary
-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. 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.
+``bindings`` records each complex variable with its two scalar replacements,
+and ``parts`` reads the real and imaginary 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. The callback passes each node to ``lowering.dispatch`` and
+returns ``Unchanged`` when no handler matches; these rewrites do not need the
+definition region kind.
The pass is built from an ordinary named function:
@@ -205,8 +198,7 @@ rewritten function; ``Some`` replaces the original. The
first ``structural_map``
walks the function with ``Lowering``. ``PostOrder`` visits children before
their
parent, so the operands of ``cmul`` are lowered before ``cmul`` itself. The
second ``structural_map`` only checks the result. It takes a tuple of three
-functions, and calls each on the nodes of its parameter type; the closure
-passes the operators to ``check_call_lowered``:
+functions, and calls each on the nodes of its parameter type:
.. code-block:: rust
@@ -217,14 +209,15 @@ passes the operators to ``check_call_lowered``:
fn check_expr(expr: Expr) -> Result<Expr> { ... }
/// Reject a call to one of the four complex operators that the rewrite did
not consume.
- fn check_call_lowered(ops: &ComplexOps, call: Call) -> Result<Call> { ... }
+ fn check_call_lowered(call: Call) -> Result<Call> { ... }
Each returns its node unchanged and raises a ``ValueError`` on an unsupported
use, such as a complex function parameter, so the pass reports such a use
before code generation.
-Finally, ``FunctionPass::new`` wraps the function as a pass, and the factory is
-exported so that Python can call it as ``lib.lower_complex()``:
+Finally, ``FunctionPass::new`` wraps the function as a pass, and
+``tvm_ffi_dll_export_typed_func!`` exports ``lower_complex``, which builds it,
+so that the Python package can call it from the loaded library:
.. literalinclude:: ../../examples/extension/src/lib.rs
:language: rust
diff --git a/docs/advanced/extending_ir_in_rust.rst
b/docs/advanced/extending_ir_in_rust.rst
index e5d286c48..fe718306a 100644
--- a/docs/advanced/extending_ir_in_rust.rst
+++ b/docs/advanced/extending_ir_in_rust.rst
@@ -48,13 +48,15 @@ directly.
With agents, compiler development iterates quickly. Along the way, a project
may
need new types and operators to describe its own concepts, such as a new number
format or a hardware feature, and passes that lower them to what the backend
-supports. To support this, TIRx provides Rust bindings, the ``tvm-tirx`` crate,
-which let you define new types, operators, and passes on the same IR objects as
-C++ and Python. An extension is a separate Rust library that builds and loads
on
-its own, without changing the C++ library, so the C++ core stays stable while
-the extension iterates quickly. Once loaded, its types and operators can be
-written in TVMScript and its passes can run in Python compilation pipelines,
-just like built-in ones.
+supports. To support this, TIRx provides the ``tvm-tirx`` crate, with which you
+define new types, operators, and passes in Rust on the same IR objects as C++
+and Python.
+
+With the ``tvm-tirx`` crate, users can extend TIRx in separate Rust libraries.
+Each extension builds and loads on its own, without changing the C++ library,
so
+the C++ core stays stable while the extension iterates quickly. Once loaded,
its
+types and operators can be written in TVMScript and its passes can run in
Python
+compilation pipelines, just like built-in ones.
Setting Up the Project
----------------------
@@ -73,9 +75,10 @@ it is installed and imported like any other package. The
example package,
├── python/my_ext/ # loads the library and registers a compilation
pipeline
└── run_example.py # compiles and runs a function that uses the extension
-The Rust library is an ordinary crate that builds a shared library. Its only
-dependency is the `tvm-tirx <https://crates.io/crates/tvm-tirx>`_ crate, which
-re-exports TVM FFI as ``tvm_tirx::tvm_ffi``:
+The Rust library is an ordinary crate that builds a shared library. It depends
+on the `tvm-tirx <https://crates.io/crates/tvm-tirx>`_ crate, which re-exports
+TVM FFI as ``tvm_tirx::tvm_ffi``, and also lists ``tvm-ffi`` itself, because
+the code that ``#[dispatch]`` generates refers to the crate by that name:
.. literalinclude:: ../../examples/extension/Cargo.toml
:language: toml
@@ -106,7 +109,7 @@ multiplication into real arithmetic:
(a + bi)(c + di) = (ac - bd) + (ad + bc)i.
-The following function multiplies two arrays of complex numbers element by
+The following function multiplies two tensors of complex numbers element by
element. ``A`` and ``B`` hold the real and imaginary parts of the first
operand,
``C`` and ``D`` those of the second:
@@ -116,9 +119,9 @@ element. ``A`` and ``B`` hold the real and imaginary parts
of the first operand,
:end-before: # end kernel
``my_ext`` is the extension's Python package, which holds the builders of its
-type and operators; :doc:`using_ir_extensions` shows how it provides them.
``z``, ``w``, and ``product`` are complex
-values that exist only in the IR.
-The function's inputs and outputs are ordinary ``float32`` arrays, so no
+type and operators; :doc:`using_ir_extensions` shows how it provides them.
+``z``, ``w``, and ``product`` are complex values that exist only in the IR.
+The function's inputs and outputs are ordinary ``float32`` tensors, so no
runtime storage format for complex numbers is needed.
Next, :doc:`defining_ir_extensions` defines the type, operators, and pass this
diff --git a/docs/advanced/using_ir_extensions.rst
b/docs/advanced/using_ir_extensions.rst
index 29d3ba284..04e72f180 100644
--- a/docs/advanced/using_ir_extensions.rst
+++ b/docs/advanced/using_ir_extensions.rst
@@ -107,6 +107,6 @@ Next Steps
Complex addition or conjugation needs one more operator and one more rewrite
rule in ``LowerComplex``, registered and written the same way. The example
keeps complex values inside a function: parameters and outputs are ordinary
-arrays, and complex numbers have no runtime storage format or calling
+tensors, and complex numbers have no runtime storage format or calling
convention. See the :doc:`Rust API reference <../reference/rust/index>` for
the full Rust API.
diff --git a/docs/reference/rust/index.rst b/docs/reference/rust/index.rst
index 4905526cc..07cfed0f1 100644
--- a/docs/reference/rust/index.rst
+++ b/docs/reference/rust/index.rst
@@ -18,8 +18,9 @@
Rust API
========
-The `tvm-tirx <https://crates.io/crates/tvm-tirx>`__ crate binds the TIRx
compiler
-objects over TVM FFI. Its modules follow the Python package: ``tvm_tirx::ir``
is
+The `tvm-tirx <https://crates.io/crates/tvm-tirx>`__ crate provides TIRx in
Rust:
+the compiler objects over TVM FFI, and the definition of new types, operators,
and
+passes. Its modules follow the Python package: ``tvm_tirx::ir`` is
:doc:`tvm_tirx.ir <../python/ir>`, ``tvm_tirx::tirx`` is :doc:`tvm_tirx.tirx
<../python/tirx>`, and so on. Operators and types can also be defined from
Rust in
``tvm_ffi::static_init_block!`` blocks, which run when their library loads; the
@@ -45,7 +46,8 @@ Crates
- Description
* - 📦 `tvm-tirx <generated/tvm_tirx/index.html>`_
- The IR, TIRx, symbolic analysis, target, backend, runtime, codegen and
script
- bindings, and the Rust IR constructors and operators.
+ bindings, the Rust IR constructors and operators, and the definition of
new
+ types, operators, and passes.
* - 📦 `tvm-ffi <generated/tvm_ffi/index.html>`_
- The TVM FFI objects, functions, containers and ``static_init_block!``
the
bindings are built on, re-exported as ``tvm_tirx::tvm_ffi``.
diff --git a/examples/extension/Cargo.toml b/examples/extension/Cargo.toml
index 12ef76e38..a22b57344 100644
--- a/examples/extension/Cargo.toml
+++ b/examples/extension/Cargo.toml
@@ -26,7 +26,8 @@ publish = false
crate-type = ["cdylib"]
[dependencies]
-# Use this checkout so the Rust API and native libraries share the same FFI
revision.
+# This example builds against this checkout. A standalone extension depends on
+# the published releases instead, e.g. tvm-tirx = "<version>".
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 f630aae1f..778150c0d 100644
--- a/examples/extension/src/lib.rs
+++ b/examples/extension/src/lib.rs
@@ -39,7 +39,7 @@ use tvm_tirx::tvm_ffi::{
#[rs_class("my_ext.ComplexType", final)]
pub struct ComplexTypeObj {
base: TypeObj,
- pub elem: PrimType,
+ pub dtype: PrimType,
}
def_object_ref!(pub ComplexType, Type, ComplexTypeObj);
@@ -97,20 +97,20 @@ fn complex_arg(call: &CallObj, index: usize) ->
Result<ComplexType> {
.clone()
.try_cast()
.map_err(|_| type_error(&format!("{} expects a complex value",
op.name.as_str())))?;
- float_part(&ty.elem.clone().into())?;
+ float_part(&ty.dtype.clone().into())?;
Ok(ty)
}
fn infer_cmul(call: &CallObj) -> Result<Type> {
check_call(call, 2)?;
let (a, b) = (complex_arg(call, 0)?, complex_arg(call, 1)?);
- require(a.elem.dtype == b.elem.dtype, "my_ext.cmul expects complex values
of the same type")?;
+ require(a.dtype.dtype == b.dtype.dtype, "my_ext.cmul expects complex
values of the same type")?;
Ok(a.into())
}
fn infer_part(call: &CallObj) -> Result<Type> {
check_call(call, 1)?;
- Ok(complex_arg(call, 0)?.elem.clone().into())
+ Ok(complex_arg(call, 0)?.dtype.clone().into())
}
// Explicit result types bypass inference, so validation checks the inputs and
result again.
@@ -161,33 +161,17 @@ tvm_ffi::static_init_block! {
.set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
}
-// begin operators
-struct ComplexOps {
- complex: Op,
- cmul: Op,
- real: Op,
- imag: Op,
-}
-
-impl ComplexOps {
- fn get() -> Result<Self> {
- Ok(Self {
- complex: Op::get("my_ext.complex")?,
- cmul: Op::get("my_ext.cmul")?,
- real: Op::get("my_ext.real")?,
- imag: Op::get("my_ext.imag")?,
- })
- }
-
- /// Whether `op`, such as a call's `op`, is one of the four operators.
- fn contains(&self, op: &Expr) -> bool {
- [&self.complex, &self.cmul, &self.real,
&self.imag].into_iter().any(|own| op.same_as(own))
+/// Whether `op`, such as a call's `op`, is one of the four complex operators.
+fn is_complex_op(op: &Expr) -> Result<bool> {
+ for name in ["my_ext.complex", "my_ext.cmul", "my_ext.real",
"my_ext.imag"] {
+ if op.same_as(&Op::get(name)?) {
+ return Ok(true);
+ }
}
+ Ok(false)
}
-// end operators
struct Lowering {
- ops: ComplexOps,
// Each complex local variable is replaced by two scalar local variables.
bindings: Vec<(Var, Var, Var)>,
}
@@ -195,7 +179,7 @@ struct Lowering {
impl Lowering {
fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> {
if let Some(call) = expr.downcast_ref::<CallObj>() {
- if call.op.same_as(&self.ops.complex) {
+ if call.op.same_as(&Op::get("my_ext.complex")?) {
return Ok((call.args.get(0)?, call.args.get(1)?));
}
}
@@ -216,11 +200,10 @@ impl Lowering {
// begin lowering
fn map_call(&mut self, call: Call) -> Result<Any> {
- let ops = &self.ops;
- if ops.contains(&call.op) {
+ if is_complex_op(&call.op)? {
validate_call(&call)?;
}
- if call.op.same_as(&ops.cmul) {
+ if call.op.same_as(&Op::get("my_ext.cmul")?) {
let (a, b) = self.parts(&call.args.get(0)?)?;
let (c, d) = self.parts(&call.args.get(1)?)?;
// (a + bi)(c + di) = (ac - bd) + (ad + bc)i
@@ -228,7 +211,7 @@ impl Lowering {
let im = &a * &d + &b * &c;
return Ok(Call::new(
None,
- ops.complex.clone(),
+ Op::get("my_ext.complex")?,
[re, im],
None,
Vec::<Type>::new(),
@@ -236,9 +219,10 @@ impl Lowering {
)?
.into());
}
- if call.op.same_as(&ops.real) || call.op.same_as(&ops.imag) {
+ let real = call.op.same_as(&Op::get("my_ext.real")?);
+ if real || call.op.same_as(&Op::get("my_ext.imag")?) {
let (re, im) = self.parts(&call.args.get(0)?)?;
- return Ok(if call.op.same_as(&ops.real) { re.into() } else {
im.into() });
+ return Ok(if real { re.into() } else { im.into() });
}
Ok(call.into())
}
@@ -247,8 +231,8 @@ impl Lowering {
if let Some(ty) = bind.var.ty.downcast_ref::<ComplexTypeObj>() {
let (re, im) = self.parts(&bind.value)?;
let name = bind.var.name.as_str();
- let re_var = Var::new(&format!("{name}_re"), ty.elem.clone(),
bind.var.loc.clone());
- let im_var = Var::new(&format!("{name}_im"), ty.elem.clone(),
bind.var.loc.clone());
+ let re_var = Var::new(&format!("{name}_re"), ty.dtype.clone(),
bind.var.loc.clone());
+ let im_var = Var::new(&format!("{name}_im"), ty.dtype.clone(),
bind.var.loc.clone());
self.bindings.push((bind.var.clone(), re_var.clone(),
im_var.clone()));
let binds = [
Bind::new(re_var, re, bind.loc.clone())?,
@@ -280,8 +264,8 @@ fn check_expr(expr: Expr) -> Result<Expr> {
}
/// Reject a call to one of the four complex operators that the rewrite did
not consume.
-fn check_call_lowered(ops: &ComplexOps, call: Call) -> Result<Call> {
- if ops.contains(&call.op) {
+fn check_call_lowered(call: Call) -> Result<Call> {
+ if is_complex_op(&call.op)? {
return Err(Error::new(
tvm_ffi::VALUE_ERROR,
"LowerComplex left an unsupported complex call",
@@ -297,16 +281,15 @@ fn lower_function(
_module: IRModule,
_ctx: PassContext,
) -> Result<Option<Function>> {
- let mut lowering = Lowering { ops: ComplexOps::get()?, bindings:
Vec::new() };
+ let mut lowering = Lowering { bindings: Vec::new() };
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,
- (check_type, |call| check_call_lowered(ops, call), check_expr),
+ (check_type, check_call_lowered, check_expr),
WalkOrder::PostOrder,
)?;
let function: Function = checked.try_into()?;
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 42702a814..c2722e2fa 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -21,7 +21,7 @@ version = "0.1.0-alpha.5"
edition = "2021"
# The oldest toolchain that builds the crate and tvm-ffi; CMake checks cargo
against it.
rust-version = "1.84"
-description = "Rust bindings for the TVM TIRx compiler"
+description = "Rust API for the TVM TIRx compiler: IR bindings and IR
extensions"
license = "Apache-2.0"
homepage = "https://tvm.apache.org/tirx/";
documentation = "https://tvm.apache.org/tirx/reference/rust/";
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index ecd457a24..c31d497cf 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -17,9 +17,10 @@
* under the License.
*/
-//! Rust bindings for the TIRx compiler objects, over TVM FFI.
+//! The TIRx compiler in Rust: its IR objects over TVM FFI, and the definition
of new types,
+//! operators and passes that extend it.
//!
-//! The bindings are generated by `tvm-ffi-stubgen --target rust` from the
reflection metadata of
+//! The IR bindings are generated by `tvm-ffi-stubgen --target rust` from the
reflection metadata of
//! `libtvm_tirx_compiler`. The modules follow the Python package `tvm_tirx`:
//!
//! - [`ir`]: the shared IR (expressions, statements, types, functions,
modules and operators),
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index 00c692fa1..a8a54f970 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -177,7 +177,7 @@ EXTENSION_CHECKS = """
a_var, b_var = ir.Var("a", "float32"), ir.Var("b", "float32")
z = my_ext.complex(a_var, b_var)
assert isinstance(z.ty, my_ext.Complex)
- assert tvm_ffi.structural_equal(z.ty.elem, ir.PrimType("float32"))
+ assert tvm_ffi.structural_equal(z.ty.dtype, ir.PrimType("float32"))
assert tvm_ffi.structural_equal(ir.load_json(ir.save_json(z.ty)), z.ty)
text = complex_product.script()
parsed = tvm_tirx.script.parse(text, extra_vars={"T": T, "my_ext": my_ext})
