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 16e90115a [FEAT][RUST] Add PrimExpr and IntExpr typed expression views
(#167)
16e90115a is described below
commit 16e90115a605c8ba8767a064e082ed51e17bf853
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Nov 2 18:46:15 2026 -0500
[FEAT][RUST] Add PrimExpr and IntExpr typed expression views (#167)
## Summary
C++ `PrimExpr` and `IntExpr` are typed views: an `Expr` whose `ty` is a
`PrimType`, or a scalar
integer type for `IntExpr`. They have no node class of their own; the
type is checked where an
expression is converted into the view. The Rust crate bound C++
`PrimExpr` fields as `Expr`,
checked the operand types inside each operator, and had `sig::PrimExpr`
and `sig::IntExpr` as
marker types for operator signatures only.
This PR adds `tvm_tirx::ir::PrimExpr` and `IntExpr` as reference types.
An expression is checked
once, where it becomes a view, and the type then holds for every use.
Before:
```rust
// `index` may be any expression: the function checks its type, and so does
every operator.
fn next_index(index: &Expr) -> Result<(Expr, DLDataType)> {
let next = Expr::try_from(index + 1)?;
let Some(ty) = next.ty.as_node::<PrimTypeObj>() else {
return Err(Error::new(TYPE_ERROR, "expected a PrimExpr", ""));
};
let dtype = ty.dtype;
Ok((next, dtype))
}
let i = Expr::from(var);
let load = TensorLoad::new(&a, [&i], UNKNOWN_LOC)?; // load.indices:
Array<Expr>
let (next, dtype) = next_index(&i)?;
```
After:
```rust
// The signature says what `index` is, and the operators build a `PrimExpr`.
fn next_index(index: &IntExpr) -> Result<(PrimExpr, DLDataType)> {
let next = PrimExpr::try_from(index + 1)?;
let dtype = next.dtype();
Ok((next, dtype))
}
let i: IntExpr = Expr::from(var).try_into()?; // checked here, once
let load = TensorLoad::new(&a, [&i], UNKNOWN_LOC)?; // load.indices:
Array<PrimExpr>
let (next, dtype) = next_index(&i)?;
```
The `?` on `PrimExpr::try_from(index + 1)` reports only the errors of
the operators, such as
mismatched vector lanes; the type of the result needs no check.
---
rust/src/ir/expr.rs | 466 ++++++++++++++++++++++++++++++++++++++++++++++---
rust/src/ir/op_def.rs | 57 +-----
rust/src/ir/prim/op.rs | 80 ++++++---
rust/tests/ir/expr.rs | 126 ++++++++++++-
4 files changed, 627 insertions(+), 102 deletions(-)
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index bad821470..4054d68f6 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -38,8 +38,24 @@
// tvm-ffi-stubgen(custom-new): ir.Var
// tvm-ffi-stubgen(custom-new): ir.LambdaExpr
// tvm-ffi-stubgen(custom-new): ir.Range
+// Constants and loads are typed by a `PrimType`: their allocators take one,
and they convert into
+// a `PrimExpr` without a check.
+// tvm-ffi-stubgen(field): ir.IntImm.ty -> PrimType
+// tvm-ffi-stubgen(field): ir.FloatImm.ty -> PrimType
+// tvm-ffi-stubgen(field): ir.TensorLoad.ty -> PrimType
+// tvm-ffi-stubgen(field): ir.TensorLoad.indices -> Array<PrimExpr>
+// tvm-ffi-stubgen(upcast): ir.IntImm -> PrimExpr
+// tvm-ffi-stubgen(upcast): ir.FloatImm -> PrimExpr
+// tvm-ffi-stubgen(upcast): ir.TensorLoad -> PrimExpr
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
+// tvm-ffi-stubgen(import-object): crate::ir::prim::op::ExprResult
+// tvm-ffi-stubgen(import-object): crate::ir::r#type::PrimTypeObj
+// tvm-ffi-stubgen(import-object): tvm_ffi::Error
+// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCast
+// tvm-ffi-stubgen(import-object): tvm_ffi::object::is_instance_of
+// tvm-ffi-stubgen(import-object): tvm_ffi::tvm_ffi_sys::TVMFFIAny
+// tvm-ffi-stubgen(import-object): tvm_ffi::type_traits::AnyCompatible
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
@@ -66,12 +82,14 @@ use crate::ir::location::Location;
use crate::ir::location::UnknownLoc;
use crate::ir::op::OpObj;
use crate::ir::prim::AddObj;
+use crate::ir::prim::op::ExprResult;
use crate::ir::r#type::AnyType;
use crate::ir::r#type::FuncType;
use crate::ir::r#type::FuncTypeObj;
use crate::ir::r#type::MissingType;
use crate::ir::r#type::MissingTypeObj;
use crate::ir::r#type::PrimType;
+use crate::ir::r#type::PrimTypeObj;
use crate::ir::r#type::StringType;
use crate::ir::r#type::TupleType;
use crate::ir::r#type::TupleTypeObj;
@@ -84,15 +102,20 @@ use tvm_ffi::AnyView;
use tvm_ffi::Array;
use tvm_ffi::BigInt;
use tvm_ffi::DLDataType;
+use tvm_ffi::Error;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
+use tvm_ffi::ObjectRefCast;
use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
use tvm_ffi::UnsafeInit;
+use tvm_ffi::object::is_instance_of;
use tvm_ffi::reflection::drop_if_defined;
+use tvm_ffi::tvm_ffi_sys::TVMFFIAny;
use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
+use tvm_ffi::type_traits::AnyCompatible;
// tvm-ffi-stubgen(end)
// tvm-ffi-stubgen(begin): object/ir.Expr
@@ -330,6 +353,396 @@ unsafe impl UnsafeInit for OpaqueExprObj {
tvm_ffi::impl_object_upcast!(OpaqueExpr => Expr);
// tvm-ffi-stubgen(end)
+//
---------------------------------------------------------------------------------------------
+// Typed views: `PrimExpr` and `IntExpr` (C++ `TypedExpr`)
+
+/// Define a view: a reference to an `ExprObj` whose type `$accepts` accepts,
named `$type_str`
+/// like the C++ type, which an `Any` holding no expression converts into by
`$fallback`.
+macro_rules! typed_expr_view {
+ (
+ $(#[$meta:meta])*
+ $name:ident, $type_str:literal, accepts = $accepts:path, fallback =
$fallback:path
+ ) => {
+ $(#[$meta])*
+ #[repr(transparent)]
+ #[derive(Clone)]
+ pub struct $name {
+ base: ObjectArc<ExprObj>,
+ }
+
+ // SAFETY: `base` is the reference; `from_data` requires its caller to
have checked it.
+ unsafe impl ObjectRefCore for $name {
+ type ContainerType = ExprObj;
+
+ #[inline]
+ fn data(this: &Self) -> &ObjectArc<ExprObj> {
+ &this.base
+ }
+
+ #[inline]
+ fn into_data(this: Self) -> ObjectArc<ExprObj> {
+ this.base
+ }
+
+ /// The view of `data`.
+ ///
+ /// # Safety
+ ///
+ /// The type of `data` is one the view accepts.
+ #[inline]
+ unsafe fn from_data(data: ObjectArc<ExprObj>) -> Self {
+ Self { base: data }
+ }
+ }
+
+ // SAFETY: a view is an `Expr` whose `check_any_strict` accepts fewer
values; the other
+ // conversions are the `Expr` ones.
+ unsafe impl AnyCompatible for $name {
+ #[inline]
+ unsafe fn is_undefined(field: *const Self) -> bool {
+ unsafe {
ObjectArc::is_null(::core::ptr::addr_of!((*field).base)) }
+ }
+
+ fn type_str() -> std::string::String {
+ $type_str.into()
+ }
+
+ fn get_mismatch_type_info(data: &TVMFFIAny) -> std::string::String
{
+ typed_expr_mismatch(data)
+ }
+
+ #[inline(always)]
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ unsafe { <Expr as
AnyCompatible>::copy_to_any_view(src.as_expr(), data) }
+ }
+
+ #[inline(always)]
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ unsafe { <Expr as AnyCompatible>::move_to_any(src.into(),
data) }
+ }
+
+ #[inline(always)]
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ unsafe { prim_dtype(data) }.is_some_and($accepts)
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self
{
+ let expr = unsafe { <Expr as
AnyCompatible>::copy_from_any_view_after_check(data) };
+ Self { base: ObjectRefCore::into_data(expr) }
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ let expr = unsafe { <Expr as
AnyCompatible>::move_from_any_after_check(data) };
+ Self { base: ObjectRefCore::into_data(expr) }
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) ->
std::result::Result<Self, ()> {
+ unsafe {
+ if Self::check_any_strict(data) {
+ Ok(Self::copy_from_any_view_after_check(data))
+ } else {
+ $fallback(data).ok_or(())
+ }
+ }
+ }
+ }
+
+ tvm_ffi::impl_try_from_any!($name);
+ tvm_ffi::impl_arg_into_ref!($name);
+ tvm_ffi::impl_into_arg_holder_default!($name);
+ crate::def_expr_op_over_loads!($name);
+
+ impl From<&$name> for $name {
+ fn from(value: &$name) -> Self {
+ value.clone()
+ }
+ }
+
+ /// The expression object, to read, as [`Expr`] dereferences to it.
+ impl Deref for $name {
+ type Target = ExprObj;
+
+ fn deref(&self) -> &ExprObj {
+ &self.base
+ }
+ }
+
+ impl $name {
+ /// This expression as an [`Expr`], borrowed.
+ #[inline]
+ pub fn as_expr(&self) -> &Expr {
+ // SAFETY: `Expr` is a `repr(C)` struct of the one
`ObjectArc<ExprObj>` the view
+ // holds, and every value of the view is an `Expr`.
+ unsafe { &*(self as *const Self as *const Expr) }
+ }
+
+ /// The type of this expression (C++ `PrimExpr::ty()`), which the
view guarantees.
+ pub fn ty(&self) -> PrimType {
+ self.base.ty.clone().try_cast().expect("a view is typed by a
`PrimType`")
+ }
+
+ /// The dtype of this expression's type.
+ pub fn dtype(&self) -> DLDataType {
+ let ty = self.base.ty.downcast_ref::<PrimTypeObj>();
+ ty.expect("a view is typed by a `PrimType`").dtype
+ }
+ }
+
+ /// The view of an expression whose type it accepts; a TypeError names
any other.
+ impl TryFrom<Expr> for $name {
+ type Error = Error;
+
+ fn try_from(value: Expr) -> Result<Self> {
+ value.try_cast()
+ }
+ }
+
+ impl TryFrom<&Expr> for $name {
+ type Error = Error;
+
+ fn try_from(value: &Expr) -> Result<Self> {
+ value.clone().try_cast()
+ }
+ }
+ };
+}
+
+typed_expr_view! {
+ /// An expression typed by a `PrimType` (C++ `PrimExpr`): a constant, a
variable, an operator
+ /// on scalars or vectors, or a call that returns one.
+ ///
+ /// `PrimExpr` and [`IntExpr`] are typed views (C++ `TypedExpr`). A view
holds an `ir.Expr` as
+ /// [`Expr`] does, and adds what the object class does not say: the type
of the expression. An
+ /// `ir.Call` is a `PrimExpr` when its `ty` is a `PrimType`, and an
`IntExpr` when that type is
+ /// a scalar integer. The views are checked when they are made, so a
function that takes one
+ /// needs no check of its own:
+ ///
+ /// - A view converts into the views and the [`Expr`] it refines for free
(`From`): an
+ /// `IntExpr` into a `PrimExpr` and both into an `Expr`. Borrowed, an
`IntExpr` is a
+ /// `PrimExpr` through `AsRef`, and both views read their object through
`Deref`, as an
+ /// `Expr` does.
+ /// - An `Expr` converts into a view only through a check (`TryFrom`,
+ /// `ObjectRefCast::try_cast`), as an `Any` does. A failed check names
the expression and its
+ /// type, `ir.StringImm[ty=ir.StringType]`.
+ /// - [`IntImm`], [`FloatImm`] and [`TensorLoad`] are typed by a
`PrimType` when they are made,
+ /// so they convert into a `PrimExpr` for free. Rust literals convert as
they convert into an
+ /// `Expr`: an `i32` and an `f32` as C++ `PrimExpr(int32_t)` and
`PrimExpr(float)` convert
+ /// them.
+ /// - The operators build a `PrimExpr`: each step of an [`ExprResult`] is
checked, so
+ /// `PrimExpr::try_from(a + b)?` reports only the errors of the
operators.
+ /// - The view holds its object privately and gives it out only to read,
so its type cannot be
+ /// changed; the one way past the check is the `unsafe`
[`ObjectRefCore::from_data`].
+ ///
+ /// ```compile_fail
+ /// # use tvm_tirx::ir::{Expr, PrimExpr};
+ /// fn widen(expr: Expr) -> PrimExpr {
+ /// expr // An `Expr` is not a `PrimExpr` until it is checked.
+ /// }
+ /// ```
+ ///
+ /// ```compile_fail
+ /// # use tvm_tirx::ir::{IntExpr, PrimExpr};
+ /// fn narrow(expr: PrimExpr) -> IntExpr {
+ /// expr.into() // Nor is a `PrimExpr` an `IntExpr`.
+ /// }
+ /// ```
+ PrimExpr, "ir.PrimExpr", accepts = any_prim, fallback = prim_literal
+}
+
+typed_expr_view! {
+ /// A [`PrimExpr`] of a scalar signed or unsigned integer type of any
width (C++ `IntExpr`),
+ /// such as an index or an extent.
+ IntExpr, "ir.IntExpr", accepts = scalar_int, fallback = int_literal
+}
+
+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;
+
+ fn try_from(value: PrimExpr) -> Result<Self> {
+ value.try_cast()
+ }
+}
+
+impl TryFrom<&PrimExpr> for IntExpr {
+ type Error = Error;
+
+ fn try_from(value: &PrimExpr) -> Result<Self> {
+ value.clone().try_cast()
+ }
+}
+
+/// The built expression, checked to be an integer, or the first error raised
while building it.
+impl TryFrom<ExprResult> for IntExpr {
+ type Error = Error;
+
+ fn try_from(value: ExprResult) -> Result<Self> {
+ PrimExpr::try_from(value)?.try_cast()
+ }
+}
+
+// Rust literals, as they convert into an `Expr`. C++ has the constructors
`PrimExpr(int32_t)` and
+// `PrimExpr(float)`; the other literals are constants of their own type.
+
+/// An `i32` is an `int32` constant (C++ `PrimExpr(int32_t)`).
+impl From<i32> for PrimExpr {
+ fn from(value: i32) -> Self {
+ IntImm::int32(value).into()
+ }
+}
+
+/// An `i64` is an `int64` constant. (C++ converts an `int64_t` held by an
`Any` to `int32` when
+/// the value fits.)
+impl From<i64> for PrimExpr {
+ fn from(value: i64) -> Self {
+ IntImm::int64(value).into()
+ }
+}
+
+/// A `bool` is a boolean constant, as C++ converts a `bool` held by an `Any`.
(C++ promotes a
+/// `bool` argument to `PrimExpr(int32_t)`, an `int32` constant.)
+impl From<bool> for PrimExpr {
+ fn from(value: bool) -> Self {
+ IntImm::bool(value).into()
+ }
+}
+
+/// An `f32` is a `float32` constant (C++ `PrimExpr(float)`).
+impl From<f32> for PrimExpr {
+ fn from(value: f32) -> Self {
+ FloatImm::float32(value).into()
+ }
+}
+
+/// An `f64` is a `float64` constant, as C++ types a `double` next to an
expression. (C++ converts
+/// a `double` held by an `Any` to `float32`.)
+impl From<f64> for PrimExpr {
+ fn from(value: f64) -> Self {
+ FloatImm::float64(value).into()
+ }
+}
+
+/// An `i32` is an `int32` constant. (C++ `IntExpr` converts only an int held
by an `Any`, as
+/// `PrimExpr` does.)
+impl From<i32> for IntExpr {
+ fn from(value: i32) -> Self {
+ int_constant(IntImm::int32(value))
+ }
+}
+
+/// An `i64` is an `int64` constant.
+impl From<i64> for IntExpr {
+ fn from(value: i64) -> Self {
+ int_constant(IntImm::int64(value))
+ }
+}
+
+/// The `IntExpr` of an integer constant, which `int32` and `int64` constants
are.
+fn int_constant(constant: IntImm) -> IntExpr {
+ debug_assert!(scalar_int(PrimExpr::from(&constant).dtype()));
+ // SAFETY: the constant is typed by a scalar integer type.
+ unsafe {
IntExpr::from_data(ObjectRefCore::into_data(Expr::from(constant))) }
+}
+
+// The checks and conversions of the views.
+
+/// Every `PrimType`.
+fn any_prim(_: DLDataType) -> bool {
+ true
+}
+
+/// A scalar signed or unsigned integer type.
+fn scalar_int(dtype: DLDataType) -> bool {
+ dtype.lanes == 1 && matches!(dtype.code, support::INT | support::UINT)
+}
+
+/// The expression `data` holds, if it holds one.
+///
+/// # Safety
+///
+/// `data` is a live FFI value.
+unsafe fn expr_of(data: &TVMFFIAny) -> Option<&ExprObj> {
+ let is_object = data.type_index >=
TVMFFITypeIndex::kTVMFFIStaticObjectBegin as i32;
+ // SAFETY: an object of an `ir.Expr` type is an `ExprObj`.
+ (is_object && is_instance_of::<ExprObj>(data.type_index))
+ .then(|| unsafe { &*(data.data_union.v_obj as *const ExprObj) })
+}
+
+/// The dtype of the expression `data` holds, if it holds one typed by a
`PrimType`.
+///
+/// # Safety
+///
+/// `data` is a live FFI value.
+unsafe fn prim_dtype(data: &TVMFFIAny) -> Option<DLDataType> {
+ Some(unsafe { expr_of(data) }?.ty.downcast_ref::<PrimTypeObj>()?.dtype)
+}
+
+/// What a view names a value it rejects (C++
`TypeTraits<TypedExpr<T>>::GetMismatchTypeInfo`):
+/// an expression with its type, `ir.StringImm[ty=ir.StringType]`, or the type
key of another value.
+fn typed_expr_mismatch(data: &TVMFFIAny) -> std::string::String {
+ // SAFETY: a mismatch is reported for a live value.
+ let Some(expr) = (unsafe { expr_of(data) }) else { return
support::type_key(data.type_index) };
+ 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()),
+ };
+ format!("{}[ty={ty}]", support::type_key(data.type_index))
+}
+
+/// The constant of the int, bool or float `data` holds, as C++ converts an
`Any` into a
+/// `PrimExpr`: a bool is a boolean constant, an int an `int32` constant when
it fits and an
+/// `int64` one otherwise, and a float a `float32` constant.
+///
+/// # Safety
+///
+/// `data` is a live FFI value.
+unsafe fn prim_literal(data: &TVMFFIAny) -> Option<PrimExpr> {
+ let index = data.type_index;
+ // SAFETY: the type index says which member of the union holds the value.
+ unsafe {
+ if index == TVMFFITypeIndex::kTVMFFIBool as i32 {
+ Some(IntImm::bool(data.data_union.v_int64 != 0).into())
+ } else if index == TVMFFITypeIndex::kTVMFFIInt as i32 {
+ Some(int_imm(data.data_union.v_int64).into())
+ } else if index == TVMFFITypeIndex::kTVMFFIFloat as i32 {
+ let value = data.data_union.v_float64;
+ FloatImm::from_dtype(dtype::float(32), value,
UnknownLoc::new()).ok().map(Into::into)
+ } else {
+ None
+ }
+ }
+}
+
+/// The constant of the int `data` holds, as [`prim_literal`] makes it.
+///
+/// # Safety
+///
+/// `data` is a live FFI value.
+unsafe fn int_literal(data: &TVMFFIAny) -> Option<IntExpr> {
+ // SAFETY: an int is held in `v_int64`.
+ (data.type_index == TVMFFITypeIndex::kTVMFFIInt as i32)
+ .then(|| int_constant(int_imm(unsafe { data.data_union.v_int64 })))
+}
+
+/// An `int32` constant when `value` fits, and an `int64` one otherwise (C++
+/// `TypeTraits<IntImm>::ConvertFallbackValue`).
+fn int_imm(value: i64) -> IntImm {
+ match i32::try_from(value) {
+ Ok(value) => IntImm::int32(value),
+ Err(_) => IntImm::int64(value),
+ }
+}
+
// tvm-ffi-stubgen(begin): object/ir.GlobalVar
/// Complete: reflected fields fill [40, 56) exactly.
#[repr(C)]
@@ -592,7 +1005,7 @@ impl TupleGetItem {
pub struct TensorLoadObj {
base: ExprObj,
pub source: Expr,
- pub indices: Array<Expr>,
+ pub indices: Array<PrimExpr>,
}
const _: () = {
@@ -621,8 +1034,8 @@ impl Deref for TensorLoadObj {
}
impl TensorLoadObj {
- pub(crate) fn new(loc: Location, ty: Type, source: Expr, indices:
Array<Expr>) -> Self {
- let base = ExprObj::new(loc, ty);
+ pub(crate) fn new(loc: Location, ty: PrimType, source: Expr, indices:
Array<PrimExpr>) -> Self {
+ let base = ExprObj::new(loc, ty.into());
Self { base, source, indices }
}
}
@@ -631,16 +1044,16 @@ impl TensorLoad {
/// Lossless complete-field allocation.
pub fn from_fields(
loc: impl Into<Location>,
- ty: impl Into<Type>,
+ ty: impl Into<PrimType>,
source: impl Into<Expr>,
- indices: impl Into<Array<Expr>>,
+ indices: impl Into<Array<PrimExpr>>,
) -> Self {
let obj = TensorLoadObj::new(loc.into(), ty.into(), source.into(),
indices.into());
Self { base: ObjectArc::new(obj) }
}
}
-tvm_ffi::impl_object_upcast!(TensorLoad => Expr);
+tvm_ffi::impl_object_upcast!(TensorLoad => Expr, TensorLoad => PrimExpr);
// tvm-ffi-stubgen(end)
impl TensorLoad {
@@ -654,20 +1067,23 @@ impl TensorLoad {
static INFER_TYPE: support::TypeAttr =
support::TypeAttr::new("__tvm_tensor_load_infer_type__");
let (source, loc) = (source.into_expr()?, loc.into());
- let indices =
-
indices.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- for index in indices.iter() {
- support::prim_dtype(&index)?;
- }
+ let indices = indices
+ .into_iter()
+ .map(|index| PrimExpr::try_from(index.into_expr()?))
+ .collect::<Result<Array<PrimExpr>>>()?;
let ty_index = AnyView::from(&source.ty).type_index();
let Some(infer) = INFER_TYPE.get(ty_index) else {
let message =
format!("Type {} does not support TensorLoad",
support::type_key(ty_index));
return Err(support::check_failed(tvm_ffi::TYPE_ERROR,
"infer.has_value()", &message));
};
- let load = Self::from_fields(loc, MissingType::new(), source, indices);
- let ty: PrimType =
support::NativeHook::from_any(infer)?.call_as(&load)?;
- Ok(Self::from_fields(load.loc.clone(), ty, load.source.clone(),
load.indices.clone()))
+ // The hook types a load that has no type yet. Such a load is no
`TensorLoad`, which is
+ // typed by a `PrimType`, so it is made as an `Expr`.
+ let base = ExprObj::new(loc.clone(), MissingType::new().into());
+ let untyped = TensorLoadObj { base, source: source.clone(), indices:
indices.clone() };
+ let untyped = Expr::from(Self { base: ObjectArc::new(untyped) });
+ let ty: PrimType =
support::NativeHook::from_any(infer)?.call_as(&untyped)?;
+ Ok(Self::from_fields(loc, ty, source, indices))
}
}
@@ -774,21 +1190,21 @@ impl Deref for FloatImmObj {
}
impl FloatImmObj {
- pub(crate) fn new(loc: Location, ty: Type, value: f64) -> Self {
- let base = ConstantObj::new(loc, ty);
+ pub(crate) fn new(loc: Location, ty: PrimType, value: f64) -> Self {
+ let base = ConstantObj::new(loc, ty.into());
Self { base, value }
}
}
impl FloatImm {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>, ty: impl Into<Type>, value: f64) ->
Self {
+ pub fn new(loc: impl Into<Location>, ty: impl Into<PrimType>, value: f64)
-> Self {
let obj = FloatImmObj::new(loc.into(), ty.into(), value);
Self { base: ObjectArc::new(obj) }
}
}
-tvm_ffi::impl_object_upcast!(FloatImm => Expr, FloatImm => Constant);
+tvm_ffi::impl_object_upcast!(FloatImm => Expr, FloatImm => Constant, FloatImm
=> PrimExpr);
// tvm-ffi-stubgen(end)
impl FloatImm {
@@ -968,21 +1384,25 @@ impl Deref for IntImmObj {
}
impl IntImmObj {
- pub(crate) fn new(loc: Location, ty: Type, value: BigInt) -> Self {
- let base = ConstantObj::new(loc, ty);
+ pub(crate) fn new(loc: Location, ty: PrimType, value: BigInt) -> Self {
+ let base = ConstantObj::new(loc, ty.into());
Self { base, value }
}
}
impl IntImm {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>, ty: impl Into<Type>, value: impl
Into<BigInt>) -> Self {
+ pub fn new(
+ loc: impl Into<Location>,
+ ty: impl Into<PrimType>,
+ value: impl Into<BigInt>,
+ ) -> Self {
let obj = IntImmObj::new(loc.into(), ty.into(), value.into());
Self { base: ObjectArc::new(obj) }
}
}
-tvm_ffi::impl_object_upcast!(IntImm => Expr, IntImm => Constant);
+tvm_ffi::impl_object_upcast!(IntImm => Expr, IntImm => Constant, IntImm =>
PrimExpr);
// tvm-ffi-stubgen(end)
impl IntImm {
@@ -1045,7 +1465,7 @@ impl From<bool> for IntImm {
}
/// What converts to an expression, possibly failing: anything that converts
into an [`Expr`] (a
-/// node, a `Var`, a Rust `i32`, `f32` and the like), and an
[`ExprResult`](crate::ir::prim::ExprResult), an
+/// node, a `Var`, a Rust `i32`, `f32` and the like), and an [`ExprResult`], an
/// expression under construction, which gives the first error any of its
steps raised. The
/// constructors and functions that can fail take their expression operands as
`impl IntoExpr`, so
/// an `ExprResult` passes straight in and `?` on their result covers it.
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index dcd963940..1620956f3 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -423,10 +423,10 @@ fn call_hook<N, R: Into<Any>>(node: *const N, hook: impl
Fn(&N) -> Result<R>) ->
/// order, [`var_args`](sig::var_args) and [`var_ty_args`](sig::var_ty_args) a
variadic tail of
/// each, and [`call_attrs`](sig::call_attrs) the type of `Call.attrs`. As an
argument type, a node
/// type, such as [`StringImm`](super::StringImm) or [`Var`], accepts an
instance of it, as C++
-/// `TypeTraits<T>::CheckAnyStrict` does, and [`Expr`] or [`Type`] accepts any.
-/// [`PrimExpr`](sig::PrimExpr), [`IntExpr`](sig::IntExpr) and
[`PrimVar`](sig::PrimVar) are the
-/// C++ view types, which also check the expression's type, as is
-/// [`TensorVar`](crate::tirx::TensorVar). Implement [`ArgType`](sig::ArgType)
for another check.
+/// `TypeTraits<T>::CheckAnyStrict` does, and [`Expr`] or [`Type`] accepts
any. The typed views
+/// [`PrimExpr`](super::PrimExpr) and [`IntExpr`](super::IntExpr) also check
the expression's type;
+/// [`PrimVar`](sig::PrimVar) and [`TensorVar`](crate::tirx::TensorVar) check
a variable's type, as
+/// the C++ view types of those names do. Implement [`ArgType`](sig::ArgType)
for another check.
pub mod sig {
use super::*;
use crate::ir::expr::ExprObj;
@@ -455,43 +455,14 @@ pub mod sig {
fn type_str() -> std::string::String {
T::type_str()
}
- }
-
- /// C++ `PrimExpr`: an expression whose type is a `PrimType`.
- pub struct PrimExpr;
-
- impl ArgType for PrimExpr {
- fn check(value: &TVMFFIAny) -> bool {
- prim_dtype(value).is_some()
- }
-
- fn type_str() -> std::string::String {
- "ir.PrimExpr".into()
- }
fn mismatch(value: &TVMFFIAny) -> std::string::String {
- typed_expr_mismatch(value)
+ T::get_mismatch_type_info(value)
}
}
- /// C++ `IntExpr`: a `PrimExpr` of a scalar signed or unsigned integer
type.
- pub struct IntExpr;
-
- impl ArgType for IntExpr {
- fn check(value: &TVMFFIAny) -> bool {
- prim_dtype(value).is_some_and(|dtype| {
- dtype.lanes == 1 && matches!(dtype.code, support::INT |
support::UINT)
- })
- }
-
- fn type_str() -> std::string::String {
- "ir.IntExpr".into()
- }
-
- fn mismatch(value: &TVMFFIAny) -> std::string::String {
- typed_expr_mismatch(value)
- }
- }
+ /// The typed views, which are argument types as they are.
+ pub use crate::ir::expr::{IntExpr, PrimExpr};
/// C++ `PrimVar`: a `Var`, not a subclass, whose type is a `PrimType`.
pub struct PrimVar;
@@ -522,20 +493,6 @@ pub mod sig {
}
}
- fn prim_dtype(value: &TVMFFIAny) -> Option<tvm_ffi::DLDataType> {
- 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.downcast_ref::<PrimTypeObj>() {
- Some(prim) => support::dtype_name(prim.dtype),
- None => support::type_key(AnyView::from(&expr.ty).type_index()),
- };
- format!("{}[ty={ty}]", support::type_key(value.type_index))
- }
-
/// A required value argument of type `T` (C++ `sig::arg<T>`); [`Expr`]
accepts any.
pub fn arg<T: ArgType>(name: &str, doc: &str) -> Arg<T> {
Arg { name: String::from(name), doc: String::from(doc), ty:
PhantomData }
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 869ab9f39..4903cd3c2 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -73,6 +73,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntImm
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntImmObj
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
+// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::TensorLoad
// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::ir::location::Location
@@ -127,6 +128,7 @@ use crate::ir::expr::FloatImmObj;
use crate::ir::expr::IntImm;
use crate::ir::expr::IntImmObj;
use crate::ir::expr::IntoExpr;
+use crate::ir::expr::PrimExpr;
use crate::ir::expr::TensorLoad;
use crate::ir::expr::Var;
use crate::ir::location::Location;
@@ -358,7 +360,9 @@ impl Operand for f64 {
/// An operand next to `peer`: a literal is typed after it.
fn typed_after(operand: OperandKind, peer: &ExprResult) -> ExprResult {
let constant = |value: i64| {
- ExprResult(peer.dtype().and_then(|dtype|
make_const(PrimType::from_dtype(dtype)?, value)))
+ ExprResult::checked(
+ peer.dtype().and_then(|dtype|
make_const(PrimType::from_dtype(dtype)?, value)),
+ )
};
match operand {
OperandKind::Expr(expr) => expr,
@@ -1727,12 +1731,14 @@ binary_intrinsics! {
//
---------------------------------------------------------------------------------------------
// `ExprResult`: expressions under construction
-/// An expression being built by operators: either the expression, or the
first error any step
-/// raised. Later steps pass an error through untouched. An `ExprResult` is an
[`IntoExpr`], so it passes
-/// into any function or constructor that takes one, whose `?` then reports
the error;
-/// `Expr::try_from` takes the expression out.
+/// An expression being built by operators: either the expression, a
[`PrimExpr`], or the first
+/// error any step raised. Later steps pass an error through untouched. An
`ExprResult` is an
+/// [`IntoExpr`], so it passes into any function or constructor that takes
one, whose `?` then
+/// reports the error; `PrimExpr::try_from` takes the expression out without a
further check, and
+/// `Expr::try_from` takes it out as an `Expr`.
///
-/// Registered expression wrappers convert into an `ExprResult`. The operators
+/// Registered expression wrappers convert into an `ExprResult`; one whose
type is not a
+/// `PrimType` is a TypeError there. The operators
/// `+`, `-`, `*`, `&`, `|`, `^`, `<<`, `>>`, unary `-` and `!` are the C++
ones, typed after C++:
/// an integer literal on either side is `MakeConst(other.ty(), v)`, an `f32`
a float32 constant,
/// an `f64` a float64 one. `/` and `%` are not defined: name the division
@@ -1741,13 +1747,23 @@ binary_intrinsics! {
/// Comparisons, `&&` and `||` are [`PrimOps`] methods too, because Rust's
comparison traits
/// return `bool`.
#[derive(Clone)]
-#[must_use = "an `ExprResult` holds the result; pass it on, or convert it with
`Expr::try_from`"]
-pub struct ExprResult(Result<Expr>);
+#[must_use = "an `ExprResult` holds the result; pass it on, or convert it with
`PrimExpr::try_from`"]
+pub struct ExprResult(Result<PrimExpr>);
impl TryFrom<ExprResult> for Expr {
type Error = Error;
/// The built expression, or the first error raised while constructing it.
+ fn try_from(value: ExprResult) -> Result<Self> {
+ value.0.map(Into::into)
+ }
+}
+
+impl TryFrom<ExprResult> for PrimExpr {
+ type Error = Error;
+
+ /// The built expression, which every step checks is a `PrimExpr`, or the
first error raised
+ /// while building it.
fn try_from(value: ExprResult) -> Result<Self> {
value.0
}
@@ -1757,7 +1773,7 @@ impl ExprResult {
/// The dtype a literal next to this expression takes
(`MakeConst(this.ty(), v)`).
fn dtype(&self) -> Result<DLDataType> {
match &self.0 {
- Ok(expr) => support::prim_dtype(expr),
+ Ok(expr) => Ok(expr.dtype()),
Err(error) => Err(error.clone()),
}
}
@@ -1765,19 +1781,27 @@ impl ExprResult {
impl IntoExpr for ExprResult {
fn into_expr(self) -> Result<Expr> {
- self.0
+ self.0.map(Into::into)
}
}
impl IntoExpr for &ExprResult {
fn into_expr(self) -> Result<Expr> {
- self.0.clone()
+ self.0.clone().map(Into::into)
}
}
impl<T: Into<Expr>> From<Result<T>> for ExprResult {
fn from(value: Result<T>) -> Self {
- ExprResult(value.map(Into::into))
+ ExprResult::checked(value.map(Into::into))
+ }
+}
+
+impl ExprResult {
+ /// A step's result, checked to be a `PrimExpr`; an expression of another
type is a TypeError,
+ /// raised where it enters the chain.
+ fn checked(value: Result<Expr>) -> Self {
+ ExprResult(value.and_then(PrimExpr::try_from))
}
}
@@ -1795,7 +1819,7 @@ fn binary(
) -> ExprResult {
let rhs = typed_after(rhs.classify(), &lhs);
match (lhs.0, rhs.0) {
- (Ok(a), Ok(b)) => ExprResult(op(a, b)),
+ (Ok(a), Ok(b)) => ExprResult::checked(op(a.into(), b.into())),
(Err(error), _) | (_, Err(error)) => ExprResult(Err(error)),
}
}
@@ -1891,14 +1915,14 @@ mod operators {
impl Neg for $ty {
type Output = ExprResult;
fn neg(self) -> ExprResult {
- ExprResult(super::neg(ExprResult::from(self)))
+ ExprResult::checked(super::neg(ExprResult::from(self)))
}
}
/// C++ `!` (`logical_not`).
impl Not for $ty {
type Output = ExprResult;
fn not(self) -> ExprResult {
- ExprResult(super::logical_not(ExprResult::from(self)))
+
ExprResult::checked(super::logical_not(ExprResult::from(self)))
}
}
)+
@@ -1914,7 +1938,7 @@ macro_rules! unary_methods {
$(
$(#[$doc])*
fn $name(&self) -> ExprResult {
- ExprResult($name(self.to_expr_result()))
+ ExprResult::checked($name(self.to_expr_result()))
}
)+
};
@@ -1967,11 +1991,11 @@ pub trait PrimOps {
}
/// C++ `a && b` ([`logical_and`]).
fn logical_and(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(logical_and(self.to_expr_result(), rhs))
+ ExprResult::checked(logical_and(self.to_expr_result(), rhs))
}
/// C++ `a || b` ([`logical_or`]).
fn logical_or(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(logical_or(self.to_expr_result(), rhs))
+ ExprResult::checked(logical_or(self.to_expr_result(), rhs))
}
binary_methods! {
@@ -1999,31 +2023,31 @@ pub trait PrimOps {
/// C++ `pow(a, b)` ([`pow`]).
fn pow(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(pow(self.to_expr_result(), rhs))
+ ExprResult::checked(pow(self.to_expr_result(), rhs))
}
/// C++ `fmod(a, b)` ([`fmod`]).
fn fmod(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(fmod(self.to_expr_result(), rhs))
+ ExprResult::checked(fmod(self.to_expr_result(), rhs))
}
/// C++ `atan2(a, b)` ([`atan2`]).
fn atan2(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(atan2(self.to_expr_result(), rhs))
+ ExprResult::checked(atan2(self.to_expr_result(), rhs))
}
/// C++ `nextafter(a, b)` ([`nextafter`]).
fn nextafter(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(nextafter(self.to_expr_result(), rhs))
+ ExprResult::checked(nextafter(self.to_expr_result(), rhs))
}
/// C++ `copysign(a, b)` ([`copysign`]).
fn copysign(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(copysign(self.to_expr_result(), rhs))
+ ExprResult::checked(copysign(self.to_expr_result(), rhs))
}
/// C++ `hypot(a, b)` ([`hypot`]).
fn hypot(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(hypot(self.to_expr_result(), rhs))
+ ExprResult::checked(hypot(self.to_expr_result(), rhs))
}
/// C++ `ldexp(a, b)` ([`ldexp`]).
fn ldexp(&self, rhs: impl IntoExpr) -> ExprResult {
- ExprResult(ldexp(self.to_expr_result(), rhs))
+ ExprResult::checked(ldexp(self.to_expr_result(), rhs))
}
unary_methods! {
@@ -2105,17 +2129,17 @@ pub trait PrimOps {
/// C++ `cast(a, t)` ([`cast`]).
fn cast(&self, ty: impl Into<PrimType>) -> ExprResult {
- ExprResult(cast(self.to_expr_result(), ty))
+ ExprResult::checked(cast(self.to_expr_result(), ty))
}
/// C++ `if_then_else(self, t, f)` ([`if_then_else`]), with `self` the
condition.
fn if_then_else(&self, t: impl IntoExpr, f: impl IntoExpr) -> ExprResult {
- ExprResult(if_then_else(self.to_expr_result(), t, f))
+ ExprResult::checked(if_then_else(self.to_expr_result(), t, f))
}
/// The raw `Select(self, t, f)` ([`Select::new`]), with `self` the
condition.
fn select(&self, t: impl IntoExpr, f: impl IntoExpr) -> ExprResult {
- ExprResult(Select::new(self.to_expr_result(), t, f,
UNKNOWN_LOC).map(Into::into))
+ ExprResult::checked(Select::new(self.to_expr_result(), t, f,
UNKNOWN_LOC).map(Expr::from))
}
}
diff --git a/rust/tests/ir/expr.rs b/rust/tests/ir/expr.rs
index 6c47e2c9d..0db3b8680 100644
--- a/rust/tests/ir/expr.rs
+++ b/rust/tests/ir/expr.rs
@@ -19,10 +19,11 @@
//! Tests of `ir::expr`.
+use tvm_tirx::ir::prim::Broadcast;
use tvm_tirx::ir::*;
use tvm_tirx::tvm_ffi::{
Any, AnyView, Array, BigInt, DLDataType, Function, ObjectRefCast,
ObjectRefCore, Result,
- String as FfiString,
+ String as FfiString, TYPE_ERROR,
};
use super::{int32, int32_var, op, parse_dtype, tensor_var};
@@ -421,3 +422,126 @@ fn calls_agree_with_cpp() {
}
}
}
+
+// Typed views: `PrimExpr` and `IntExpr`.
+
+#[test]
+fn views_check_the_type() {
+ let _serial = serial();
+ let i = Expr::from(int32_var("i"));
+ let x = Expr::from(Var::new("x", PrimType::float(32), UNKNOWN_LOC));
+ let a = Expr::from(tensor_var("A"));
+ let text = Expr::from(StringImm::new("s", UNKNOWN_LOC));
+
+ // An expression is a view only when its type is one the view accepts.
+ let prim = PrimExpr::try_from(&i).unwrap();
+ assert_eq!(prim.dtype(), dtype::int(32));
+ assert_eq!(prim.ty().dtype, dtype::int(32));
+ assert!(prim.same_as(&i));
+ assert_eq!(PrimExpr::try_from(&x).unwrap().dtype(), dtype::float(32));
+ let error = PrimExpr::try_from(&a).err().unwrap();
+ assert_eq!(error.kind(), TYPE_ERROR);
+ assert!(error.message().contains("`ir.Var[ty=tirx.TensorType]` to
`ir.PrimExpr`"), "{error}");
+ let error = PrimExpr::try_from(&text).err().unwrap();
+ assert!(error.message().contains("ir.StringImm[ty=ir.StringType]"),
"{error}");
+
+ let index = IntExpr::try_from(&i).unwrap();
+ assert!(IntExpr::try_from(&x).is_err());
+ assert!(IntExpr::try_from(PrimExpr::from(true)).is_err());
+ // A node type whose constructor takes any type converts only through a
check.
+ let vector = Broadcast::new(&i, int32_lanes(4), UNKNOWN_LOC).unwrap();
+ let vector: PrimExpr = vector.try_cast().unwrap();
+ assert!(IntExpr::try_from(vector).is_err());
+
+ // A view widens for free and reads as the views it refines.
+ let widened: PrimExpr = index.clone().into();
+ let expr: Expr = index.clone().into();
+ assert!(widened.same_as(&expr));
+ assert_eq!(index.dtype(), dtype::int(32));
+ assert!(index.as_expr().same_as(&i));
+ let borrowed: &PrimExpr = index.as_ref();
+ assert!(borrowed.same_as(&i));
+ // A view reads the fields of its object.
+ let ty: &Type = &index.ty;
+ assert_eq!(ty.downcast_ref::<PrimTypeObj>().unwrap().dtype,
dtype::int(32));
+}
+
+#[test]
+fn constants_and_loads_are_views() {
+ let _serial = serial();
+ // Constants and loads are typed by a `PrimType`, so they convert without
a check.
+ assert_eq!(PrimExpr::from(IntImm::int64(1)).dtype(), dtype::int(64));
+ assert_eq!(PrimExpr::from(FloatImm::float32(0.5)).dtype(),
dtype::float(32));
+ let load = TensorLoad::new(tensor_var("A"), [int32_var("i")],
UNKNOWN_LOC).unwrap();
+ let indices: &Array<PrimExpr> = &load.indices;
+ assert_eq!(indices.get(0).unwrap().dtype(), dtype::int(32));
+ assert_eq!(PrimExpr::from(&load).dtype(), dtype::float(32));
+
+ // A load whose index is not a `PrimExpr` is not made.
+ let text = StringImm::new("s", UNKNOWN_LOC);
+ let error = TensorLoad::new(tensor_var("A"), [text],
UNKNOWN_LOC).err().unwrap();
+ assert_eq!(error.kind(), TYPE_ERROR);
+
+ // Rust literals.
+ assert_eq!(PrimExpr::from(1).dtype(), dtype::int(32));
+ assert_eq!(PrimExpr::from(1i64).dtype(), dtype::int(64));
+ assert_eq!(PrimExpr::from(true).dtype(), dtype::bool());
+ assert_eq!(PrimExpr::from(1.5f32).dtype(), dtype::float(32));
+ assert_eq!(PrimExpr::from(1.5f64).dtype(), dtype::float(64));
+ assert_eq!(IntExpr::from(7).dtype(), dtype::int(32));
+ assert_eq!(IntExpr::from(7i64).dtype(), dtype::int(64));
+}
+
+#[test]
+fn any_converts_as_in_cpp() {
+ let _serial = serial();
+ // An expression is checked; a bool, an int or a float becomes a constant,
as C++
+ // `TypeTraits<PrimExpr>` converts it: an int that fits `int32` is one, a
float is `float32`.
+ let i = int32_var("i");
+ assert!(PrimExpr::try_from(Any::from(i.clone())).unwrap().same_as(&i));
+ assert!(PrimExpr::try_from(Any::from(tensor_var("A"))).is_err());
+ assert_eq!(PrimExpr::try_from(Any::from(true)).unwrap().dtype(),
dtype::bool());
+ assert_eq!(PrimExpr::try_from(Any::from(3i64)).unwrap().dtype(),
dtype::int(32));
+ assert_eq!(PrimExpr::try_from(Any::from(1i64 << 40)).unwrap().dtype(),
dtype::int(64));
+ assert_eq!(PrimExpr::try_from(Any::from(0.5f64)).unwrap().dtype(),
dtype::float(32));
+
assert!(PrimExpr::try_from(Any::from(tvm_tirx::tvm_ffi::String::from("s"))).is_err());
+ assert_eq!(IntExpr::try_from(Any::from(3i64)).unwrap().dtype(),
dtype::int(32));
+ assert!(IntExpr::try_from(Any::from(true)).is_err());
+ assert!(IntExpr::try_from(Any::from(0.5f64)).is_err());
+
+ // An array of views checks each element, and converts one of literals
element by element.
+ let mixed: Array<Any> = [Any::from(i.clone()),
Any::from(2i64)].into_iter().collect();
+ let views: Array<PrimExpr> = Any::from(mixed).try_into().unwrap();
+ assert!(views.get(0).unwrap().same_as(&i));
+ assert_eq!(views.get(1).unwrap().dtype(), dtype::int(32));
+ let exprs: Array<Expr> =
+ [Expr::from(i), StringImm::new("s",
UNKNOWN_LOC).into()].into_iter().collect();
+ assert!(Array::<PrimExpr>::try_from(Any::from(exprs)).is_err());
+}
+
+#[test]
+fn operators_take_views() {
+ let _serial = serial();
+ let i = IntExpr::try_from(Expr::from(int32_var("i"))).unwrap();
+ let j = int32_var_expr("j");
+ let sum = PrimExpr::try_from(&i + &j + 1).unwrap();
+ assert_eq!(sum.dtype(), dtype::int(32));
+ assert_eq!(IntExpr::try_from(i.clone() * 2).unwrap().dtype(),
dtype::int(32));
+ // An operand that is not a `PrimExpr` is a TypeError where it enters the
chain.
+ let error = PrimExpr::try_from(Expr::from(tensor_var("A")) +
1).err().unwrap();
+ assert_eq!(error.kind(), TYPE_ERROR);
+ assert!(error.message().contains("ir.Var[ty=tirx.TensorType]"), "{error}");
+ // `ObjectRefCast` downcasts between views as it does between node types.
+ let back: IntExpr = sum.try_cast().unwrap();
+ assert_eq!(back.dtype(), dtype::int(32));
+}
+
+/// `int32` lanes of a vector.
+fn int32_lanes(lanes: i32) -> Expr {
+ IntImm::int32(lanes).into()
+}
+
+/// An `int32` variable as a `PrimExpr`.
+fn int32_var_expr(name: &str) -> PrimExpr {
+ PrimExpr::try_from(Expr::from(int32_var(name))).unwrap()
+}
