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 3edd4f1d9 [Rust] Borrow references in IntoExpr; bump tvm-ffi to
a21f5fa7 (#203)
3edd4f1d9 is described below
commit 3edd4f1d905481984a278ab56f24f10d80fe3878
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Nov 5 10:16:58 2026 -0500
[Rust] Borrow references in IntoExpr; bump tvm-ffi to a21f5fa7 (#203)
Bumps `3rdparty/tvm-ffi` to a21f5fa7, which brings in apache/tvm-ffi#930
(borrowed upcasts through `AsRef`) and #931 (`downcast_ref`,
`#[dispatch]` and `match_any!` borrowing reference types in place), and
uses them in TIRx.
- **`IntoExpr` borrows references.** The trait has an associated
`Output: Borrow<Expr> + Into<Expr>`. An owned value or a literal
converts into an owned `Expr`, as before. A reference to an expression
type (`&Expr`, `&Var`, `&PrimExpr`, ...) borrows as `&Expr` with no
reference-count change, so `analyzer.can_prove(&cond, ..)` no longer
clones. Read-only functions borrow with `.borrow()`. Constructors take
ownership with `into_owned_expr()`, which clones only a borrowed
expression. Owned values keep a blanket impl. References are listed per
expression type, because blanket impls over `T` and `&T` overlap under
Rust's coherence rules.
- **Typed views.** `typed_expr_view!` gives `PrimExpr` and `IntExpr` the
in-place borrow impls that `#[derive(ObjectRef)]` generates
(`downcast_ref`, `&PrimExpr` handlers and `PrimExpr(ref x)` arms). The
hand-written `AsRef<PrimExpr> for IntExpr` is removed in favor of the
one the upcast macro now generates.
- **Examples.** The IR guide's `match_any!`, `#[dispatch]`, walk and
visit examples borrow reference types (`Add(ref add)`, `&Add`, `&For`,
`&If`) in place of object types, and the guard examples call
`can_prove(&branch.condition, ..)`.
---
3rdparty/tvm-ffi | 2 +-
docs/ir_guides/working_with_ir_in_rust.rst | 14 +--
examples/ir_guides/rust/src/lib.rs | 35 ++++----
rust/src/ir/expr.rs | 133 ++++++++++++++++++++++++-----
rust/src/ir/module.rs | 3 +-
rust/src/ir/prim.rs | 3 +-
rust/src/ir/prim/expr.rs | 21 +++--
rust/src/ir/prim/op.rs | 47 +++++-----
rust/src/ir/prim/vector_expr.rs | 9 +-
rust/src/ir/stmt.rs | 17 ++--
rust/src/macro/typed_expr_view.rs | 34 ++++++++
rust/src/sym/analyzer.rs | 19 +++--
rust/src/sym/int_set.rs | 13 +--
rust/src/sym/iter_affine_map.rs | 28 +++---
rust/src/tirx/analysis/analysis.rs | 5 +-
rust/src/tirx/index_map.rs | 5 +-
rust/src/tirx/layout.rs | 27 ++++--
rust/src/tirx/op.rs | 6 +-
rust/src/tirx/tensor.rs | 3 +-
rust/src/tirx/tile/dispatch.rs | 6 +-
rust/src/tirx/tile/instruction.rs | 3 +-
21 files changed, 295 insertions(+), 138 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 55b4fad05..a21f5fa79 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 55b4fad05c39996f506c170985da2381134ef362
+Subproject commit a21f5fa79329cc63b46c0cf087ebcf7ac4516686
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 9c6741245..575560be5 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -249,8 +249,8 @@ 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. For read-only inspection, ``value.downcast_ref::<AddObj>()``
- returns ``Option<&AddObj>`` without consuming or cloning the reference.
+ increment. For read-only inspection, ``value.downcast_ref::<Add>()``
+ returns ``Option<&Add>`` 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
@@ -355,9 +355,9 @@ generates a ``match_any!`` and a ``Dispatcher``
implementation:
pub fn dispatch(&self, value: AnyView<'_>, nested: bool) ->
Option<String> {
match_any! {
value {
- AddObj(ref node) => Some(self.print_add(node,
nested).into()),
- VarObj(ref node) => Some(self.print_var(node).into()),
- IntImmObj(ref node) => Some(self.print_imm(node).into()),
+ Add(ref node) => Some(self.print_add(node, nested).into()),
+ Var(ref node) => Some(self.print_var(node).into()),
+ IntImm(ref node) => Some(self.print_imm(node).into()),
_ => None,
}
}
@@ -453,8 +453,8 @@ The parameter list of a callback decides what it receives:
- Receives
* - ``|store: TensorStore|``
- Each write, as a managed reference.
- * - ``|store: &TensorStoreObj|``
- - Each write, as a borrowed object, without changing its reference
+ * - ``|store: &TensorStore|``
+ - Each write, borrowed in place, without changing its reference
count.
* - ``|value: AnyView<'_>|``
- Every value the walk reaches.
diff --git a/examples/ir_guides/rust/src/lib.rs
b/examples/ir_guides/rust/src/lib.rs
index b380c1fe7..c285e9e60 100644
--- a/examples/ir_guides/rust/src/lib.rs
+++ b/examples/ir_guides/rust/src/lib.rs
@@ -24,11 +24,8 @@
use std::collections::HashMap;
-use tvm_tirx::ir::prim::{make_const, Add, AddObj, Mul, MulObj};
-use tvm_tirx::ir::{
- Expr, For, ForObj, If, IfObj, IntImmObj, PrimType, TensorStore,
TensorStoreObj, Var, VarObj,
- UNKNOWN_LOC,
-};
+use tvm_tirx::ir::prim::{make_const, Add, Mul};
+use tvm_tirx::ir::{Expr, For, If, IntImm, PrimType, TensorStore, Var,
UNKNOWN_LOC};
use tvm_tirx::script::script;
use tvm_tirx::sym::{Analyzer, ConstraintContext, ProofStrength};
use
tvm_tirx::tirx::structural_map_with_constraint_context::structural_map_with_constraint_context;
@@ -92,10 +89,10 @@ fn casts() -> Result<Strings> {
fn describe(expr: &Expr) -> Result<String> {
let text = match_any! {
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(),
+ Add(ref add) => format!("({} + {})", describe(&add.a)?,
describe(&add.b)?),
+ Mul(ref mul) => format!("({} * {})", describe(&mul.a)?,
describe(&mul.b)?),
+ Var(ref var) => var.name.as_str().to_owned(),
+ IntImm(ref imm) => imm.value.to_string(),
_ => "<other>".to_owned(),
}
};
@@ -116,7 +113,7 @@ struct Printer;
#[dispatch(print, (AnyView<'_>, bool), String)]
impl Printer {
- fn print_add(&self, add: &AddObj, nested: bool) -> String {
+ fn print_add(&self, add: &Add, nested: bool) -> String {
let text = format!("{} + {}", self.show(&add.a, true),
self.show(&add.b, true));
if nested {
format!("({text})")
@@ -125,11 +122,11 @@ impl Printer {
}
}
- fn print_var(&self, var: &VarObj) -> String {
+ fn print_var(&self, var: &Var) -> String {
var.name.as_str().to_owned()
}
- fn print_imm(&self, imm: &IntImmObj) -> String {
+ fn print_imm(&self, imm: &IntImm) -> String {
imm.value.to_string()
}
@@ -220,12 +217,12 @@ struct Counter {
#[dispatch(walk, (AnyView<'_>, DefRegionKind), Result<WalkResult>)]
impl Counter {
- fn walk_var(&mut self, _var: &VarObj, _kind: DefRegionKind) ->
Result<WalkResult> {
+ fn walk_var(&mut self, _var: &Var, _kind: DefRegionKind) ->
Result<WalkResult> {
self.vars += 1;
Ok(WalkResult::Advance)
}
- fn walk_add(&mut self, _add: &AddObj) -> Result<WalkResult> {
+ fn walk_add(&mut self, _add: &Add) -> Result<WalkResult> {
self.adds += 1;
Ok(WalkResult::Advance)
}
@@ -252,7 +249,7 @@ struct LoopNest {
impl LoopNest {
fn visit_for(
&mut self,
- loop_: &ForObj,
+ loop_: &For,
visitor: &mut StructuralVisitorView<'_, Self>,
) -> Result<Option<VisitInterrupt>> {
self.loops.push(loop_.var.name.as_str().to_owned());
@@ -261,7 +258,7 @@ impl LoopNest {
result
}
- fn visit_store(&mut self, store: &TensorStoreObj) ->
Result<Option<VisitInterrupt>> {
+ fn visit_store(&mut self, store: &TensorStore) ->
Result<Option<VisitInterrupt>> {
let loops = self.loops.join(", ");
let write = format!("{} inside [{}]", script(&store.dest)?.as_str(),
loops);
self.writes.push(write);
@@ -319,9 +316,9 @@ fn proven_guards(func: Function) -> Result<Array<bool>> {
let mut proven = Vec::new();
structural_walk_with_constraint_context(
&func,
- |branch: &IfObj, constraints: &ConstraintContext| ->
Result<WalkResult> {
+ |branch: &If, constraints: &ConstraintContext| -> Result<WalkResult> {
analyzer.sync_constraint_context(constraints)?;
- proven.push(analyzer.can_prove(branch.condition.clone(),
ProofStrength::Default)?);
+ proven.push(analyzer.can_prove(&branch.condition,
ProofStrength::Default)?);
Ok(WalkResult::Advance)
},
WalkOrder::PreOrder,
@@ -337,7 +334,7 @@ fn remove_redundant_guards(func: Function) ->
Result<Function> {
func,
|branch: If, constraints: &ConstraintContext| ->
Result<UnchangedOr<Any>> {
analyzer.sync_constraint_context(constraints)?;
- Ok(if analyzer.can_prove(branch.condition.clone(),
ProofStrength::Default)? {
+ Ok(if analyzer.can_prove(&branch.condition,
ProofStrength::Default)? {
branch.then_case.clone().into()
} else {
Unchanged.into()
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index ec3f01d45..7b215a8c7 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -109,6 +109,7 @@ use crate::ir::r#type::Type;
use crate::ir::support;
use std::ops::Deref;
use std::ptr::addr_of_mut;
+use std::borrow::Borrow;
use tvm_ffi::Any;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
@@ -265,16 +266,112 @@ impl Expr {
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
pub trait IntoExpr {
+ /// The expression: an owned [`Expr`] for an owned value or a literal, or
a borrowed
+ /// `&Expr` for a reference, which converts without a reference-count
change.
+ ///
+ /// A function that only reads the expression borrows it with `.borrow()`;
one that keeps
+ /// it takes ownership with `.into()`, which clones only a borrowed
expression.
+ type Output: Borrow<Expr> + Into<Expr>;
+
/// Convert into an expression, or return the first error an
[`ExprResult`] raised.
- fn into_expr(self) -> Result<Expr>;
+ fn into_expr(self) -> Result<Self::Output>;
+
+ /// Convert into an owned expression, as
[`Cow::into_owned`](std::borrow::Cow::into_owned)
+ /// does: an owned value moves, and only a borrowed expression is cloned.
+ fn into_owned_expr(self) -> Result<Expr>
+ where
+ Self: Sized,
+ {
+ self.into_expr().map(Into::into)
+ }
}
-impl<T: Into<Expr>> IntoExpr for T {
+// An owned expression, typed view or literal converts into an owned `Expr`.
+impl<T: AnyCompatible + Into<Expr>> IntoExpr for T {
+ type Output = Expr;
+
fn into_expr(self) -> Result<Expr> {
Ok(self.into())
}
}
+// A reference to an expression, a typed view or a node borrows as `&Expr`
through the upcast's
+// `AsRef`, without a reference-count change.
+macro_rules! impl_into_expr_ref {
+ ($($t:ty),*) => {$(
+ impl<'a> IntoExpr for &'a $t {
+ type Output = &'a Expr;
+
+ fn into_expr(self) -> Result<&'a Expr> {
+ Ok(self.as_ref())
+ }
+ }
+ )*};
+}
+// Every expression type with an upcast to `Expr`. A reference borrows as
`&Expr`; the owned
+// blanket above cannot also cover references, because two blanket impls over
`T` and `&T`
+// overlap for a downstream type. Add a new expression type here.
+impl_into_expr_ref!(
+ Expr,
+ crate::ir::expr::Call,
+ crate::ir::expr::Constant,
+ crate::ir::expr::DataTypeImm,
+ crate::ir::expr::FloatImm,
+ crate::ir::expr::GenericConst,
+ crate::ir::expr::GlobalVar,
+ crate::ir::expr::IntExpr,
+ crate::ir::expr::IntImm,
+ crate::ir::expr::LambdaExpr,
+ crate::ir::expr::OpaqueExpr,
+ crate::ir::expr::PrimExpr,
+ crate::ir::expr::StagingExpr,
+ crate::ir::expr::StringImm,
+ crate::ir::expr::TensorLoad,
+ crate::ir::expr::TensorRegion,
+ crate::ir::expr::Tuple,
+ crate::ir::expr::TupleGetItem,
+ crate::ir::expr::Var,
+ crate::ir::function::BaseFunc,
+ crate::ir::op::Op,
+ crate::ir::prim::expr::Add,
+ crate::ir::prim::expr::And,
+ crate::ir::prim::expr::BitwiseAnd,
+ crate::ir::prim::expr::BitwiseNot,
+ crate::ir::prim::expr::BitwiseOr,
+ crate::ir::prim::expr::BitwiseXor,
+ crate::ir::prim::expr::Cast,
+ crate::ir::prim::expr::Div,
+ crate::ir::prim::expr::EQ,
+ crate::ir::prim::expr::FloorDiv,
+ crate::ir::prim::expr::FloorMod,
+ crate::ir::prim::expr::GE,
+ crate::ir::prim::expr::GT,
+ crate::ir::prim::expr::LE,
+ crate::ir::prim::expr::LShift,
+ crate::ir::prim::expr::LT,
+ crate::ir::prim::expr::Let,
+ crate::ir::prim::expr::Max,
+ crate::ir::prim::expr::Min,
+ crate::ir::prim::expr::Mod,
+ crate::ir::prim::expr::Mul,
+ crate::ir::prim::expr::NE,
+ crate::ir::prim::expr::Not,
+ crate::ir::prim::expr::Or,
+ crate::ir::prim::expr::RShift,
+ crate::ir::prim::expr::Select,
+ crate::ir::prim::expr::Sub,
+ crate::ir::prim::vector_expr::Broadcast,
+ crate::ir::prim::vector_expr::Ramp,
+ crate::ir::prim::vector_expr::Shuffle,
+ crate::sym::analyzer::CanonicalExpr,
+ crate::sym::analyzer::SplitExpr,
+ crate::sym::analyzer::SumExpr,
+ crate::sym::iter_affine_map::IterMapExpr,
+ crate::sym::iter_affine_map::IterSplitExpr,
+ crate::sym::iter_affine_map::IterSumExpr,
+ crate::tirx::function::Function
+);
+
// Rust literals, as C++ converts them into an expression: an `i32` and an
`f32` through the
// constructors `PrimExpr(int32_t)` and `PrimExpr(float)`, the others as
constants of their type.
@@ -379,15 +476,6 @@ crate::r#macro::typed_expr_view! {
tvm_ffi::impl_object_upcast!(PrimExpr => Expr, IntExpr => Expr, IntExpr =>
PrimExpr);
-/// An `IntExpr` borrowed as the `PrimExpr` it is, as C++ `IntExpr` derives
from `PrimExpr`.
-impl AsRef<PrimExpr> for IntExpr {
- fn as_ref(&self) -> &PrimExpr {
- // SAFETY: both views are `repr(transparent)` over an
`ObjectArc<ExprObj>`, and every
- // scalar integer type is a `PrimType`.
- unsafe { &*(self as *const IntExpr as *const PrimExpr) }
- }
-}
-
impl TryFrom<PrimExpr> for IntExpr {
type Error = Error;
@@ -913,7 +1001,7 @@ impl TupleGetItem {
index: i32,
loc: impl Into<Location>,
) -> Result<Self> {
- let tuple_value = tuple_value.into_expr()?;
+ let tuple_value = tuple_value.into_owned_expr()?;
if index < 0 {
let message = format!(
"Index out of bounds: Tuple {} cannot be accessed with
negative index {index}",
@@ -1029,10 +1117,10 @@ impl TensorLoad {
) -> Result<Self> {
static INFER_TYPE: support::TypeAttr =
support::TypeAttr::new("__tvm_tensor_load_infer_type__");
- let (source, loc) = (source.into_expr()?, loc.into());
+ let (source, loc) = (source.into_owned_expr()?, loc.into());
let indices = indices
.into_iter()
- .map(|index| PrimExpr::try_from(index.into_expr()?))
+ .map(|index| PrimExpr::try_from(index.into_owned_expr()?))
.collect::<Result<Array<PrimExpr>>>()?;
let ty_index = AnyView::from(&source.ty).type_index();
let Some(infer) = INFER_TYPE.get(ty_index) else {
@@ -1811,8 +1899,9 @@ impl Call {
ty_args: impl IntoIterator<Item = impl Into<Type>>,
loc: impl Into<Location>,
) -> Result<Self> {
- let (op, attrs, loc) = (op.into_expr()?, attrs.cloned(), loc.into());
- let args =
args.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let (op, attrs, loc) = (op.into_owned_expr()?, attrs.cloned(),
loc.into());
+ let args =
+
args.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let ty_args =
ty_args.into_iter().map(Into::into).collect::<Array<Type>>();
if let Some(ty) = ty {
return Ok(Self::from_fields(loc, ty.clone(), op, args, attrs,
ty_args));
@@ -2089,8 +2178,10 @@ impl LambdaExpr {
///
/// Returns a `ValueError` if the number of arguments is not the number of
variables.
pub fn apply(&self, arguments: impl IntoIterator<Item = impl IntoExpr>) ->
Result<Expr> {
- let arguments =
-
arguments.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let arguments = arguments
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("ir.LambdaExprApply")
.call_tuple((self, &arguments))?
.try_into()
@@ -2152,7 +2243,7 @@ impl Range {
/// Create the range `[begin, end)` (C++ `Range(begin, end, loc)`);
/// [`from_min_extent`](Self::from_min_extent) takes the extent instead of
the end.
pub fn new(begin: impl IntoExpr, end: impl IntoExpr, loc: impl
Into<Location>) -> Result<Self> {
- let (begin, end) = (begin.into_expr()?, end.into_expr()?);
+ let (begin, end) = (begin.into_owned_expr()?, end.into_owned_expr()?);
let extent = Self::get_extent_from_begin_end(begin.clone(), end)?;
Ok(Self::from_fields(begin, extent, loc.into()))
}
@@ -2163,7 +2254,7 @@ impl Range {
extent: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (min_value, extent) = (min_value.into_expr()?,
extent.into_expr()?);
+ let (min_value, extent) = (min_value.into_owned_expr()?,
extent.into_owned_expr()?);
support::prim_dtype(&min_value)?;
support::prim_dtype(&extent)?;
Ok(Self::from_fields(min_value, extent, loc.into()))
@@ -2172,7 +2263,7 @@ impl Range {
/// Return the extent `end - begin` of the range `[begin, end)`, folded
/// (C++ `Range::GetExtentFromBeginEnd`).
pub fn get_extent_from_begin_end(begin: impl IntoExpr, end: impl IntoExpr)
-> Result<Expr> {
- let (begin, end) = (begin.into_expr()?, end.into_expr()?);
+ let (begin, end) = (begin.into_owned_expr()?, end.into_owned_expr()?);
support::prim_dtype(&begin)?;
support::prim_dtype(&end)?;
if let Some(add) = end.downcast_ref::<AddObj>() {
diff --git a/rust/src/ir/module.rs b/rust/src/ir/module.rs
index 7a60a16a0..ba0d91c68 100644
--- a/rust/src/ir/module.rs
+++ b/rust/src/ir/module.rs
@@ -53,6 +53,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::attrs::DictAttrs;
use crate::ir::expr::Expr;
use crate::ir::expr::GlobalVar;
@@ -173,7 +174,7 @@ impl IRModule {
/// Returns a `TypeError` if `expr` is not a function.
pub fn from_expr(expr: impl IntoExpr, functions: &Map<GlobalVar,
BaseFunc>) -> Result<Self> {
tvm_ffi::cached_global_func!("ir.IRModuleFromExpr")
- .call_tuple((&expr.into_expr()?, functions))?
+ .call_tuple((expr.into_expr()?.borrow(), functions))?
.try_into()
}
diff --git a/rust/src/ir/prim.rs b/rust/src/ir/prim.rs
index 81c69fb24..b8fa95599 100644
--- a/rust/src/ir/prim.rs
+++ b/rust/src/ir/prim.rs
@@ -42,6 +42,7 @@ pub use op::*;
pub use vector_expr::*;
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
use tvm_ffi::Result;
@@ -51,6 +52,6 @@ use tvm_ffi::Result;
/// structural equality for that (`prim.expr_deep_equal`).
pub fn expr_deep_equal(lhs: impl IntoExpr, rhs: impl IntoExpr) -> Result<bool>
{
tvm_ffi::cached_global_func!("prim.expr_deep_equal")
- .call_tuple((&lhs.into_expr()?, &rhs.into_expr()?))?
+ .call_tuple((lhs.into_expr()?.borrow(), rhs.into_expr()?.borrow()))?
.try_into()
}
diff --git a/rust/src/ir/prim/expr.rs b/rust/src/ir/prim/expr.rs
index 104971a08..ee8878a1c 100644
--- a/rust/src/ir/prim/expr.rs
+++ b/rust/src/ir/prim/expr.rs
@@ -185,7 +185,7 @@ macro_rules! binary_new {
b: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
let ty = binary_type(stringify!($node), $operands, &a,
&b)?;
Ok(Self::from_fields(loc.into(), ty, a, b))
}
@@ -267,7 +267,7 @@ impl Cast {
value: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (ty, value) = (ty.into(), value.into_expr()?);
+ let (ty, value) = (ty.into(), value.into_owned_expr()?);
let value_dtype = support::prim_dtype(&value)?;
if ty.dtype.lanes != value_dtype.lanes {
let condition = "ty->dtype.lanes == value_expr_ty->dtype.lanes";
@@ -672,7 +672,7 @@ impl BitwiseNot {
///
/// Returns a `TypeError` if `a` is not of an integer or boolean type.
pub fn new(a: impl IntoExpr, loc: impl Into<Location>) -> Result<Self> {
- let a = a.into_expr()?;
+ let a = a.into_owned_expr()?;
if !matches!(support::prim_dtype(&a)?.code, support::INT |
support::UINT | support::BOOL) {
let condition = "a_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) || \
a_ty.MatchesCode(DLDataTypeCode::kDLBool)";
@@ -1743,7 +1743,7 @@ impl And {
/// Returns an `InternalError` if an operand is not boolean, or a
`TypeError` if the operand
/// types differ.
pub fn new(a: impl IntoExpr, b: impl IntoExpr, loc: impl Into<Location>)
-> Result<Self> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
let ty = logical_type(&a, &b)?;
Ok(Self::from_fields(loc.into(), ty, a, b))
}
@@ -1817,7 +1817,7 @@ impl Or {
/// Returns an `InternalError` if an operand is not boolean, or a
`TypeError` if the operand
/// types differ.
pub fn new(a: impl IntoExpr, b: impl IntoExpr, loc: impl Into<Location>)
-> Result<Self> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
let ty = logical_type(&a, &b)?;
Ok(Self::from_fields(loc.into(), ty, a, b))
}
@@ -1884,7 +1884,7 @@ impl Not {
///
/// Returns an `InternalError` if `a` is not boolean.
pub fn new(a: impl IntoExpr, loc: impl Into<Location>) -> Result<Self> {
- let a = a.into_expr()?;
+ let a = a.into_owned_expr()?;
let a_dtype = support::prim_dtype(&a)?;
if a_dtype.code != support::BOOL {
let condition = "a_ty.MatchesCode(DLDataTypeCode::kDLBool)";
@@ -1985,8 +1985,11 @@ impl Select {
false_value: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (condition, true_value, false_value) =
- (condition.into_expr()?, true_value.into_expr()?,
false_value.into_expr()?);
+ let (condition, true_value, false_value) = (
+ condition.into_owned_expr()?,
+ true_value.into_owned_expr()?,
+ false_value.into_owned_expr()?,
+ );
let condition_dtype = support::prim_dtype(&condition)?;
let true_dtype = support::prim_dtype(&true_value)?;
let false_dtype = support::prim_dtype(&false_value)?;
@@ -2101,7 +2104,7 @@ impl Let {
body: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (var, value, body) = (var.into(), value.into_expr()?,
body.into_expr()?);
+ let (var, value, body) = (var.into(), value.into_owned_expr()?,
body.into_owned_expr()?);
let value_dtype = support::prim_dtype(&value)?;
support::prim_dtype(&body)?;
let Some(var_ty) = var.ty.downcast_ref::<PrimTypeObj>() else {
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 3cc62ee4a..c146e76d6 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -239,14 +239,21 @@ impl ExprResult {
}
impl IntoExpr for ExprResult {
+ type Output = Expr;
+
fn into_expr(self) -> Result<Expr> {
self.0.map(Into::into)
}
}
-impl IntoExpr for &ExprResult {
- fn into_expr(self) -> Result<Expr> {
- self.0.clone().map(Into::into)
+impl<'a> IntoExpr for &'a ExprResult {
+ type Output = &'a Expr;
+
+ fn into_expr(self) -> Result<&'a Expr> {
+ match &self.0 {
+ Ok(expr) => Ok(expr.as_ref()),
+ Err(error) => Err(error.clone()),
+ }
}
}
@@ -737,7 +744,7 @@ fn operands(a: impl Operand, b: impl Operand) ->
Result<(Expr, Expr)> {
}
(a, b) => (on_its_own(a), on_its_own(b)),
};
- Ok((a.into_expr()?, b.into_expr()?))
+ Ok((a.into_owned_expr()?, b.into_owned_expr()?))
}
/// The functions without a location: each builds at the unknown location
through its `_at`
@@ -1013,7 +1020,7 @@ pub fn cast_at(
ty: impl Into<PrimType>,
loc: impl Into<Location>,
) -> Result<Expr> {
- let (value, ty, loc) = (value.into_expr()?, ty.into(), loc.into());
+ let (value, ty, loc) = (value.into_owned_expr()?, ty.into(), loc.into());
let (from, to) = (support::prim_dtype(&value)?, ty.dtype);
if from == to {
return Ok(value);
@@ -1145,7 +1152,7 @@ pub fn match_types_at(
b: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<(Expr, Expr)> {
- let (a, b, loc) = (a.into_expr()?, b.into_expr()?, loc.into());
+ let (a, b, loc) = (a.into_owned_expr()?, b.into_owned_expr()?, loc.into());
let (a_dtype, b_dtype) = (support::prim_dtype(&a)?,
support::prim_dtype(&b)?);
if a_dtype == b_dtype {
return Ok((a, b));
@@ -1317,7 +1324,7 @@ pub fn eq_at(a: impl Operand, b: impl Operand, loc: impl
Into<Location>) -> Resu
/// Build [`neg`] at the location `loc`.
pub fn neg_at(a: impl IntoExpr, loc: impl Into<Location>) -> Result<Expr> {
- let (a, loc) = (a.into_expr()?, loc.into());
+ let (a, loc) = (a.into_owned_expr()?, loc.into());
let dtype = support::prim_dtype(&a)?;
if let Some(imm) = a.downcast_ref::<IntImmObj>() {
let value = -imm.value.clone();
@@ -1489,7 +1496,7 @@ pub fn logical_and_at(
b: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Expr> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
check_operands(&a, &b, "&&", "logical AND", "boolean", (&[support::BOOL],
BOOL_CODES))?;
match const_fold::and(&a, &b)? {
Some(folded) => Ok(folded),
@@ -1499,7 +1506,7 @@ pub fn logical_and_at(
/// Build [`logical_or`] at the location `loc`.
pub fn logical_or_at(a: impl IntoExpr, b: impl IntoExpr, loc: impl
Into<Location>) -> Result<Expr> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
check_operands(&a, &b, "||", "logical OR", "boolean", (&[support::BOOL],
BOOL_CODES))?;
match const_fold::or(&a, &b)? {
Some(folded) => Ok(folded),
@@ -1509,7 +1516,7 @@ pub fn logical_or_at(a: impl IntoExpr, b: impl IntoExpr,
loc: impl Into<Location
/// Build [`logical_not`] at the location `loc`.
pub fn logical_not_at(a: impl IntoExpr, loc: impl Into<Location>) ->
Result<Expr> {
- let a = a.into_expr()?;
+ let a = a.into_owned_expr()?;
check_code(&a, "a", &[support::BOOL], BOOL_CODES, || {
format!(
"Expected boolean argument for ! operator (logical NOT), but
received {} of type {}",
@@ -1545,7 +1552,7 @@ fn fold_conditions(
name: &str,
op: fn(Expr, Expr, Location) -> Result<Expr>,
) -> Result<Expr> {
- let mut conds = conds.into_iter().map(IntoExpr::into_expr);
+ let mut conds = conds.into_iter().map(IntoExpr::into_owned_expr);
let Some(first) = conds.next() else {
let message = format!("{name} must take at least 1 argument");
return Err(Error::new(VALUE_ERROR, &message, ""));
@@ -1560,7 +1567,7 @@ pub fn if_then_else_at(
false_value: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Expr> {
- let (cond, loc) = (cond.into_expr()?, loc.into());
+ let (cond, loc) = (cond.into_owned_expr()?, loc.into());
if support::prim_dtype(&cond)?.code != support::BOOL {
let condition = "cond.ty().MatchesCode(DLDataTypeCode::kDLBool)";
let message = "if_then_else only accept the condition to be boolean
type.";
@@ -1717,7 +1724,7 @@ bitwise!(
/// Build [`bitwise_not`] at the location `loc`.
pub fn bitwise_not_at(a: impl IntoExpr, loc: impl Into<Location>) ->
Result<Expr> {
- let a = a.into_expr()?;
+ let a = a.into_owned_expr()?;
check_code(&a, "a", &INTEGER_OR_BOOL, INTEGER_OR_BOOL_CODES, || {
format!(
"Expected integer or boolean argument for ~ operator (bitwise
NOT), but received {} \
@@ -1758,7 +1765,7 @@ macro_rules! unary {
#[doc = concat!("Build [`", stringify!($name), "`] at the location
`loc`.")]
pub fn $at(x: impl IntoExpr, loc: impl Into<Location>) ->
Result<Expr> {
- $imp(x.into_expr()?, loc.into())
+ $imp(x.into_owned_expr()?, loc.into())
}
)+
};
@@ -2031,12 +2038,12 @@ unlocated! {
/// Build [`pow`] at the location `loc`.
pub fn pow_at(x: impl IntoExpr, y: impl IntoExpr, loc: impl Into<Location>) ->
Result<Expr> {
- float_binary(x.into_expr()?, y.into_expr()?, "prim.pow", "power",
loc.into())
+ float_binary(x.into_owned_expr()?, y.into_owned_expr()?, "prim.pow",
"power", loc.into())
}
/// Build [`fmod`] at the location `loc`.
pub fn fmod_at(x: impl IntoExpr, y: impl IntoExpr, loc: impl Into<Location>)
-> Result<Expr> {
- float_binary(x.into_expr()?, y.into_expr()?, "prim.fmod", "fmod",
loc.into())
+ float_binary(x.into_owned_expr()?, y.into_owned_expr()?, "prim.fmod",
"fmod", loc.into())
}
/// Build [`fma`] at the location `loc`.
@@ -2046,9 +2053,9 @@ pub fn fma_at(
z: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Expr> {
- let x = x.into_expr()?;
+ let x = x.into_owned_expr()?;
let ty = x.ty.clone();
- call(Some(&ty), "prim.fma", vec![x, y.into_expr()?, z.into_expr()?],
loc.into())
+ call(Some(&ty), "prim.fma", vec![x, y.into_owned_expr()?,
z.into_owned_expr()?], loc.into())
}
//
---------------------------------------------------------------------------------------------
@@ -2159,9 +2166,9 @@ macro_rules! binary_intrinsics {
y: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Expr> {
- let x = x.into_expr()?;
+ let x = x.into_owned_expr()?;
let ty = x.ty.clone();
- call(Some(&ty), $op, vec![x, y.into_expr()?], loc.into())
+ call(Some(&ty), $op, vec![x, y.into_owned_expr()?], loc.into())
}
)+
};
diff --git a/rust/src/ir/prim/vector_expr.rs b/rust/src/ir/prim/vector_expr.rs
index b0b163556..816f5aeea 100644
--- a/rust/src/ir/prim/vector_expr.rs
+++ b/rust/src/ir/prim/vector_expr.rs
@@ -169,7 +169,8 @@ impl Ramp {
lanes: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (base, stride, lanes) = (base.into_expr()?, stride.into_expr()?,
lanes.into_expr()?);
+ let (base, stride, lanes) =
+ (base.into_owned_expr()?, stride.into_owned_expr()?,
lanes.into_owned_expr()?);
let base_ty = support::prim_type(&base)?;
let stride_ty = support::prim_type(&stride)?;
support::prim_dtype(&lanes)?;
@@ -276,7 +277,7 @@ impl Broadcast {
lanes: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (value, lanes) = (value.into_expr()?, lanes.into_expr()?);
+ let (value, lanes) = (value.into_owned_expr()?,
lanes.into_owned_expr()?);
let value_dtype = support::prim_dtype(&value)?;
support::prim_dtype(&lanes)?;
if !support::is_scalar(value_dtype) {
@@ -362,9 +363,9 @@ impl Shuffle {
loc: impl Into<Location>,
) -> Result<Self> {
let vectors =
-
vectors.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+
vectors.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let indices =
-
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+
indices.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let dtypes = vectors
.iter()
.map(|vector| support::prim_dtype(&vector))
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index 2847ec2a4..5ead61de4 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -244,7 +244,7 @@ impl Bind {
value: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Self> {
- let (var, value) = (var.into(), value.into_expr()?);
+ let (var, value) = (var.into(), value.into_owned_expr()?);
if !support::structural_equal(&value.ty, &var.ty)? {
let condition = "ffi::StructuralEqual()(value->ty, var->ty)";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
condition, ""));
@@ -337,7 +337,7 @@ impl AssertStmt {
message_parts: impl IntoIterator<Item = impl Into<StringImm>>,
loc: impl Into<Location>,
) -> Result<Self> {
- let (condition, error_kind) = (condition.into_expr()?,
error_kind.into());
+ let (condition, error_kind) = (condition.into_owned_expr()?,
error_kind.into());
let message_parts =
message_parts.into_iter().map(Into::into).collect::<Array<StringImm>>();
let condition_ty = support::prim_type(&condition)?;
if condition_ty.dtype.code != support::BOOL {
@@ -477,7 +477,7 @@ impl For {
loc: impl Into<Location>,
) -> Result<Self> {
let (var, begin, end, step) =
- (var.into(), begin.into_expr()?, end.into_expr()?,
step.into_expr()?);
+ (var.into(), begin.into_owned_expr()?, end.into_owned_expr()?,
step.into_owned_expr()?);
let var_expr: Expr = var.clone().into();
let var_ty = support::prim_type(&var_expr)?;
for expr in [&begin, &end, &step] {
@@ -635,7 +635,7 @@ impl While {
body: impl Into<Stmt>,
loc: impl Into<Location>,
) -> Result<Self> {
- let condition = condition.into_expr()?;
+ let condition = condition.into_owned_expr()?;
let condition_dtype = support::prim_dtype(&condition)?;
let body = into_body(body.into())?;
if !support::is_scalar(condition_dtype) {
@@ -723,9 +723,9 @@ impl TensorStore {
) -> Result<Self> {
static VALIDATE: support::TypeAttr =
support::TypeAttr::new("__tvm_tensor_store_validate__");
- let (dest, value) = (dest.into_expr()?, value.into_expr()?);
+ let (dest, value) = (dest.into_owned_expr()?,
value.into_owned_expr()?);
let indices =
-
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+
indices.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
for index in indices.iter() {
support::prim_dtype(&index)?;
}
@@ -850,7 +850,8 @@ impl RegionStmt {
loc: impl Into<Location>,
) -> Result<Self> {
let (op, attrs, body) = (op.into(), attrs.into(),
into_body(body.into())?);
- let args =
args.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let args =
+
args.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let body_params =
body_params.into_iter().map(Into::into).collect::<Array<Var>>();
let result_vars =
result_vars.into_iter().map(Into::into).collect::<Array<Var>>();
let name = op.name.as_str().to_string();
@@ -1110,7 +1111,7 @@ impl If {
else_case: Option<&Stmt>,
loc: impl Into<Location>,
) -> Result<Self> {
- let condition = condition.into_expr()?;
+ let condition = condition.into_owned_expr()?;
support::prim_dtype(&condition)?;
let then_case = into_body(then_case.into())?;
let else_case = else_case.map(|stmt|
into_body(stmt.clone())).transpose()?;
diff --git a/rust/src/macro/typed_expr_view.rs
b/rust/src/macro/typed_expr_view.rs
index 3c0b4d7cd..691bfd4df 100644
--- a/rust/src/macro/typed_expr_view.rs
+++ b/rust/src/macro/typed_expr_view.rs
@@ -110,6 +110,40 @@ macro_rules! typed_expr_view {
}
}
+ // As `#[derive(ObjectRef)]` generates: the view is exactly one object
pointer, so a
+ // cell or an object pointer that holds a `$name` borrows as `&$name`
in place.
+ // SAFETY: the view is `repr(transparent)` over its `ObjectArc`, and
the checked
+ // forms run `check_any_strict` first.
+ unsafe impl tvm_ffi::type_traits::RefFromAnyUnchecked for $name {
+ #[inline(always)]
+ unsafe fn ref_from_any_unchecked(cell: &TVMFFIAny) -> &Self {
+ unsafe { &*(::core::ptr::addr_of!(cell.data_union) as *const
Self) }
+ }
+ }
+
+ unsafe impl tvm_ffi::type_traits::DowncastRefFromAny for $name {
+ #[inline(always)]
+ unsafe fn downcast_ref_from_any(cell: &TVMFFIAny) -> Option<&Self>
{
+ unsafe {
tvm_ffi::type_traits::downcast_ref_from_any_in_place::<Self>(cell) }
+ }
+ }
+
+ unsafe impl tvm_ffi::type_traits::DowncastRefFromObject for $name {
+ #[inline(always)]
+ unsafe fn downcast_ref_from_object(
+ place: &*mut tvm_ffi::tvm_ffi_sys::TVMFFIObject,
+ ) -> Option<&Self> {
+ unsafe {
tvm_ffi::type_traits::downcast_ref_from_object_in_place::<Self>(place) }
+ }
+ }
+
+ impl AsRef<$name> for $name {
+ #[inline(always)]
+ fn as_ref(&self) -> &Self {
+ self
+ }
+ }
+
tvm_ffi::impl_try_from_any!($name);
tvm_ffi::impl_arg_into_ref!($name);
tvm_ffi::impl_into_arg_holder_default!($name);
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index cc69d7938..01e0044b6 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -64,6 +64,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ExitOnDrop;
use crate::ir::expr::Expr;
use crate::ir::expr::ExprObj;
@@ -165,7 +166,7 @@ impl Analyzer {
/// Return the constant bounds of `expr` under the known facts
(`sym.AnalyzerConstIntBound`).
pub fn const_int_bound(&self, expr: impl IntoExpr) ->
Result<ConstIntBound> {
tvm_ffi::cached_global_func!("sym.AnalyzerConstIntBound")
- .call_tuple((self, &expr.into_expr()?))?
+ .call_tuple((self, expr.into_expr()?.borrow()))?
.try_into()
}
@@ -179,7 +180,7 @@ impl Analyzer {
/// Return the modular set `coeff * x + base` of `expr`
(`sym.AnalyzerModularSet`).
pub fn modular_set(&self, expr: impl IntoExpr) -> Result<ModularSet> {
tvm_ffi::cached_global_func!("sym.AnalyzerModularSet")
- .call_tuple((self, &expr.into_expr()?))?
+ .call_tuple((self, expr.into_expr()?.borrow()))?
.try_into()
}
@@ -191,7 +192,7 @@ impl Analyzer {
dom_map: Option<&Map<Var, IntSet>>,
) -> Result<IntSet> {
tvm_ffi::cached_global_func!("sym.AnalyzerIntSet")
- .call_tuple((self, &expr.into_expr()?, dom_map.cloned()))?
+ .call_tuple((self, expr.into_expr()?.borrow(), dom_map.cloned()))?
.try_into()
}
@@ -199,21 +200,21 @@ impl Analyzer {
/// (`sym.AnalyzerSimplify`).
pub fn simplify(&self, expr: impl IntoExpr, steps: i32) -> Result<Expr> {
tvm_ffi::cached_global_func!("sym.AnalyzerSimplify")
- .call_tuple((self, &expr.into_expr()?, steps))?
+ .call_tuple((self, expr.into_expr()?.borrow(), steps))?
.try_into()
}
/// Simplify `expr` by rewrite rules (`sym.AnalyzerRewriteSimplify`).
pub fn rewrite_simplify(&self, expr: impl IntoExpr) -> Result<Expr> {
tvm_ffi::cached_global_func!("sym.AnalyzerRewriteSimplify")
- .call_tuple((self, &expr.into_expr()?))?
+ .call_tuple((self, expr.into_expr()?.borrow()))?
.try_into()
}
/// Simplify `expr` into canonical form (`sym.AnalyzerCanonicalSimplify`).
pub fn canonical_simplify(&self, expr: impl IntoExpr) -> Result<Expr> {
tvm_ffi::cached_global_func!("sym.AnalyzerCanonicalSimplify")
- .call_tuple((self, &expr.into_expr()?))?
+ .call_tuple((self, expr.into_expr()?.borrow()))?
.try_into()
}
@@ -242,14 +243,14 @@ impl Analyzer {
/// Return whether `expr` is provably true (`sym.AnalyzerCanProve`).
pub fn can_prove(&self, expr: impl IntoExpr, strength: ProofStrength) ->
Result<bool> {
tvm_ffi::cached_global_func!("sym.AnalyzerCanProve")
- .call_tuple((self, &expr.into_expr()?, strength as i32))?
+ .call_tuple((self, expr.into_expr()?.borrow(), strength as i32))?
.try_into()
}
/// Return whether `lhs == rhs` is provably true
(`sym.AnalyzerCanProveEqual`).
pub fn can_prove_equal(&self, lhs: impl IntoExpr, rhs: impl IntoExpr) ->
Result<bool> {
tvm_ffi::cached_global_func!("sym.AnalyzerCanProveEqual")
- .call_tuple((self, &lhs.into_expr()?, &rhs.into_expr()?))?
+ .call_tuple((self, lhs.into_expr()?.borrow(),
rhs.into_expr()?.borrow()))?
.try_into()
}
@@ -257,7 +258,7 @@ impl Analyzer {
/// (`sym.AnalyzerTryCompare`).
pub fn try_compare(&self, lhs: impl IntoExpr, rhs: impl IntoExpr) ->
Result<CompareResult> {
let result: i64 =
tvm_ffi::cached_global_func!("sym.AnalyzerTryCompare")
- .call_tuple((self, &lhs.into_expr()?, &rhs.into_expr()?))?
+ .call_tuple((self, lhs.into_expr()?.borrow(),
rhs.into_expr()?.borrow()))?
.try_into()?;
result.try_into()
}
diff --git a/rust/src/sym/int_set.rs b/rust/src/sym/int_set.rs
index 6d97f1f4a..68beb1b5c 100644
--- a/rust/src/sym/int_set.rs
+++ b/rust/src/sym/int_set.rs
@@ -52,6 +52,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::Range;
@@ -106,14 +107,14 @@ impl IntSet {
/// (`sym.IntSetVector`).
pub fn vector(vec: impl IntoExpr) -> Result<Self> {
tvm_ffi::cached_global_func!("sym.IntSetVector")
- .call_tuple((&vec.into_expr()?,))?
+ .call_tuple((vec.into_expr()?.borrow(),))?
.try_into()
}
/// Return the set that holds only `point` (`sym.IntSetSinglePoint`).
pub fn single_point(point: impl IntoExpr) -> Result<Self> {
tvm_ffi::cached_global_func!("sym.IntSetSinglePoint")
- .call_tuple((&point.into_expr()?,))?
+ .call_tuple((point.into_expr()?.borrow(),))?
.try_into()
}
}
@@ -185,7 +186,7 @@ pub fn estimate_region_lower_bound(
) -> Result<Option<Array<IntSet>>> {
let region = region.into_iter().map(Into::into).collect::<Array<Range>>();
tvm_ffi::cached_global_func!("sym.EstimateRegionLowerBound")
- .call_tuple((®ion, var_dom, &predicate.into_expr()?,
analyzer.cloned()))?
+ .call_tuple((®ion, var_dom, predicate.into_expr()?.borrow(),
analyzer.cloned()))?
.try_into()
}
@@ -199,7 +200,7 @@ pub fn estimate_region_strict_bound(
) -> Result<Option<Array<IntSet>>> {
let region = region.into_iter().map(Into::into).collect::<Array<Range>>();
tvm_ffi::cached_global_func!("sym.EstimateRegionStrictBound")
- .call_tuple((®ion, var_dom, &predicate.into_expr()?,
analyzer.cloned()))?
+ .call_tuple((®ion, var_dom, predicate.into_expr()?.borrow(),
analyzer.cloned()))?
.try_into()
}
@@ -213,7 +214,7 @@ pub fn estimate_region_upper_bound(
) -> Result<Array<IntSet>> {
let region = region.into_iter().map(Into::into).collect::<Array<Range>>();
tvm_ffi::cached_global_func!("sym.EstimateRegionUpperBound")
- .call_tuple((®ion, var_dom, &predicate.into_expr()?,
analyzer.cloned()))?
+ .call_tuple((®ion, var_dom, predicate.into_expr()?.borrow(),
analyzer.cloned()))?
.try_into()
}
@@ -244,6 +245,6 @@ pub fn deduce_bound(
relax_map: &Map<Var, IntSet>,
) -> Result<IntSet> {
tvm_ffi::cached_global_func!("sym.DeduceBound")
- .call_tuple((&var.into_expr()?, &cond.into_expr()?, hint_map,
relax_map))?
+ .call_tuple((var.into_expr()?.borrow(), cond.into_expr()?.borrow(),
hint_map, relax_map))?
.try_into()
}
diff --git a/rust/src/sym/iter_affine_map.rs b/rust/src/sym/iter_affine_map.rs
index 84ce1ea6b..dfabbe9ba 100644
--- a/rust/src/sym/iter_affine_map.rs
+++ b/rust/src/sym/iter_affine_map.rs
@@ -58,6 +58,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::expr::Expr;
use crate::ir::expr::ExprObj;
use crate::ir::expr::IntoExpr;
@@ -317,9 +318,9 @@ impl IterSplitExpr {
tvm_ffi::cached_global_func!("sym.IterSplitExpr")
.call_tuple((
&source.into(),
- &lower_factor.into_expr()?,
- &extent.into_expr()?,
- &scale.into_expr()?,
+ lower_factor.into_expr()?.borrow(),
+ extent.into_expr()?.borrow(),
+ scale.into_expr()?.borrow(),
))?
.try_into()
}
@@ -365,8 +366,9 @@ pub fn detect_iter_map(
analyzer: Option<&Analyzer>,
simplify_trivial_iterators: bool,
) -> Result<IterMapResult> {
- let indices =
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- let (predicate, analyzer) = (predicate.into_expr()?, analyzer.cloned());
+ let indices =
+
indices.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ let (predicate, analyzer) = (predicate.into_owned_expr()?,
analyzer.cloned());
tvm_ffi::cached_global_func!("sym.DetectIterMap")
.call_packed(&[
AnyView::from(&indices),
@@ -387,7 +389,7 @@ pub fn normalize_to_iter_sum(
analyzer: Option<&Analyzer>,
) -> Result<IterSumExpr> {
tvm_ffi::cached_global_func!("sym.NormalizeToIterSum")
- .call_tuple((&index.into_expr()?, input_iters, analyzer.cloned()))?
+ .call_tuple((index.into_expr()?.borrow(), input_iters,
analyzer.cloned()))?
.try_into()
}
@@ -401,8 +403,9 @@ pub fn iter_map_simplify(
analyzer: Option<&Analyzer>,
simplify_trivial_iterators: bool,
) -> Result<Array<Expr>> {
- let indices =
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- let (predicate, analyzer) = (predicate.into_expr()?, analyzer.cloned());
+ let indices =
+
indices.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ let (predicate, analyzer) = (predicate.into_owned_expr()?,
analyzer.cloned());
tvm_ffi::cached_global_func!("sym.IterMapSimplify")
.call_packed(&[
AnyView::from(&indices),
@@ -418,7 +421,7 @@ pub fn iter_map_simplify(
/// Return the plain expression of the iterator map `expr`
(`sym.NormalizeIterMapToExpr`).
pub fn normalize_iter_map_to_expr(expr: impl IntoExpr) -> Result<Expr> {
tvm_ffi::cached_global_func!("sym.NormalizeIterMapToExpr")
- .call_tuple((&expr.into_expr()?,))?
+ .call_tuple((expr.into_expr()?.borrow(),))?
.try_into()
}
@@ -437,9 +440,9 @@ pub fn subspace_divide(
simplify_trivial_iterators: bool,
) -> Result<Array<Array<IterMark>>> {
let bindings =
-
bindings.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+
bindings.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let sub_iters =
sub_iters.into_iter().map(Into::into).collect::<Array<Var>>();
- let (predicate, analyzer) = (predicate.into_expr()?, analyzer.cloned());
+ let (predicate, analyzer) = (predicate.into_owned_expr()?,
analyzer.cloned());
tvm_ffi::cached_global_func!("sym.SubspaceDivide")
.call_packed(&[
AnyView::from(&bindings),
@@ -460,7 +463,8 @@ pub fn inverse_affine_iter_map(
outputs: impl IntoIterator<Item = impl IntoExpr>,
) -> Result<Map<Var, Expr>> {
let iter_map =
iter_map.into_iter().map(Into::into).collect::<Array<IterSumExpr>>();
- let outputs =
outputs.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let outputs =
+
outputs.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("sym.InverseAffineIterMap")
.call_tuple((&iter_map, &outputs))?
.try_into()
diff --git a/rust/src/tirx/analysis/analysis.rs
b/rust/src/tirx/analysis/analysis.rs
index b7f8f88d2..f91d282e4 100644
--- a/rust/src/tirx/analysis/analysis.rs
+++ b/rust/src/tirx/analysis/analysis.rs
@@ -48,6 +48,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::Var;
@@ -105,7 +106,7 @@ pub fn detect_linear_equation(
) -> Result<Array<Expr>> {
let vars = vars.into_iter().map(Into::into).collect::<Array<Var>>();
tvm_ffi::cached_global_func!("tirx.analysis.DetectLinearEquation")
- .call_tuple((&expr.into_expr()?, &vars))?
+ .call_tuple((expr.into_expr()?.borrow(), &vars))?
.try_into()
}
@@ -118,6 +119,6 @@ pub fn detect_clip_bound(
) -> Result<Array<Option<Expr>>> {
let vars = vars.into_iter().map(Into::into).collect::<Array<Var>>();
tvm_ffi::cached_global_func!("tirx.analysis.DetectClipBound")
- .call_tuple((&expr.into_expr()?, &vars))?
+ .call_tuple((expr.into_expr()?.borrow(), &vars))?
.try_into()
}
diff --git a/rust/src/tirx/index_map.rs b/rust/src/tirx/index_map.rs
index 10fd495ba..fee79a12e 100644
--- a/rust/src/tirx/index_map.rs
+++ b/rust/src/tirx/index_map.rs
@@ -140,7 +140,7 @@ impl IndexMap {
analyzer: Option<&Analyzer>,
) -> Result<Array<Expr>> {
let indices =
-
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+
indices.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.IndexMapMapIndices")
.call_tuple((self, &indices, analyzer.cloned()))?
.try_into()
@@ -156,7 +156,8 @@ impl IndexMap {
shape: impl IntoIterator<Item = impl IntoExpr>,
analyzer: Option<&Analyzer>,
) -> Result<Array<Expr>> {
- let shape =
shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let shape =
+
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.IndexMapMapShape")
.call_tuple((self, &shape, analyzer.cloned()))?
.try_into()
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index e2eaf2e27..6b4d3e409 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -55,6 +55,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::expr::Expr;
use crate::ir::expr::IntoExpr;
use crate::ir::support;
@@ -142,7 +143,8 @@ impl Layout {
coord: impl IntoIterator<Item = impl IntoExpr>,
analyzer: &Analyzer,
) -> Result<Map<String, Expr>> {
- let coord =
coord.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let coord =
+
coord.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutApply")
.call_tuple((self, &coord, analyzer))?
.try_into()
@@ -156,7 +158,7 @@ impl Layout {
analyzer: &Analyzer,
) -> Result<Map<String, Expr>> {
tvm_ffi::cached_global_func!("tirx.LayoutApplyLinear")
- .call_tuple((self, &coord.into_expr()?, analyzer))?
+ .call_tuple((self, coord.into_expr()?.borrow(), analyzer))?
.try_into()
}
@@ -168,8 +170,10 @@ impl Layout {
shape: impl IntoIterator<Item = impl IntoExpr>,
analyzer: &Analyzer,
) -> Result<Map<String, Expr>> {
- let coord =
coord.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- let shape =
shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let coord =
+
coord.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
+ let shape =
+
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutApplyWithShape")
.call_tuple((self, &coord, &shape, analyzer))?
.try_into()
@@ -194,10 +198,14 @@ impl Layout {
inner_shape: impl IntoIterator<Item = impl IntoExpr>,
) -> Result<Layout> {
let outer = outer.into();
- let outer_shape =
-
outer_shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- let inner_shape =
-
inner_shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let outer_shape = outer_shape
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<Expr>>>()?;
+ let inner_shape = inner_shape
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutTile")
.call_tuple((self, &outer, &outer_shape, &inner_shape))?
.try_into()
@@ -291,7 +299,8 @@ impl TileLayout {
&self,
shape: impl IntoIterator<Item = impl IntoExpr>,
) -> Result<(TileLayout, Array<i64>)> {
- let shape =
shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let shape =
+
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
let pair: Array<Any> =
tvm_ffi::cached_global_func!("tirx.TileLayoutGroup")
.call_tuple((self, &shape))?
.try_into()?;
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index 23c7fc874..b4c175e91 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -185,7 +185,7 @@ pub fn logaddexp(a: impl IntoExpr, b: impl IntoExpr) ->
Result<Expr> {
/// Build [`logaddexp`] with the casts at `loc`.
pub fn logaddexp_at(a: impl IntoExpr, b: impl IntoExpr, loc: impl
Into<Location>) -> Result<Expr> {
- let (a, b) = (a.into_expr()?, b.into_expr()?);
+ let (a, b) = (a.into_owned_expr()?, b.into_owned_expr()?);
for (operand, side) in [(&a, "a"), (&b, "b")] {
if support::prim_dtype(operand)?.code != support::FLOAT {
let condition =
format!("{side}.ty().MatchesCode(DLDataTypeCode::kDLFloat)");
@@ -228,7 +228,7 @@ pub fn reinterpret_at(
value: impl IntoExpr,
loc: impl Into<Location>,
) -> Result<Expr> {
- let (ty, value, loc) = (ty.into(), value.into_expr()?, loc.into());
+ let (ty, value, loc) = (ty.into(), value.into_owned_expr()?, loc.into());
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| {
@@ -350,7 +350,7 @@ pub fn fast_erf_float_expr(arg: impl IntoExpr, bits: u8) ->
Result<Expr> {
-2.13374055278905e-04,
-1.45660718464996e-05,
];
- let x = arg.into_expr()?.min(constant(4.0)).max(constant(-4.0));
+ let x = arg.into_owned_expr()?.min(constant(4.0)).max(constant(-4.0));
let x2 = &x * &x;
// Horner's rule from the highest coefficient down, as C++ writes it out.
let mut p = &x2 * constant(alpha[6]) + constant(alpha[5]);
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index edc26732b..1b3a5231f 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -191,7 +191,8 @@ impl TensorType {
layout: Option<&Layout>,
loc: impl Into<Location>,
) -> Result<Self> {
- let shape =
shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let shape =
+
shape.into_iter().map(IntoExpr::into_owned_expr).collect::<Result<Array<Expr>>>()?;
for extent in shape.iter() {
support::prim_dtype(&extent)?;
}
diff --git a/rust/src/tirx/tile/dispatch.rs b/rust/src/tirx/tile/dispatch.rs
index 8dda2c949..495fb332d 100644
--- a/rust/src/tirx/tile/dispatch.rs
+++ b/rust/src/tirx/tile/dispatch.rs
@@ -200,8 +200,10 @@ impl DispatchContext {
annotations: &Map<String, Any>,
) -> Result<()> {
let tensor = tensor.into();
- let allocated_addr =
-
allocated_addr.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
+ let allocated_addr = allocated_addr
+ .into_iter()
+ .map(IntoExpr::into_owned_expr)
+ .collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.DispatchContextAddAllocTensor").call_tuple((
self,
&tensor,
diff --git a/rust/src/tirx/tile/instruction.rs
b/rust/src/tirx/tile/instruction.rs
index a79ae6df2..55d8d08ec 100644
--- a/rust/src/tirx/tile/instruction.rs
+++ b/rust/src/tirx/tile/instruction.rs
@@ -36,6 +36,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(begin): import-section
+use std::borrow::Borrow;
use crate::ir::attrs::Attrs;
use crate::ir::attrs::AttrsObj;
use crate::ir::expr::Call;
@@ -206,7 +207,7 @@ tvm_ffi::impl_object_upcast!(TrnAttrs => Attrs);
/// Returns an error if `value` is neither.
pub fn as_tensor_region(value: impl IntoExpr) -> Result<TensorRegion> {
tvm_ffi::cached_global_func!("tirx.AsTensorRegion")
- .call_tuple((&value.into_expr()?,))?
+ .call_tuple((value.into_expr()?.borrow(),))?
.try_into()
}
