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 3e67e126e [Rust] Update pinned FFI and migrate borrowed IR APIs (#165)
3e67e126e is described below
commit 3e67e126e8096349c57261ce891b8aa8051d2562
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Nov 2 14:01:54 2026 -0500
[Rust] Update pinned FFI and migrate borrowed IR APIs (#165)
Update the pinned TVM-FFI source and migrate Rust IR APIs to borrowed
downcasts, AnyView callbacks, and consuming MutationView traversal. Keep
Rust examples and native libraries on the checkout dependency, and
update the matching and mutation guide examples.
Refresh the pinned ARM64 and x86-64 CUDA manylinux images used by
Jenkins.
---
3rdparty/tvm-ffi | 2 +-
ci/jenkins/docker-images.ini | 4 +--
docs/advanced/extending_ir_in_rust.rst | 13 ++++++----
docs/ir_guides/working_with_ir_in_rust.rst | 41 ++++++++++++++++++------------
examples/extension/Cargo.toml | 3 ++-
examples/extension/README.md | 12 ++++++---
examples/extension/src/lib.rs | 6 ++---
examples/ir_guides/rust/src/lib.rs | 25 ++++++++++--------
rust/src/ir/expr.rs | 16 ++++++------
rust/src/ir/op_def.rs | 10 ++++----
rust/src/ir/prim/const_fold.rs | 4 +--
rust/src/ir/prim/expr.rs | 2 +-
rust/src/ir/prim/op.rs | 24 ++++++++---------
rust/src/ir/prim/vector_expr.rs | 10 ++++----
rust/src/ir/stmt.rs | 15 ++++++-----
rust/src/ir/support.rs | 4 +--
rust/src/ir/type.rs | 2 +-
rust/src/lib.rs | 2 +-
rust/src/tirx/function.rs | 2 +-
rust/src/tirx/op.rs | 12 ++++-----
rust/tests/common/mod.rs | 4 +--
rust/tests/ir/op_def.rs | 2 +-
rust/tests/ir/prim/op.rs | 22 ++++++++--------
rust/tests/ir/type.rs | 14 +++++-----
rust/tests/tirx/op.rs | 2 +-
tests/python/ir/test_op_rust.py | 10 +++-----
26 files changed, 140 insertions(+), 123 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 69be48a29..3380316d0 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 69be48a2931505255733bfc5a18d6e0e5cc1a033
+Subproject commit 3380316d02864ff51513a5edd353b6d1f2e86d40
diff --git a/ci/jenkins/docker-images.ini b/ci/jenkins/docker-images.ini
index e2066d3fb..ff0dcf843 100644
--- a/ci/jenkins/docker-images.ini
+++ b/ci/jenkins/docker-images.ini
@@ -18,8 +18,8 @@
# This data file is read during when Jenkins runs job to determine docker
images.
[jenkins]
ci_tag: 20260705-142349-cfb98e93
-ci_arm:
quay.io/manylinux_cuda/manylinux_2_28_aarch64_cuda13_1@sha256:485e71e17a8c57a0a2fac0f6bd579b0819fad755f25020945d08a5e3f50e877e
+ci_arm:
quay.io/manylinux_cuda/manylinux_2_28_aarch64_cuda13_1@sha256:84c3d68c69ff277332a1c5adafe89da11d9b48b131793e97b5e767d8d62c2d89
ci_cpu: %(ci_gpu)s
-ci_gpu:
quay.io/manylinux_cuda/manylinux_2_28_x86_64_cuda13_1@sha256:7ca06e770bf3479e3334aa7bb30543a0b454a70c05890da6ac3d0cecaefd06c8
+ci_gpu:
quay.io/manylinux_cuda/manylinux_2_28_x86_64_cuda13_1@sha256:fc956755dd27491d7948d6d2974c59833e64c3c73564c0a577a39b8c2de86ebc
# Only EXTRA needs the legacy OpenCL/Vulkan/ROCm toolchains.
ci_extra: tlcpack/ci-gpu:%(ci_tag)s
diff --git a/docs/advanced/extending_ir_in_rust.rst
b/docs/advanced/extending_ir_in_rust.rst
index ce8c3ac03..ab19d4609 100644
--- a/docs/advanced/extending_ir_in_rust.rst
+++ b/docs/advanced/extending_ir_in_rust.rst
@@ -75,15 +75,18 @@ re-exports TVM FFI as ``tvm_tirx::tvm_ffi``:
:language: toml
:start-at: [lib]
-With Rust 1.84 or newer, build and install the package from the TIRx
-repository root:
+First build and install TIRx from this checkout. With Rust 1.84 or newer,
+install the extension from the repository root in the same Python environment:
.. code-block:: bash
- pip install ./examples/extension
+ uv pip install scikit-build-core
+ uv pip install --no-build-isolation --no-deps ./examples/extension
-The package depends on ``apache-tvm-tirx``, and the build links the library
-against the TIRx native libraries that this package provides.
+The example uses this checkout's Rust crate and pinned TVM FFI sources.
+The package depends on ``apache-tvm-tirx`` at run time; disabling build
+isolation keeps its native build aligned with these sources instead of
+installing a released package in a temporary build environment.
Case Study: Complex Numbers
---------------------------
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 4f6bcc3d8..65fa4a086 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -248,7 +248,9 @@ requested ``prim.Mul``.
``try_cast`` consumes the managed reference it converts, so we clone
``value`` before the first cast; the clone costs one reference-count
- increment.
+ increment. For read-only inspection, ``value.downcast_ref::<AddObj>()``
+ returns ``Option<&AddObj>`` without consuming or cloning the reference.
+ ``Any`` and ``AnyView`` provide the same borrowed downcast.
Casting to one type at a time gets tedious when a node can be one of
many; that is what matching is for.
@@ -279,9 +281,9 @@ The output is:
((x * 2) + 1)
<other>
-Each arm names a managed reference type and binds the matching node to it,
-and the first matching arm runs. The ``Add`` and ``Mul`` arms recurse into
-the operands, the ``Var`` and ``IntImm`` arms end the recursion, and
+Each arm names an object type and borrows the matching node with ``ref``;
+the first matching arm runs. The ``AddObj`` and ``MulObj`` arms recurse into
+the operands, the ``VarObj`` and ``IntImmObj`` arms end the recursion, and
``x - 1`` falls to the ``_`` arm because ``describe`` has no arm for
subtraction.
@@ -292,10 +294,12 @@ types it knows, and ``_`` lets programs with new types
pass through.
.. note::
- Like ``try_cast``, ``match_any!`` consumes what it matches, so
- ``describe`` passes it ``expr.clone()``. ``match_any!`` works on any
- TVM-FFI object, such as tensors, shapes, and typed arrays, and on
- ``Any`` values. Plain values such as integers take the ``_`` arm.
+ ``Obj(ref node)`` borrows the underlying object without incrementing its
+ reference count. ``describe`` converts ``expr`` to an ``AnyView``. An
+ owning arm such as ``Add(add)`` still produces a managed reference.
+ ``match_any!`` also works on ``Any`` and ``AnyView`` values, including
+ tensors, shapes, and typed arrays. Plain values such as integers take
+ the ``_`` arm in this example.
A match looks at one node at a time; to process every node in a program
while keeping state, we dispatch.
@@ -408,7 +412,7 @@ The parameter list of a callback decides what it receives:
* - ``|store: &TensorStoreObj|``
- Each write, as a borrowed object, without changing its reference
count.
- * - ``|value: &StructuralView|``
+ * - ``|value: AnyView<'_>|``
- Every value the walk reaches.
* - ``|var: Var, kind: DefRegionKind|``
- Each variable, with its definition region kind.
@@ -546,8 +550,12 @@ 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 ``default_mutate`` for the children and
-``def_region_kind()``. Let's reverse the loop
+``VisitContext``, with ``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,
+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:
@@ -569,8 +577,9 @@ The output is:
for i in range(n):
B[n - 1 - i] = A[n - 1 - i] + T.float32(1.0)
-``mutate_for`` records the replacement, then ``default_mutate`` rewrites
-the loop's children through the same struct, so ``mutate_var`` sees the
+``mutate_for`` records the replacement, then
+``default_maybe_inplace_mutate`` rewrites the loop's children through the
+same struct, so ``mutate_var`` sees the
replacement before any use of ``i``.
.. note::
@@ -602,7 +611,7 @@ Quick Reference
- What it does
* - |match_any!|_
- Branches on the node type; the first matching arm runs
- * - |try_cast|_, |as_node|_
+ * - |try_cast|_, |downcast_ref|_
- Downcasts to one type, as a managed reference in a ``Result`` or a
borrowed object in an ``Option``
* - ``.into()``
@@ -666,7 +675,7 @@ Quick Reference
.. _match_any!: ../reference/rust/generated/tvm_ffi/macro.match_any.html
.. |Var::new| replace:: ``Var::new``
.. _Var::new:
../reference/rust/generated/tvm_tirx/ir/expr/struct.Var.html#method.new
-.. |as_node| replace:: ``as_node``
-.. _as_node:
../reference/rust/generated/tvm_ffi/object/trait.ObjectRefCore.html#method.as_node
+.. |downcast_ref| replace:: ``downcast_ref``
+.. _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
diff --git a/examples/extension/Cargo.toml b/examples/extension/Cargo.toml
index 809a036b8..61beb674c 100644
--- a/examples/extension/Cargo.toml
+++ b/examples/extension/Cargo.toml
@@ -26,4 +26,5 @@ publish = false
crate-type = ["cdylib"]
[dependencies]
-tvm-tirx = "0.1.0-alpha.5"
+# Use this checkout so the Rust API and native libraries share the same FFI
revision.
+tvm-tirx = { path = "../../rust" }
diff --git a/examples/extension/README.md b/examples/extension/README.md
index 25dcb9c0e..8e2928abe 100644
--- a/examples/extension/README.md
+++ b/examples/extension/README.md
@@ -30,12 +30,16 @@ pipeline. The tutorial in `docs/advanced/` walks through
the code.
## Install and Run
-With Rust 1.84 or newer, from the repository root:
+After building and installing TIRx from this checkout, with Rust 1.84 or newer,
+run from the repository root in the same Python environment:
```bash
-pip install ./examples/extension
+uv pip install scikit-build-core
+uv pip install --no-build-isolation --no-deps ./examples/extension
python examples/extension/run_example.py
```
-The build compiles the library against the published `tvm-tirx` crate and the
-`apache-tvm-tirx` package, which the extension also depends on at run time.
+The build uses this checkout's `tvm-tirx` crate and pinned TVM FFI sources.
+Disabling build isolation keeps the native libraries from the same checkout
+available to cargo, instead of installing a released package in a temporary
+build environment.
diff --git a/examples/extension/src/lib.rs b/examples/extension/src/lib.rs
index 5fa78ebb7..c2368002c 100644
--- a/examples/extension/src/lib.rs
+++ b/examples/extension/src/lib.rs
@@ -194,7 +194,7 @@ struct Lowering {
impl Lowering {
fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> {
- if let Some(call) = expr.as_node::<CallObj>() {
+ if let Some(call) = expr.downcast_ref::<CallObj>() {
if call.op.same_as(&self.ops.complex) {
return Ok((call.args.get(0)?, call.args.get(1)?));
}
@@ -244,7 +244,7 @@ impl Lowering {
}
fn map_bind(&mut self, bind: Bind) -> Result<Stmt> {
- if let Some(ty) = bind.var.ty.as_node::<ComplexTypeObj>() {
+ 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());
@@ -263,7 +263,7 @@ impl Lowering {
/// Reject a complex type, wherever it appears: on a variable, a parameter or
an expression.
fn check_type(ty: Type) -> Result<Type> {
- if ty.as_node::<ComplexTypeObj>().is_some() {
+ if ty.downcast_ref::<ComplexTypeObj>().is_some() {
return Err(Error::new(
tvm_ffi::VALUE_ERROR,
"LowerComplex left an unsupported complex type",
diff --git a/examples/ir_guides/rust/src/lib.rs
b/examples/ir_guides/rust/src/lib.rs
index 8c9a74ce4..c906ab4c6 100644
--- a/examples/ir_guides/rust/src/lib.rs
+++ b/examples/ir_guides/rust/src/lib.rs
@@ -24,15 +24,16 @@
use std::collections::HashMap;
-use tvm_tirx::ir::prim::{make_const, Add, AddObj, Mul};
+use tvm_tirx::ir::prim::{make_const, Add, AddObj, Mul, MulObj};
use tvm_tirx::ir::{
- Expr, For, ForObj, IntImm, PrimType, TensorStore, TensorStoreObj, Var,
VarObj, UNKNOWN_LOC,
+ Expr, For, ForObj, IntImmObj, PrimType, TensorStore, TensorStoreObj, Var,
VarObj, UNKNOWN_LOC,
};
use tvm_tirx::script::script;
use tvm_tirx::tirx::Function;
use tvm_tirx::tvm_ffi::{
- self, dispatch, match_any, Any, Array, DefRegionKind, Mutator,
ObjectIdentity, ObjectRefCast,
- ObjectRefCore, Result, VisitCallbacks, VisitContext, VisitInterrupt,
WalkOrder, WalkResult,
+ self, dispatch, match_any, Any, AnyView, Array, DefRegionKind,
MutationView, Mutator,
+ ObjectIdentity, ObjectRefCast, ObjectRefCore, Result, VisitCallbacks,
VisitContext,
+ VisitInterrupt, WalkOrder, WalkResult,
};
/// A Python list of strings.
@@ -86,11 +87,11 @@ fn casts() -> Result<Strings> {
// [describe.begin]
fn describe(expr: &Expr) -> Result<String> {
let text = match_any! {
- expr.clone() {
- Add(add) => format!("({} + {})", describe(&add.a)?,
describe(&add.b)?),
- Mul(mul) => format!("({} * {})", describe(&mul.a)?,
describe(&mul.b)?),
- Var(var) => var.name.as_str().to_owned(),
- IntImm(imm) => imm.value.to_string(),
+ AnyView::from(expr) {
+ AddObj(ref add) => format!("({} + {})", describe(&add.a)?,
describe(&add.b)?),
+ MulObj(ref mul) => format!("({} * {})", describe(&mul.a)?,
describe(&mul.b)?),
+ VarObj(ref var) => var.name.as_str().to_owned(),
+ IntImmObj(ref imm) => imm.value.to_string(),
_ => "<other>".to_owned(),
}
};
@@ -240,10 +241,12 @@ struct ReverseLoops {
#[dispatch(mutate)]
impl ReverseLoops {
- fn mutate_for(&mut self, loop_: For, mutator: &mut Mutator) -> Result<Any>
{
+ 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_mutate(self, &loop_)
+ mutator.default_maybe_inplace_mutate(self, value)
}
fn mutate_var(&mut self, var: Var, mutator: &mut Mutator) -> Expr {
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index 2dbd74e67..bad821470 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -463,7 +463,7 @@ impl Tuple {
) -> Self {
let fields =
fields.into_iter().map(Into::into).collect::<Array<Expr>>();
let ty: Type =
- match fields.iter().any(|field|
field.ty.as_node::<MissingTypeObj>().is_some()) {
+ match fields.iter().any(|field|
field.ty.downcast_ref::<MissingTypeObj>().is_some()) {
true => MissingType::new().into(),
false => {
let field_types = fields.iter().map(|field|
field.ty.clone());
@@ -556,7 +556,7 @@ impl TupleGetItem {
&message,
));
}
- let ty: Type = match tuple_value.ty.as_node::<TupleTypeObj>() {
+ let ty: Type = match tuple_value.ty.downcast_ref::<TupleTypeObj>() {
Some(tuple_type) => {
let size = tuple_type.fields.len();
if index as usize >= size {
@@ -1213,7 +1213,7 @@ impl GenericConst {
) -> Result<Self> {
use TVMFFITypeIndex::*;
let (ty, value) = (ty.into(), value.into());
- if ty.as_node::<MissingTypeObj>().is_some() {
+ if ty.downcast_ref::<MissingTypeObj>().is_some() {
let condition = "!ty.as<MissingType>().has_value()";
let message = "GenericConst requires an expression type";
return Err(support::check_failed(tvm_ffi::TYPE_ERROR, condition,
message));
@@ -1477,7 +1477,7 @@ impl Call {
/// Check the call against the contract its operator registered: the
operator's validator, if
/// it has one (C++ `Call::Validate`).
pub fn validate(&self) -> Result<()> {
- let Some(op) = self.op.as_node::<OpObj>() else { return Ok(()) };
+ let Some(op) = self.op.downcast_ref::<OpObj>() else { return Ok(()) };
if op.validator().type_index() == TVMFFITypeIndex::kTVMFFINone as i32 {
return Ok(());
}
@@ -1488,8 +1488,8 @@ impl Call {
/// `TFixedReturnType`, else what its `FInferType` infers; the return type
of a callee of a
/// function type; otherwise a missing type (C++ `Call::ReinferType`).
pub fn reinfer_type(&self) -> Result<Type> {
- let Some(op) = self.op.as_node::<OpObj>() else {
- return match self.op.ty.as_node::<FuncTypeObj>() {
+ let Some(op) = self.op.downcast_ref::<OpObj>() else {
+ return match self.op.ty.downcast_ref::<FuncTypeObj>() {
Some(func_type) => Ok(func_type.ret_type.clone()),
None => Ok(MissingType::new().into()),
};
@@ -1809,7 +1809,7 @@ impl Range {
let (begin, end) = (begin.into_expr()?, end.into_expr()?);
support::prim_dtype(&begin)?;
support::prim_dtype(&end)?;
- if let Some(add) = end.as_node::<AddObj>() {
+ if let Some(add) = end.downcast_ref::<AddObj>() {
if crate::ir::prim::expr_deep_equal(add.a.clone(), begin.clone())?
{
return Ok(add.b.clone());
}
@@ -1818,7 +1818,7 @@ impl Range {
}
}
let is_zero =
- begin.as_node::<IntImmObj>().is_some_and(|imm| imm.value ==
BigInt::from(0i64));
+ begin.downcast_ref::<IntImmObj>().is_some_and(|imm| imm.value ==
BigInt::from(0i64));
match is_zero {
true => Ok(end),
false => crate::ir::prim::op::sub(end, begin),
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index c91b09345..dcd963940 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -509,7 +509,7 @@ pub mod sig {
/// Whether `value` is a `Var`, not a subclass, whose type is an `N`.
pub fn var_of_type<N: ObjectCore>(value: &TVMFFIAny) -> bool {
value.type_index == crate::ir::expr::VarObj::type_index()
- && expr(value).is_some_and(|var| var.ty.as_node::<N>().is_some())
+ && expr(value).is_some_and(|var|
var.ty.downcast_ref::<N>().is_some())
}
/// The expression `value` holds, if it holds one.
@@ -523,13 +523,13 @@ pub mod sig {
}
fn prim_dtype(value: &TVMFFIAny) -> Option<tvm_ffi::DLDataType> {
- Some(expr(value)?.ty.as_node::<PrimTypeObj>()?.dtype)
+ Some(expr(value)?.ty.downcast_ref::<PrimTypeObj>()?.dtype)
}
/// C++ `TypeTraits<TypedExpr<T>>::GetMismatchTypeInfo`: an expression
with its type.
fn typed_expr_mismatch(value: &TVMFFIAny) -> std::string::String {
let Some(expr) = expr(value) else { return
support::type_key(value.type_index) };
- let ty = match expr.ty.as_node::<PrimTypeObj>() {
+ let ty = match expr.ty.downcast_ref::<PrimTypeObj>() {
Some(prim) => support::dtype_name(prim.dtype),
None => support::type_key(AnyView::from(&expr.ty).type_index()),
};
@@ -713,7 +713,7 @@ pub mod sig {
fn check(call: &CallObj, op: &OpObj, _: (usize, usize), _: (usize,
usize)) -> Result<()> {
let actual = match &call.attrs {
- Some(attrs) if attrs.as_node::<T::ContainerType>().is_some()
=> return Ok(()),
+ Some(attrs) if
attrs.downcast_ref::<T::ContainerType>().is_some() => return Ok(()),
Some(attrs) => {
format!("`{}`",
support::type_key(AnyView::from(attrs).type_index()))
}
@@ -841,7 +841,7 @@ pub mod sig {
"a signature has at most one variadic tail of each kind and
one call_attrs"
)
};
- let Some(op) = call.op.as_node::<OpObj>() else {
+ let Some(op) = call.op.downcast_ref::<OpObj>() else {
return Err(Error::new(TYPE_ERROR, "Invalid Op Call", ""));
};
let count = |label: &str, expected: usize, tail: bool, unit: &str,
got: usize| {
diff --git a/rust/src/ir/prim/const_fold.rs b/rust/src/ir/prim/const_fold.rs
index 57b15a73a..466f64829 100644
--- a/rust/src/ir/prim/const_fold.rs
+++ b/rust/src/ir/prim/const_fold.rs
@@ -80,12 +80,12 @@ pub(super) fn is_in_float32_range(value: f64) -> bool {
/// The value of an integer constant.
pub(super) fn int(expr: &Expr) -> Option<&BigInt> {
- expr.as_node::<IntImmObj>().map(|imm| &imm.value)
+ expr.downcast_ref::<IntImmObj>().map(|imm| &imm.value)
}
/// The value of a float constant.
fn float(expr: &Expr) -> Option<f64> {
- expr.as_node::<FloatImmObj>().map(|imm| imm.value)
+ expr.downcast_ref::<FloatImmObj>().map(|imm| imm.value)
}
/// C++ `IntImm(dtype, value)`.
diff --git a/rust/src/ir/prim/expr.rs b/rust/src/ir/prim/expr.rs
index 323793b7b..393b5d5b1 100644
--- a/rust/src/ir/prim/expr.rs
+++ b/rust/src/ir/prim/expr.rs
@@ -1961,7 +1961,7 @@ impl Let {
let (var, value, body) = (var.into(), value.into_expr()?,
body.into_expr()?);
let value_dtype = support::prim_dtype(&value)?;
support::prim_dtype(&body)?;
- let Some(var_ty) = var.ty.as_node::<PrimTypeObj>() else {
+ let Some(var_ty) = var.ty.downcast_ref::<PrimTypeObj>() else {
let message = format!(
"Cannot treat type `{}` as type `ir.PrimType`",
support::type_key(AnyView::from(&var.ty).type_index())
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index f1b092f72..869ab9f39 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -242,7 +242,7 @@ fn call(ty: Option<&Type>, op: &'static str, args:
Vec<Expr>, loc: Location) ->
/// Whether `expr` is a call of `prim.vscale`.
fn is_vscale_call(expr: &Expr) -> Result<bool> {
- let Some(call) = expr.as_node::<CallObj>() else { return Ok(false) };
+ let Some(call) = expr.downcast_ref::<CallObj>() else { return Ok(false) };
Ok(call.op.same_as(&support::cached_op("prim.vscale")?))
}
@@ -573,7 +573,7 @@ fn fold_cast(value: &Expr, to: DLDataType) ->
Result<Option<Expr>> {
let float = |value: f64, loc: &Location| -> Result<Option<Expr>> {
Ok(Some(FloatImm::from_dtype(to, value, loc.clone())?.into()))
};
- if let Some(imm) = value.as_node::<IntImmObj>() {
+ if let Some(imm) = value.downcast_ref::<IntImmObj>() {
if integer {
if const_fold::is_representable(&imm.value, to) {
return int(imm.value.clone(), &imm.loc);
@@ -594,7 +594,7 @@ fn fold_cast(value: &Expr, to: DLDataType) ->
Result<Option<Expr>> {
return float(source, &imm.loc);
}
}
- } else if let Some(imm) = value.as_node::<FloatImmObj>() {
+ } else if let Some(imm) = value.downcast_ref::<FloatImmObj>() {
let from = support::prim_dtype(value)?;
if imm.value.is_finite() && from.bits == 16 {
let bfloat = match from.code {
@@ -690,11 +690,11 @@ pub fn cast_at(
if !lanes_match {
return Err(support::check_failed(INTERNAL_ERROR, "lanes_match", ""));
}
- if let Some(broadcast) = value.as_node::<BroadcastObj>() {
+ if let Some(broadcast) = value.downcast_ref::<BroadcastObj>() {
let value = cast_at(broadcast.value.clone(), element, loc.clone())?;
return Ok(Broadcast::new(value, broadcast.lanes.clone(), loc)?.into());
}
- if let Some(ramp) = value.as_node::<RampObj>() {
+ if let Some(ramp) = value.downcast_ref::<RampObj>() {
// Casting the base and the stride is valid only when every lane is
the same integer in
// both types: widening must not remove an overflow of the source, and
a floating-point
// cast need not be affine.
@@ -937,13 +937,13 @@ pub fn eq_at(a: impl Operand, b: impl Operand, loc: impl
Into<Location>) -> Resu
pub fn neg_at(a: impl IntoExpr, loc: impl Into<Location>) -> Result<Expr> {
let (a, loc) = (a.into_expr()?, loc.into());
let dtype = support::prim_dtype(&a)?;
- if let Some(imm) = a.as_node::<IntImmObj>() {
+ if let Some(imm) = a.downcast_ref::<IntImmObj>() {
let value = -imm.value.clone();
if const_fold::is_representable(&value, dtype) {
return Ok(IntImm::from_dtype(dtype, value, loc)?.into());
}
}
- if let Some(imm) = a.as_node::<FloatImmObj>() {
+ if let Some(imm) = a.downcast_ref::<FloatImmObj>() {
return Ok(FloatImm::from_dtype(dtype, -imm.value, loc)?.into());
}
sub(make_const_at(PrimType::from_dtype(dtype)?, 0, loc)?, a)
@@ -1163,7 +1163,7 @@ pub fn if_then_else_at(
return Err(support::check_failed(INTERNAL_ERROR, condition, message));
}
let (true_value, false_value) = match_types_at(true_value, false_value,
loc.clone())?;
- if let Some(imm) = cond.as_node::<IntImmObj>() {
+ if let Some(imm) = cond.downcast_ref::<IntImmObj>() {
return Ok(if imm.value.is_zero() { false_value } else { true_value });
}
let ty = true_value.ty.clone();
@@ -1311,7 +1311,7 @@ fn rounding(x: Expr, op: &'static str, fold: fn(f64) ->
f64, loc: Location) -> R
if matches!(dtype.code, support::INT | support::UINT | support::BOOL) {
return Ok(x);
}
- if let Some(imm) = x.as_node::<FloatImmObj>() {
+ if let Some(imm) = x.downcast_ref::<FloatImmObj>() {
return Ok(FloatImm::from_dtype(dtype, fold(imm.value),
imm.loc.clone())?.into());
}
let ty = x.ty.clone();
@@ -1379,7 +1379,7 @@ fn absolute(x: Expr, loc: Location) -> Result<Expr> {
let dtype = support::prim_dtype(&x)?;
match dtype.code {
support::INT => integer_abs(x, loc),
- support::FLOAT | support::BFLOAT => match x.as_node::<FloatImmObj>() {
+ support::FLOAT | support::BFLOAT => match
x.downcast_ref::<FloatImmObj>() {
Some(imm) => Ok(FloatImm::from_dtype(dtype, imm.value.abs(),
imm.loc.clone())?.into()),
None => {
let ty = x.ty.clone();
@@ -1401,7 +1401,7 @@ fn absolute(x: Expr, loc: Location) -> Result<Expr> {
/// value of its type; anything else is `select(x >= 0, x, -x)`.
fn integer_abs(x: Expr, loc: Location) -> Result<Expr> {
let dtype = support::prim_dtype(&x)?;
- if let Some(imm) = x.as_node::<IntImmObj>() {
+ if let Some(imm) = x.downcast_ref::<IntImmObj>() {
let value = match imm.value.is_negative() {
true => -imm.value.clone(),
false => imm.value.clone(),
@@ -1444,7 +1444,7 @@ fn finite(x: Expr, loc: Location) -> Result<Expr> {
}
fn likely_hint(cond: Expr, loc: Location) -> Result<Expr> {
- if cond.as_node::<IntImmObj>().is_some() {
+ if cond.downcast_ref::<IntImmObj>().is_some() {
return Ok(cond);
}
let ty = cond.ty.clone();
diff --git a/rust/src/ir/prim/vector_expr.rs b/rust/src/ir/prim/vector_expr.rs
index 90f52a5d5..b506d63fd 100644
--- a/rust/src/ir/prim/vector_expr.rs
+++ b/rust/src/ir/prim/vector_expr.rs
@@ -372,7 +372,7 @@ impl Shuffle {
/// The vector type of `element` elements over `lanes`, and the lanes a `Ramp`
or `Broadcast` keeps:
/// a constant count as an `int32` constant, or a scalable count as `vscale *
factor`.
fn vector_lanes(element: DLDataType, lanes: Expr) -> Result<(DLDataType,
Expr)> {
- if let Some(imm) = lanes.as_node::<IntImmObj>() {
+ if let Some(imm) = lanes.downcast_ref::<IntImmObj>() {
let Some(count) = imm.value.to_i64().and_then(|count|
i32::try_from(count).ok()) else {
return Err(tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, "bad
optional access", ""));
};
@@ -411,17 +411,17 @@ pub(crate) fn vscale_lanes(factor: i32) -> Result<Expr> {
/// `factor` when `lanes` is `factor * vscale` or `vscale * factor` (C++
`ExtractVscaleFactor`).
fn vscale_factor(lanes: &Expr) -> Result<Option<i32>> {
- let Some(mul) = lanes.as_node::<MulObj>() else { return Ok(None) };
+ let Some(mul) = lanes.downcast_ref::<MulObj>() else { return Ok(None) };
let vscale = support::cached_op("prim.vscale")?;
let is_vscale =
- |expr: &Expr| expr.as_node::<CallObj>().is_some_and(|call|
call.op.same_as(&vscale));
+ |expr: &Expr| expr.downcast_ref::<CallObj>().is_some_and(|call|
call.op.same_as(&vscale));
let factor = |imm: &IntImmObj| imm.value.to_i64().and_then(|value|
i32::try_from(value).ok());
- if let Some(imm) = mul.a.as_node::<IntImmObj>() {
+ if let Some(imm) = mul.a.downcast_ref::<IntImmObj>() {
if is_vscale(&mul.b) {
return Ok(factor(imm));
}
}
- if let Some(imm) = mul.b.as_node::<IntImmObj>() {
+ if let Some(imm) = mul.b.downcast_ref::<IntImmObj>() {
if is_vscale(&mul.a) {
return Ok(factor(imm));
}
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index 21a7457f4..d31039ae0 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -470,7 +470,7 @@ impl For {
if ty.dtype == var_ty.dtype {
return Ok(expr);
}
- if let Some(imm) = expr.as_node::<IntImmObj>() {
+ if let Some(imm) = expr.downcast_ref::<IntImmObj>() {
if !int_imm_value_fits(&imm.value, var_ty.dtype)? {
let message = format!(
"Literal value {} is not representable in the loop
variable's dtype ({})",
@@ -865,7 +865,7 @@ fn region_body_params(op: &Op, args: &Array<Expr>, attrs:
&DictAttrs) -> Result<
let params: Array<Var> = get_body_params.call_as(&signature)?;
let mut definitions = std::collections::HashSet::new();
for param in params.iter() {
- if param.ty.as_node::<MissingTypeObj>().is_some() {
+ if param.ty.downcast_ref::<MissingTypeObj>().is_some() {
let condition = "!param->ty.as<MissingType>().has_value()";
let message = format!("FRegionGetBodyParams for {name} must return
typed variables");
return Err(support::check_failed(tvm_ffi::VALUE_ERROR, condition,
&message));
@@ -951,12 +951,13 @@ impl SeqStmt {
/// `seq` with each nested sequence replaced by its statements and the no-ops
removed.
fn flatten(seq: Array<Stmt>) -> Array<Stmt> {
- let nested_or_no_op = |stmt: &Stmt| stmt.as_node::<SeqStmtObj>().is_some()
|| is_no_op(stmt);
+ let nested_or_no_op =
+ |stmt: &Stmt| stmt.downcast_ref::<SeqStmtObj>().is_some() ||
is_no_op(stmt);
if !seq.iter().any(|stmt| nested_or_no_op(&stmt)) {
return seq;
}
fn append(stmt: Stmt, flattened: &mut Vec<Stmt>) {
- match stmt.as_node::<SeqStmtObj>() {
+ match stmt.downcast_ref::<SeqStmtObj>() {
Some(nested) => nested.seq.iter().for_each(|child| append(child,
flattened)),
None if is_no_op(&stmt) => {}
None => flattened.push(stmt),
@@ -969,15 +970,15 @@ fn flatten(seq: Array<Stmt>) -> Array<Stmt> {
/// Whether `stmt` is `Evaluate(0)`, the no-op a sequence drops (C++
`IsSeqStmtNoOp`).
fn is_no_op(stmt: &Stmt) -> bool {
- stmt.as_node::<EvaluateObj>()
- .and_then(|evaluate| evaluate.value.as_node::<IntImmObj>())
+ stmt.downcast_ref::<EvaluateObj>()
+ .and_then(|evaluate| evaluate.value.downcast_ref::<IntImmObj>())
.is_some_and(|imm| imm.value == BigInt::from(0i64))
}
/// A statement as a body, as C++ converts a `Stmt` to a `SeqStmt`: a sequence
as it is, any other
/// statement as the sequence of it at its location (C++ `SeqStmt(stmt)`).
fn into_body(stmt: Stmt) -> Result<SeqStmt> {
- if stmt.as_node::<SeqStmtObj>().is_some() {
+ if stmt.downcast_ref::<SeqStmtObj>().is_some() {
return Ok(tvm_ffi::ObjectRefCast::try_cast(stmt).expect("a SeqStmt"));
}
let loc = stmt.loc.clone();
diff --git a/rust/src/ir/support.rs b/rust/src/ir/support.rs
index de9a6a910..309c47cfc 100644
--- a/rust/src/ir/support.rs
+++ b/rust/src/ir/support.rs
@@ -214,7 +214,7 @@ pub(crate) fn shared<T: Clone + 'static>(
/// The type of a primitive expression. C++ takes such an argument as a
`PrimExpr`, whose
/// conversion accepts only an expression of a `PrimType`.
pub(crate) fn prim_type(expr: &Expr) -> Result<PrimType> {
- match expr.ty.as_node::<PrimTypeObj>() {
+ match expr.ty.downcast_ref::<PrimTypeObj>() {
Some(_) =>
Ok(tvm_ffi::ObjectRefCast::try_cast(expr.ty.clone()).expect("a PrimType")),
None => Err(not_a_prim_expr(expr)),
}
@@ -222,7 +222,7 @@ pub(crate) fn prim_type(expr: &Expr) -> Result<PrimType> {
/// The dtype of a primitive expression ([`prim_type`]).
pub(crate) fn prim_dtype(expr: &Expr) -> Result<DLDataType> {
- match expr.ty.as_node::<PrimTypeObj>() {
+ match expr.ty.downcast_ref::<PrimTypeObj>() {
Some(ty) => Ok(ty.dtype),
None => Err(not_a_prim_expr(expr)),
}
diff --git a/rust/src/ir/type.rs b/rust/src/ir/type.rs
index 651106b48..d0cdbb57e 100644
--- a/rust/src/ir/type.rs
+++ b/rust/src/ir/type.rs
@@ -683,7 +683,7 @@ impl PtrType {
/// (C++ `PtrType(element_type, scope)`).
pub fn new(element_type: impl Into<Type>, scope: &MemoryScope) ->
Result<Self> {
let element_type = element_type.into();
- if element_type.as_node::<MissingTypeObj>().is_some() {
+ if element_type.downcast_ref::<MissingTypeObj>().is_some() {
let condition = "!element_type.as<MissingType>().has_value()";
let message = "PtrType element_type cannot be Type::Missing()";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
condition, message));
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 83fc8b6fb..ecd457a24 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -205,7 +205,7 @@
//! fn main() {
//! let pair = PairType::new(PrimType::int(32), PrimType::float(32));
//! let ty: Type = pair.into(); // an `ir.Type` that C++ sees as
`my_ext.PairType`
-//! assert!(ty.as_node::<PairTypeObj>().is_some());
+//! assert!(ty.downcast_ref::<PairTypeObj>().is_some());
//! }
//! ```
//!
diff --git a/rust/src/tirx/function.rs b/rust/src/tirx/function.rs
index e74fcf037..dddf997f3 100644
--- a/rust/src/tirx/function.rs
+++ b/rust/src/tirx/function.rs
@@ -147,7 +147,7 @@ impl Function {
) -> Self {
let params =
params.into_iter().map(Into::into).collect::<Array<Var>>();
let ret_type: Type = ret_type.into();
- let ret_type = match ret_type.as_node::<MissingTypeObj>() {
+ let ret_type = match ret_type.downcast_ref::<MissingTypeObj>() {
Some(_) => TupleType::new(Vec::<Type>::new(),
UnknownLoc::new()).into(),
None => ret_type,
};
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index 7ea2f6bd0..8c13ebb5c 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -208,12 +208,12 @@ pub fn reinterpret_at(
loc: impl Into<Location>,
) -> Result<Expr> {
let (ty, value, loc) = (ty.into(), value.into_expr()?, loc.into());
- let target_is_pointer = ty.as_node::<PtrTypeObj>().is_some();
- let source_is_pointer = value.ty.as_node::<PtrTypeObj>().is_some();
+ let target_is_pointer = ty.downcast_ref::<PtrTypeObj>().is_some();
+ let source_is_pointer = value.ty.downcast_ref::<PtrTypeObj>().is_some();
let type_error = |condition: &str, message: std::string::String| {
Err(support::check_failed(TYPE_ERROR, condition, &message))
};
- if value.as_node::<StringImmObj>().is_some() {
+ if value.downcast_ref::<StringImmObj>().is_some() {
if !target_is_pointer {
let message = format!(
"String reinterpret requires a pointer target, but got {}",
@@ -221,9 +221,9 @@ pub fn reinterpret_at(
);
return type_error("target_ty.as<PtrTypeObj>()", message);
}
- } else if let Some(target) = ty.as_node::<PrimTypeObj>() {
+ } else if let Some(target) = ty.downcast_ref::<PrimTypeObj>() {
let target_dtype = target.dtype;
- if let Some(source) = value.ty.as_node::<PrimTypeObj>() {
+ if let Some(source) = value.ty.downcast_ref::<PrimTypeObj>() {
let source_dtype = source.dtype;
if source_dtype == target_dtype {
return Ok(value);
@@ -279,7 +279,7 @@ pub fn reinterpret_at(
);
return type_error("target_ty.as<PtrTypeObj>()", message);
}
- if let Some(source) = value.ty.as_node::<PrimTypeObj>() {
+ if let Some(source) = value.ty.downcast_ref::<PrimTypeObj>() {
if !is_scalar_int64(source.dtype) {
let condition = "source_dtype.value().IsScalar() &&
source_dtype.value().bits() == 64 && \
source_dtype.value().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)";
diff --git a/rust/tests/common/mod.rs b/rust/tests/common/mod.rs
index 1a350b0f3..7bbf19db1 100644
--- a/rust/tests/common/mod.rs
+++ b/rust/tests/common/mod.rs
@@ -110,8 +110,8 @@ fn agree_up_to<T: AnyCompatible>(
);
if let (Ok(rust), Ok(native)) = (Expr::try_from(rust),
Expr::try_from(native)) {
assert!(structural_equal(&rust.ty, &native.ty), "{what}: the
types differ");
- let shared = native.ty.as_node::<PrimTypeObj>().is_some()
- || native.ty.as_node::<MissingTypeObj>().is_some();
+ let shared = native.ty.downcast_ref::<PrimTypeObj>().is_some()
+ || native.ty.downcast_ref::<MissingTypeObj>().is_some();
if shared {
assert!(rust.ty.same_as(&native.ty), "{what}: the shared
type is not shared");
}
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 10d521e50..b112eedc7 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -94,7 +94,7 @@ fn first_arg_type(call: &CallObj) -> Result<Type> {
}
fn only_variables(call: &CallObj) -> Result<()> {
- match call.args.iter().all(|arg| arg.as_node::<VarObj>().is_some()) {
+ match call.args.iter().all(|arg| arg.downcast_ref::<VarObj>().is_some()) {
true => Ok(()),
false => Err(Error::new(VALUE_ERROR, "arguments must be variables",
"")),
}
diff --git a/rust/tests/ir/prim/op.rs b/rust/tests/ir/prim/op.rs
index 0179290ae..d7617af8e 100644
--- a/rust/tests/ir/prim/op.rs
+++ b/rust/tests/ir/prim/op.rs
@@ -428,10 +428,10 @@ fn located_variants() {
let located: prim::FloorDiv = op::floordiv_at(&x, 4,
here.clone()).unwrap().try_cast().unwrap();
assert!(located.loc.same_as(&here));
let unlocated: prim::FloorDiv = op::floordiv(&x,
4).unwrap().try_cast().unwrap();
- assert!(unlocated.loc.as_node::<UnknownLocObj>().is_some());
+ assert!(unlocated.loc.downcast_ref::<UnknownLocObj>().is_some());
// A folded constant is at the unknown location, as C++ `TryConstFold`
builds it.
let folded: IntImm = op::add_at(1, 2,
here.clone()).unwrap().try_cast().unwrap();
- assert!(folded.loc.as_node::<UnknownLocObj>().is_some());
+ assert!(folded.loc.downcast_ref::<UnknownLocObj>().is_some());
let cast: Cast = op::cast_at(&x, PrimType::int(64),
here.clone()).unwrap().try_cast().unwrap();
assert!(cast.loc.same_as(&here));
let call: Call =
@@ -789,14 +789,14 @@ fn cast_folds_and_broadcasts_as_cpp() {
assert!(same(imm("int32", 5).cast(prim("int64")), &imm("int64", 5)));
let node = |value: &Expr, name: &str|
value.cast(prim(name)).into_expr().unwrap();
let broadcast = node(&Expr::from(&f), "int32x4");
- let element = &broadcast.as_node::<BroadcastObj>().unwrap().value;
- assert!(element.as_node::<CastObj>().is_some());
- assert!(node(&imm("int32", -5), "uint8").as_node::<CastObj>().is_some());
- assert!(node(&constant, "int64x4").as_node::<RampObj>().is_some());
- assert!(node(&overflowing, "int64x4").as_node::<CastObj>().is_some());
- assert!(node(&ramp(x.clone().into()),
"int64x4").as_node::<CastObj>().is_some());
- assert!(node(&constant, "float32x4").as_node::<CastObj>().is_some());
- assert!(node(&constant, "int16x4").as_node::<CastObj>().is_some());
+ let element = &broadcast.downcast_ref::<BroadcastObj>().unwrap().value;
+ assert!(element.downcast_ref::<CastObj>().is_some());
+ assert!(node(&imm("int32", -5),
"uint8").downcast_ref::<CastObj>().is_some());
+ assert!(node(&constant, "int64x4").downcast_ref::<RampObj>().is_some());
+ assert!(node(&overflowing, "int64x4").downcast_ref::<CastObj>().is_some());
+ assert!(node(&ramp(x.clone().into()),
"int64x4").downcast_ref::<CastObj>().is_some());
+ assert!(node(&constant, "float32x4").downcast_ref::<CastObj>().is_some());
+ assert!(node(&constant, "int16x4").downcast_ref::<CastObj>().is_some());
let raw: Expr = Cast::new(PrimType::int(64), &x,
UNKNOWN_LOC).unwrap().into();
assert!(same(x.cast(PrimType::int(64)), &raw));
@@ -938,6 +938,6 @@ fn make_const() {
let vector = op::make_const(scalable.clone(), 1).unwrap();
assert!(structural_equal(&vector.ty, &scalable));
let broadcast: prim::Broadcast = vector.try_cast().unwrap();
- let lanes = broadcast.lanes.as_node::<prim::MulObj>().expect("vscale * 4");
+ let lanes = broadcast.lanes.downcast_ref::<prim::MulObj>().expect("vscale
* 4");
assert_eq!(int_value(lanes.b.clone()), 4);
}
diff --git a/rust/tests/ir/type.rs b/rust/tests/ir/type.rs
index 9eea3620a..adc3f0036 100644
--- a/rust/tests/ir/type.rs
+++ b/rust/tests/ir/type.rs
@@ -30,7 +30,7 @@ use tvm_tirx::ir::{
use tvm_tirx::tvm_ffi::{self, 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,
StructuralView, WalkOrder,
+ ObjectCore, ObjectRefCore, Result, ShallowCopy, String as FfiString,
WalkOrder,
};
use super::{scope};
@@ -263,8 +263,8 @@ fn
rust_type_cpp_deep_copy_goes_through_the_shallow_copy_and_the_setters() {
struct IncrementIntegers;
impl MapDispatch for IncrementIntegers {
- fn dispatch_map(&mut self, value: &StructuralView, _: DefRegionKind) ->
Option<Result<Any>> {
- value.cast::<i64>().map(|integer| Ok(Any::from(integer + 1)))
+ fn dispatch_map(&mut self, value: AnyView<'_>, _: DefRegionKind) ->
Option<Result<Any>> {
+ value.try_as::<i64>().map(|integer| Ok(Any::from(integer + 1)))
}
}
@@ -303,7 +303,7 @@ fn rust_type_converts_through_any_and_the_type_hierarchy() {
let _serial = serial();
let value = int_float_pair(1, "");
let as_type: Type = value.clone().into();
- assert!(as_type.as_node::<PairTypeObj>().is_some());
+ assert!(as_type.downcast_ref::<PairTypeObj>().is_some());
let any = Any::from(as_type);
assert_eq!(any.type_index(), PairTypeObj::type_index());
let back = PairType::try_from(any).unwrap();
@@ -315,7 +315,7 @@ fn rust_type_of_a_variable_takes_part_in_ir() {
let _serial = serial();
// The type of a variable, as `PrimType` or `TupleType` is one.
let x = Var::new("x", int_float_pair(7, "seven"), UNKNOWN_LOC);
- assert_eq!(x.ty.as_node::<PairTypeObj>().unwrap().tag, 7);
+ assert_eq!(x.ty.downcast_ref::<PairTypeObj>().unwrap().tag, 7);
assert!(structural_equal(&x.ty, &int_float_pair(7, "other")));
assert!(!structural_equal(&x.ty, &int_float_pair(8, "seven")));
@@ -323,8 +323,8 @@ fn rust_type_of_a_variable_takes_part_in_ir() {
let stmt = Evaluate::new(x.clone(), UNKNOWN_LOC);
let copy =
global("ffi.DeepCopy").call_packed(&[AnyView::from(&stmt)]).unwrap();
let copy = Evaluate::try_from(copy).unwrap();
- let copied = copy.value.as_node::<VarObj>().unwrap();
- let ty = copied.ty.as_node::<PairTypeObj>().unwrap();
+ let copied = copy.value.downcast_ref::<VarObj>().unwrap();
+ let ty = copied.ty.downcast_ref::<PairTypeObj>().unwrap();
assert_eq!((ty.tag, ty.note.as_str()), (7, "seven"));
assert!(structural_equal(&copied.ty, &x.ty));
}
diff --git a/rust/tests/tirx/op.rs b/rust/tests/tirx/op.rs
index 30641eadc..0e8214bd7 100644
--- a/rust/tests/tirx/op.rs
+++ b/rust/tests/tirx/op.rs
@@ -125,7 +125,7 @@ fn fast_erf_follows_the_native_lowering() {
let x = var("x", ty.clone());
let rust = fast_erf_float_expr(&x, bits).unwrap();
assert!(structural_equal(&rust.ty, &ty));
- assert!(rust.as_node::<tvm_tirx::ir::prim::DivObj>().is_some(), "p /
q");
+ assert!(rust.downcast_ref::<tvm_tirx::ir::prim::DivObj>().is_some(),
"p / q");
if let Some(rule) = &rule {
let ty: Type = ty.into();
let call: Expr = Call::new(
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index ee3436bd8..00c692fa1 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -23,7 +23,6 @@ with them. A copy of the library is a second library with the
same blocks, as li
the same crate have.
"""
-import json
import os
import re
import shutil
@@ -92,12 +91,9 @@ def define_ir_library():
@pytest.fixture(scope="module")
def extension_package(tmp_path_factory):
- """The my_ext package of examples/extension, laid out as its wheel
installs it, built against
- rust/ in place of the published tvm-tirx crate."""
- rust_dir = json.dumps(str(RUST_DIR)) # a TOML string, as `--config` takes
- debug_dir = cargo_build(
- EXTENSION_DIR / "Cargo.toml", "--config",
f"patch.crates-io.tvm-tirx.path={rust_dir}"
- )
+ """The my_ext package of examples/extension, built from its declared
source dependency and
+ laid out as its wheel installs it."""
+ debug_dir = cargo_build(EXTENSION_DIR / "Cargo.toml")
root = tmp_path_factory.mktemp("extension")
package = shutil.copytree(EXTENSION_DIR / "python" / "my_ext", root /
"my_ext")
shutil.copy(debug_dir / library_name("my_ext"), package)
