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 dadf75073 [Rust] Give value constructors the C++ argument order;
generated full constructors are from_fields (#207)
dadf75073 is described below
commit dadf75073bfd12207c1ea72945a1ed4b7cc667f5
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Nov 5 21:58:34 2026 -0500
[Rust] Give value constructors the C++ argument order; generated full
constructors are from_fields (#207)
A type's `new` is now the constructor with the C++ one's arguments,
`loc` last. The generated complete-field constructor is `from_fields`,
which takes every field in reflection order, inherited fields first, as
`__ffi_init__` does. Eight types exposed the generated one as `new`,
with `loc` first, unlike both C++ and the other types' `new`.
## Changes
- **`IntImm::new(ty, value, loc)` and `FloatImm::new(ty, value, loc)`**
take a `PrimType` and check that the value fits it, like C++ `IntImm(ty,
value, loc)`:
```rust
IntImm::new(PrimType::uint(8), 255, loc)?
FloatImm::new(PrimType::float(16), 0.5, loc)?
```
They replace `from_dtype(DLDataType, ..)`, which is removed: constant
folding builds the `PrimType` with `PrimType::from_dtype` and calls
`new`. `int32`, `int64`, `bool`, `float32` and `float64` are unchanged.
- **Renames to `new`** for the constructors that already had the C++
arguments: `GenericConst::checked`, `While::checked`,
`TupleGetItem::from_tuple` and `DataTypeImm::from_dtype`.
- **`TensorRegion::new(source, region, ty, loc)`** is new, in C++ order.
- **Generated constructors become `from_fields`** through the stubgen
`custom-new` directive: `IntImm`, `FloatImm`, `GenericConst`,
`DataTypeImm`, `TupleGetItem`, `TensorRegion`, `While` and `PrimType`.
`PrimType` is built with `int`, `uint`, `float`, `bfloat`, `bool`,
`void` or `from_dtype`.
- **`UNKNOWN_LOC` for an unknown `loc` argument.** The crate's own calls
pass `UNKNOWN_LOC` instead of `UnknownLoc::new()`, the one spelling the
docs use. `Location::unknown()` remains for a `Location` value, such as
a native call argument.
- **Stubgen import markers** for the `IntoObjectRef` and `Borrow`
imports from #206, so regenerating the bindings keeps them.
---
rust/src/ir/expr.rs | 140 +++++++++++++++++++++++--------------
rust/src/ir/module.rs | 1 +
rust/src/ir/prim/const_fold.rs | 7 +-
rust/src/ir/prim/op.rs | 45 +++++++-----
rust/src/ir/prim/vector_expr.rs | 8 +--
rust/src/ir/stmt.rs | 17 +++--
rust/src/ir/type.rs | 11 +--
rust/src/sym/analyzer.rs | 1 +
rust/src/tirx/analysis/analysis.rs | 1 +
rust/src/tirx/function.rs | 6 +-
rust/src/tirx/op.rs | 7 +-
rust/src/tirx/tensor.rs | 2 +
rust/src/tirx/thread_scope.rs | 2 +
rust/tests/ir/expr.rs | 32 ++++++---
rust/tests/ir/main.rs | 2 +-
rust/tests/ir/op_def.rs | 2 +-
rust/tests/ir/prim/op.rs | 8 ++-
rust/tests/ir/stmt.rs | 6 +-
rust/tests/sym/main.rs | 6 +-
rust/tests/tirx/main.rs | 6 +-
rust/tests/tirx/op.rs | 2 +-
21 files changed, 187 insertions(+), 125 deletions(-)
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index 017f87278..4d027b2a5 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -40,8 +40,9 @@
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
-// `custom-new` reserves `new` for a constructor with the C++ one's semantics
(defaults, implicit
-// conversions, a derived type); the generated `from_fields` allocates the
fields as given.
+// `custom-new` reserves `new` for a constructor with the C++ one's arguments
and semantics
+// (argument order, checks, defaults, implicit conversions, a derived type);
the generated
+// `from_fields` allocates the fields as given, inherited fields first.
// tvm-ffi-stubgen(custom-new): ir.GlobalVar
// tvm-ffi-stubgen(custom-new): ir.Tuple
// tvm-ffi-stubgen(custom-new): ir.TensorLoad
@@ -50,6 +51,12 @@
// tvm-ffi-stubgen(custom-new): ir.Var
// tvm-ffi-stubgen(custom-new): ir.LambdaExpr
// tvm-ffi-stubgen(custom-new): ir.Range
+// tvm-ffi-stubgen(custom-new): ir.TupleGetItem
+// tvm-ffi-stubgen(custom-new): ir.FloatImm
+// tvm-ffi-stubgen(custom-new): ir.IntImm
+// tvm-ffi-stubgen(custom-new): ir.GenericConst
+// tvm-ffi-stubgen(custom-new): ir.DataTypeImm
+// tvm-ffi-stubgen(custom-new): ir.TensorRegion
// 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
@@ -75,7 +82,7 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex
-// tvm-ffi-stubgen(import-object): crate::ir::location::UnknownLoc
+// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(import-object): crate::ir::op::OpObj
// tvm-ffi-stubgen(import-object): crate::ir::prim::AddObj
// tvm-ffi-stubgen(import-object): crate::ir::r#type::AnyType
@@ -96,7 +103,7 @@
use crate::ir::attrs::Attrs;
use crate::ir::dtype;
use crate::ir::location::Location;
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
use crate::ir::op::OpObj;
use crate::ir::prim::AddObj;
use crate::ir::prim::op::ExprResult;
@@ -628,7 +635,7 @@ unsafe fn prim_literal(data: &TVMFFIAny) ->
Option<PrimExpr> {
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)
+ FloatImm::new(PrimType::float(32), value,
UNKNOWN_LOC).ok().map(Into::into)
} else {
None
}
@@ -927,7 +934,7 @@ impl Tuple {
true => MissingType::new().into(),
false => {
let field_types = fields.iter().map(|field|
field.ty.clone());
- TupleType::new(field_types, UnknownLoc::new()).into()
+ TupleType::new(field_types, UNKNOWN_LOC).into()
}
};
Self::from_fields(loc.into(), ty, fields)
@@ -980,7 +987,7 @@ impl TupleGetItemObj {
impl TupleGetItem {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
loc: impl Into<Location>,
ty: impl Into<Type>,
tuple_value: impl Into<Expr>,
@@ -1002,11 +1009,7 @@ impl TupleGetItem {
///
/// Returns an `IndexError` if `index` is negative, or not less than the
number of fields of a
/// tuple of known type.
- pub fn from_tuple(
- tuple_value: impl IntoExpr,
- index: i32,
- loc: impl Into<Location>,
- ) -> Result<Self> {
+ pub fn new(tuple_value: impl IntoExpr, index: i32, loc: impl
Into<Location>) -> Result<Self> {
let tuple_value = tuple_value.into_owned_expr()?;
if index < 0 {
let message = format!(
@@ -1044,7 +1047,7 @@ impl TupleGetItem {
}
None => MissingType::new().into(),
};
- Ok(Self::new(loc.into(), ty, tuple_value, index))
+ Ok(Self::from_fields(loc, ty, tuple_value, index))
}
}
@@ -1255,7 +1258,7 @@ impl FloatImmObj {
impl FloatImm {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>, ty: impl Into<PrimType>, value: f64)
-> Self {
+ pub fn from_fields(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) }
}
@@ -1265,35 +1268,43 @@ tvm_ffi::impl_object_upcast!(FloatImm => Expr, FloatImm
=> Constant, FloatImm =>
// tvm-ffi-stubgen(end)
impl FloatImm {
- /// Create the floating-point constant `value` of type `dtype`, checked to
fit it
- /// (C++ `FloatImm(PrimType(dtype), value, loc)`).
+ /// Create the floating-point constant `value` of type `ty`, checked to
fit it
+ /// (C++ `FloatImm(ty, value, loc)`).
+ ///
+ /// ```
+ /// use tvm_tirx::ir::{FloatImm, PrimType, UNKNOWN_LOC};
+ /// let half = FloatImm::new(PrimType::float(16), 0.5, UNKNOWN_LOC)?;
+ /// assert_eq!(half.value, 0.5);
+ /// assert!(FloatImm::new(PrimType::int(32), 0.5, UNKNOWN_LOC).is_err());
+ /// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+ /// ```
///
/// # Errors
///
- /// Returns a `ValueError` if `dtype` is not a scalar floating-point type,
or a finite `value`
- /// is out of its range.
- pub fn from_dtype(dtype: DLDataType, value: f64, loc: impl Into<Location>)
-> Result<Self> {
- let ty = PrimType::from_dtype(dtype)?;
- check_float_imm(dtype, value)?;
- Ok(Self::new(loc.into(), ty, value))
+ /// Returns a `ValueError` if `ty` is not a scalar floating-point type, or
a finite `value` is
+ /// out of its range.
+ pub fn new(ty: impl Into<PrimType>, value: f64, loc: impl Into<Location>)
-> Result<Self> {
+ let ty = ty.into();
+ check_float_imm(ty.dtype, value)?;
+ Ok(Self::from_fields(loc, ty, value))
}
/// Create the `float32` constant `value` at the unknown location
/// (C++ `FloatImm(PrimType::Float(32), value)`).
pub fn float32(value: f32) -> Self {
- Self::builtin(dtype::float(32), value as f64)
+ Self::builtin(PrimType::float(32), value as f64)
}
/// Create the `float64` constant `value` at the unknown location
/// (C++ `FloatImm(PrimType::Float(64), value)`).
pub fn float64(value: f64) -> Self {
- Self::builtin(dtype::float(64), value)
+ Self::builtin(PrimType::float(64), value)
}
- /// [`FloatImm::from_dtype`] of a value its type holds, which fails only
when the compiler
+ /// [`FloatImm::new`] of a value its type holds, which fails only when the
compiler
/// library is not loaded.
- fn builtin(dtype: DLDataType, value: f64) -> Self {
- support::infallible(Self::from_dtype(dtype, value, UnknownLoc::new()))
+ fn builtin(ty: PrimType, value: f64) -> Self {
+ support::infallible(Self::new(ty, value, UNKNOWN_LOC))
}
}
@@ -1441,7 +1452,7 @@ impl IntImmObj {
impl IntImm {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
loc: impl Into<Location>,
ty: impl Into<PrimType>,
value: impl Into<BigInt>,
@@ -1455,42 +1466,50 @@ tvm_ffi::impl_object_upcast!(IntImm => Expr, IntImm =>
Constant, IntImm => PrimE
// tvm-ffi-stubgen(end)
impl IntImm {
- /// Create the integer constant `value` of the scalar integer or boolean
type `dtype`, checked
- /// to fit it (C++ `IntImm(PrimType(dtype), value, loc)`).
+ /// Create the integer constant `value` of the scalar integer or boolean
type `ty`, checked to
+ /// fit it (C++ `IntImm(ty, value, loc)`).
+ ///
+ /// ```
+ /// use tvm_tirx::ir::{IntImm, PrimType, UNKNOWN_LOC};
+ /// let byte = IntImm::new(PrimType::uint(8), 255, UNKNOWN_LOC)?;
+ /// assert_eq!(byte.value.to_string(), "255");
+ /// assert!(IntImm::new(PrimType::uint(8), 256, UNKNOWN_LOC).is_err());
+ /// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+ /// ```
///
/// # Errors
///
- /// Returns a `ValueError` if `dtype` is not a scalar integer or boolean
type, or `value` is out
- /// of its range.
- pub fn from_dtype(
- dtype: DLDataType,
+ /// Returns a `ValueError` if `ty` is not a scalar integer or boolean
type, or `value` is out of
+ /// its range.
+ pub fn new(
+ ty: impl Into<PrimType>,
value: impl Into<BigInt>,
loc: impl Into<Location>,
) -> Result<Self> {
- let (value, ty) = (value.into(), PrimType::from_dtype(dtype)?);
- check_int_imm(dtype, &value)?;
- Ok(Self::new(loc.into(), ty, value))
+ let (ty, value) = (ty.into(), value.into());
+ check_int_imm(ty.dtype, &value)?;
+ Ok(Self::from_fields(loc, ty, value))
}
/// Create the `int32` constant `value` at the unknown location (C++
`IntImm::Int32(value)`).
pub fn int32(value: i32) -> Self {
- Self::builtin(dtype::int(32), value)
+ Self::builtin(PrimType::int(32), value)
}
/// Create the `int64` constant `value` at the unknown location (C++
`IntImm::Int64(value)`).
pub fn int64(value: i64) -> Self {
- Self::builtin(dtype::int(64), value)
+ Self::builtin(PrimType::int(64), value)
}
/// Create the boolean constant `value` at the unknown location (C++
`IntImm::Bool(value)`).
pub fn bool(value: bool) -> Self {
- Self::builtin(dtype::bool(), value)
+ Self::builtin(PrimType::bool(), value)
}
- /// [`IntImm::from_dtype`] of a value its type holds, which fails only
when the compiler
+ /// [`IntImm::new`] of a value its type holds, which fails only when the
compiler
/// library is not loaded.
- fn builtin(dtype: DLDataType, value: impl Into<BigInt>) -> Self {
- support::infallible(Self::from_dtype(dtype, value, UnknownLoc::new()))
+ fn builtin(ty: PrimType, value: impl Into<BigInt>) -> Self {
+ support::infallible(Self::new(ty, value, UNKNOWN_LOC))
}
}
@@ -1632,7 +1651,11 @@ impl GenericConstObj {
impl GenericConst {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>, ty: impl Into<Type>, value: impl
Into<Any>) -> Self {
+ pub fn from_fields(
+ loc: impl Into<Location>,
+ ty: impl Into<Type>,
+ value: impl Into<Any>,
+ ) -> Self {
let obj = GenericConstObj::new(loc.into(), ty.into(), value.into());
Self { base: ObjectArc::new(obj) }
}
@@ -1649,7 +1672,7 @@ impl GenericConst {
///
/// Returns a `TypeError` if `ty` is a [`MissingType`], or `value` is an
integer, a boolean, a
/// float or a string.
- pub fn checked(
+ pub fn new(
ty: impl Into<Type>,
value: impl Into<Any>,
loc: impl Into<Location>,
@@ -1672,7 +1695,7 @@ impl GenericConst {
let message = "Primitive literals use IntImm, FloatImm, or
StringImm";
return Err(support::check_failed(tvm_ffi::TYPE_ERROR, condition,
message));
}
- Ok(Self::new(loc.into(), ty, value))
+ Ok(Self::from_fields(loc, ty, value))
}
}
@@ -1721,7 +1744,7 @@ impl DataTypeImmObj {
impl DataTypeImm {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
loc: impl Into<Location>,
ty: impl Into<Type>,
value: impl Into<DLDataType>,
@@ -1737,9 +1760,9 @@ tvm_ffi::impl_object_upcast!(DataTypeImm => Expr,
DataTypeImm => Constant);
impl DataTypeImm {
/// Create the data type `value` as a constant of type [`AnyType`] (C++
`DataTypeImm(value,
/// loc)`).
- pub fn from_dtype(value: DLDataType, loc: impl Into<Location>) -> Self {
- let ty = AnyType::new(UnknownLoc::new());
- Self::new(loc.into(), ty, value)
+ pub fn new(value: DLDataType, loc: impl Into<Location>) -> Self {
+ let ty = AnyType::new(UNKNOWN_LOC);
+ Self::from_fields(loc, ty, value)
}
}
@@ -1804,7 +1827,7 @@ tvm_ffi::impl_object_upcast!(StringImm => Expr, StringImm
=> Constant);
impl StringImm {
/// Create the string constant `value`, of type [`StringType`] (C++
`StringImm(value, loc)`).
pub fn new(value: &str, loc: impl Into<Location>) -> Self {
- let ty = StringType::new(UnknownLoc::new());
+ let ty = StringType::new(UNKNOWN_LOC);
Self::from_fields(loc.into(), ty, String::from(value))
}
}
@@ -2008,7 +2031,7 @@ impl TensorRegionObj {
impl TensorRegion {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
loc: impl Into<Location>,
ty: impl Into<Type>,
source: impl Into<Expr>,
@@ -2022,7 +2045,18 @@ impl TensorRegion {
tvm_ffi::impl_object_upcast!(TensorRegion => Expr);
// tvm-ffi-stubgen(end)
-// `ir.TensorRegion` stores its arguments as given: the generated
`TensorRegion::new` is the same.
+impl TensorRegion {
+ /// Create the region `region` of `source`, of type `ty` (C++
`TensorRegion(source, region, ty,
+ /// loc)`).
+ pub fn new(
+ source: impl Into<Expr>,
+ region: impl Into<Array<Range>>,
+ ty: impl Into<Type>,
+ loc: impl Into<Location>,
+ ) -> Self {
+ Self::from_fields(loc, ty, source, region)
+ }
+}
// tvm-ffi-stubgen(begin): object/ir.Var
/// Complete: reflected fields fill [40, 56) exactly.
diff --git a/rust/src/ir/module.rs b/rust/src/ir/module.rs
index a8c2f0f73..45b38253a 100644
--- a/rust/src/ir/module.rs
+++ b/rust/src/ir/module.rs
@@ -52,6 +52,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ir::attrs::DictAttrs;
diff --git a/rust/src/ir/prim/const_fold.rs b/rust/src/ir/prim/const_fold.rs
index e97563a8e..44fdb18c1 100644
--- a/rust/src/ir/prim/const_fold.rs
+++ b/rust/src/ir/prim/const_fold.rs
@@ -30,7 +30,8 @@
use tvm_ffi::{BigInt, DLDataType, ObjectRefCore, Result, INTERNAL_ERROR};
use crate::ir::expr::{FloatImm, FloatImmObj, IntImm, IntImmObj, PrimExpr};
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
+use crate::ir::r#type::PrimType;
use crate::ir::support;
/// The folded expression, or `None` when there is nothing to fold.
@@ -90,7 +91,7 @@ fn float(expr: &PrimExpr) -> Option<f64> {
/// C++ `IntImm(dtype, value)`.
fn int_imm(dtype: DLDataType, value: BigInt) -> Folded {
- Ok(Some(IntImm::from_dtype(dtype, value, UnknownLoc::new())?.into()))
+ Ok(Some(IntImm::new(PrimType::from_dtype(dtype)?, value,
UNKNOWN_LOC)?.into()))
}
/// C++ `IntImm(result_ty, value)` of an exact result, or nothing to fold when
it is not a value of
@@ -104,7 +105,7 @@ fn exact(dtype: DLDataType, value: BigInt) -> Folded {
/// C++ `FloatImm(dtype, value)`.
fn float_imm(dtype: DLDataType, value: f64) -> Folded {
- Ok(Some(FloatImm::from_dtype(dtype, value, UnknownLoc::new())?.into()))
+ Ok(Some(FloatImm::new(PrimType::from_dtype(dtype)?, value,
UNKNOWN_LOC)?.into()))
}
/// C++ `IntImm::Bool(value)`.
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index be39bb7bb..e4e6d7e10 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -836,7 +836,7 @@ pub(crate) fn make_const_scalar(
let message = format!("cannot make uint from negative value
{value}");
return Err(support::check_failed(INTERNAL_ERROR, "value >= 0",
&message));
}
- return Ok(IntImm::from_dtype(dtype, value, loc.clone())?.into());
+ return Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
}
if code == support::FLOAT
|| is_float8(code)
@@ -848,7 +848,7 @@ pub(crate) fn make_const_scalar(
ConstValue::Int(value) => value.to_f64()?,
ConstValue::Float(value) => value,
};
- return Ok(FloatImm::from_dtype(dtype, value, loc.clone())?.into());
+ return Ok(FloatImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
}
let ty = PrimType::from_dtype(dtype)?;
let message = format!("cannot make const for type {}", support::repr(&ty));
@@ -931,10 +931,10 @@ fn fold_cast(value: &PrimExpr, to: DLDataType) ->
Result<Option<PrimExpr>> {
_ => None,
};
let int = |value: BigInt, loc: &Location| -> Result<Option<PrimExpr>> {
- Ok(Some(IntImm::from_dtype(to, value, loc.clone())?.into()))
+ Ok(Some(IntImm::new(PrimType::from_dtype(to)?, value,
loc.clone())?.into()))
};
let float = |value: f64, loc: &Location| -> Result<Option<PrimExpr>> {
- Ok(Some(FloatImm::from_dtype(to, value, loc.clone())?.into()))
+ Ok(Some(FloatImm::new(PrimType::from_dtype(to)?, value,
loc.clone())?.into()))
};
if let Some(imm) = value.downcast_ref::<IntImmObj>() {
if integer {
@@ -1345,11 +1345,11 @@ pub fn neg_at(a: impl IntoExpr<PrimExpr>, loc: impl
Into<Location>) -> Result<Pr
if let Some(imm) = a.downcast_ref::<IntImmObj>() {
let value = -imm.value.clone();
if const_fold::is_representable(&value, dtype) {
- return Ok(IntImm::from_dtype(dtype, value, loc)?.into());
+ return Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc)?.into());
}
}
if let Some(imm) = a.downcast_ref::<FloatImmObj>() {
- return Ok(FloatImm::from_dtype(dtype, -imm.value, loc)?.into());
+ return Ok(FloatImm::new(PrimType::from_dtype(dtype)?, -imm.value,
loc)?.into());
}
sub(make_const_at(PrimType::from_dtype(dtype)?, 0, loc)?, a)
}
@@ -1636,7 +1636,7 @@ fn fold_shift(
}
if let (Some(x), Some(y)) = (pa, pb) {
if let Some(value) = shift(x, y)? {
- return Ok(Some(IntImm::from_dtype(dtype, value,
loc.clone())?.into()));
+ return Ok(Some(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into()));
}
}
if pb.is_some_and(BigInt::is_zero) {
@@ -1729,7 +1729,7 @@ macro_rules! bitwise {
if let (true, Some(x), Some(y)) =
(index, const_fold::int(&a), const_fold::int(&b))
{
- return Ok(IntImm::from_dtype(dtype, x $op y, loc)?.into());
+ return Ok(IntImm::new(PrimType::from_dtype(dtype)?, x $op
y, loc)?.into());
}
Ok($node::new(a, b, loc)?.into())
}
@@ -1788,7 +1788,9 @@ fn rounding(
return Ok(x);
}
if let Some(imm) = x.downcast_ref::<FloatImmObj>() {
- return Ok(FloatImm::from_dtype(dtype, fold(imm.value),
imm.loc.clone())?.into());
+ return Ok(
+ FloatImm::new(PrimType::from_dtype(dtype)?, fold(imm.value),
imm.loc.clone())?.into()
+ );
}
let ty = x.ty.clone();
call(Some(&ty), op, vec![x], loc)
@@ -1864,7 +1866,10 @@ fn absolute(x: PrimExpr, loc: Location) ->
Result<PrimExpr> {
match dtype.code {
support::INT => integer_abs(x, loc),
support::FLOAT | support::BFLOAT => match
x.downcast_ref::<FloatImmObj>() {
- Some(imm) => Ok(FloatImm::from_dtype(dtype, imm.value.abs(),
imm.loc.clone())?.into()),
+ Some(imm) => {
+ Ok(FloatImm::new(PrimType::from_dtype(dtype)?,
imm.value.abs(), imm.loc.clone())?
+ .into())
+ }
None => {
let ty = x.ty.clone();
call(Some(&ty), "prim.fabs", vec![x], loc)
@@ -1891,7 +1896,7 @@ fn integer_abs(x: PrimExpr, loc: Location) ->
Result<PrimExpr> {
false => imm.value.clone(),
};
if const_fold::is_representable(&value, dtype) {
- return Ok(IntImm::from_dtype(dtype, value,
imm.loc.clone())?.into());
+ return Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
imm.loc.clone())?.into());
}
}
let zero = make_const(PrimType::from_dtype(dtype)?, 0)?;
@@ -1975,8 +1980,10 @@ pub fn max_value_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result
let dtype = ty.dtype;
check_scalar(dtype)?;
let (code, bits) = (dtype.code, dtype.bits);
- let int = |value: BigInt| Ok(IntImm::from_dtype(dtype, value,
loc.clone())?.into());
- let float = |value: f64| Ok(FloatImm::from_dtype(dtype, value,
loc.clone())?.into());
+ let int =
+ |value: BigInt| Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
+ let float =
+ |value: f64| Ok(FloatImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
match code {
support::INT if bits == 64 => int(BigInt::from(i64::MAX)),
support::INT if bits < 64 => int(BigInt::from((1i64 << (bits - 1)) -
1)),
@@ -2003,8 +2010,10 @@ pub fn min_value_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result
let dtype = ty.dtype;
check_scalar(dtype)?;
let (code, bits) = (dtype.code, dtype.bits);
- let int = |value: BigInt| Ok(IntImm::from_dtype(dtype, value,
loc.clone())?.into());
- let float = |value: f64| Ok(FloatImm::from_dtype(dtype, value,
loc.clone())?.into());
+ let int =
+ |value: BigInt| Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
+ let float =
+ |value: f64| Ok(FloatImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into());
match code {
support::INT if bits == 64 => int(BigInt::from(i64::MIN)),
support::INT if bits < 64 => int(BigInt::from(-(1i64 << (bits - 1)))),
@@ -2031,9 +2040,11 @@ pub fn infinity_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result<
let dtype = ty.dtype;
check_scalar(dtype)?;
match (dtype.code, dtype.bits) {
- (support::FLOAT, 64) => Ok(FloatImm::from_dtype(dtype, f64::INFINITY,
loc)?.into()),
+ (support::FLOAT, 64) => {
+ Ok(FloatImm::new(PrimType::from_dtype(dtype)?, f64::INFINITY,
loc)?.into())
+ }
(support::FLOAT, 32 | 16) => {
- Ok(FloatImm::from_dtype(dtype, f64::from(f32::INFINITY),
loc)?.into())
+ Ok(FloatImm::new(PrimType::from_dtype(dtype)?,
f64::from(f32::INFINITY), loc)?.into())
}
_ => {
let message = format!("Cannot decide infinity for type {}",
support::repr(&ty));
diff --git a/rust/src/ir/prim/vector_expr.rs b/rust/src/ir/prim/vector_expr.rs
index 6def1b3a2..cf894924e 100644
--- a/rust/src/ir/prim/vector_expr.rs
+++ b/rust/src/ir/prim/vector_expr.rs
@@ -72,7 +72,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::CallObj
// 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::location::UnknownLoc
+// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(import-object): crate::ir::prim::expr::Mul
// tvm-ffi-stubgen(import-object): crate::ir::prim::expr::MulObj
// tvm-ffi-stubgen(import-object): crate::ir::r#type::PrimType
@@ -95,7 +95,7 @@ use crate::ir::expr::IntImmObj;
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
use crate::ir::location::Location;
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
use crate::ir::prim::expr::Mul;
use crate::ir::prim::expr::MulObj;
use crate::ir::prim::op::cast;
@@ -467,10 +467,10 @@ pub(crate) fn vscale_lanes(factor: i32) ->
Result<PrimExpr> {
Vec::<Expr>::new(),
None,
Vec::<Type>::new(),
- UnknownLoc::new(),
+ UNKNOWN_LOC,
)?
.into();
- Ok(Mul::new(PrimExpr::try_from(vscale)?, IntImm::int32(factor),
UnknownLoc::new())?.into())
+ Ok(Mul::new(PrimExpr::try_from(vscale)?, IntImm::int32(factor),
UNKNOWN_LOC)?.into())
}
/// `factor` when `lanes` is `factor * vscale` or `vscale * factor` (C++
`ExtractVscaleFactor`).
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index abba8fa2b..7eeade055 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -60,6 +60,7 @@
// tvm-ffi-stubgen(custom-new): ir.If
// tvm-ffi-stubgen(custom-new): ir.Evaluate
// tvm-ffi-stubgen(custom-new): ir.Return
+// tvm-ffi-stubgen(custom-new): ir.While
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
@@ -71,7 +72,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::expr::Call
// 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::location::UnknownLoc
+// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(import-object): crate::ir::r#type::MissingType
// tvm-ffi-stubgen(import-object): crate::ir::r#type::MissingTypeObj
// tvm-ffi-stubgen(import-object): crate::ir::r#type::PrimType
@@ -104,7 +105,7 @@ use crate::ir::expr::PrimVar;
use crate::ir::expr::StringImm;
use crate::ir::expr::Var;
use crate::ir::location::Location;
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
use crate::ir::op::Op;
use crate::ir::r#type::MissingType;
use crate::ir::r#type::MissingTypeObj;
@@ -535,9 +536,7 @@ impl For {
&message,
));
}
- return Ok(
- IntImm::from_dtype(var_ty.dtype, imm.value.clone(),
UnknownLoc::new())?.into()
- );
+ return Ok(IntImm::new(var_ty.clone(), imm.value.clone(),
UNKNOWN_LOC)?.into());
}
if ty.dtype.bits > var_ty.dtype.bits {
let message = format!(
@@ -624,7 +623,7 @@ impl WhileObj {
impl While {
/// Lossless complete-field allocation.
- pub fn new(
+ pub fn from_fields(
loc: impl Into<Location>,
condition: impl Into<PrimExpr>,
body: impl Into<SeqStmt>,
@@ -644,7 +643,7 @@ impl While {
/// # Errors
///
/// Returns an `InternalError` if `condition` is a vector.
- pub fn checked(
+ pub fn new(
condition: impl IntoExpr<PrimExpr>,
body: impl Into<Stmt>,
loc: impl Into<Location>,
@@ -656,7 +655,7 @@ impl While {
let check = "condition.ty().IsScalar()";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR, check,
""));
}
- Ok(Self::new(loc.into(), condition, body))
+ Ok(Self::from_fields(loc, condition, body))
}
}
@@ -922,7 +921,7 @@ fn region_body_params(op: &Op, args: &Array<Expr>, attrs:
&DictAttrs) -> Result<
return Err(support::check_failed(tvm_ffi::VALUE_ERROR,
"IsRegionOp(op)", &message));
};
let signature = Call::from_fields(
- UnknownLoc::new(),
+ UNKNOWN_LOC,
MissingType::new(),
op.clone(),
args.clone(),
diff --git a/rust/src/ir/type.rs b/rust/src/ir/type.rs
index db9b64038..45ddf4cf5 100644
--- a/rust/src/ir/type.rs
+++ b/rust/src/ir/type.rs
@@ -48,11 +48,12 @@
// tvm-ffi-stubgen(custom-new): ir.PtrType
// tvm-ffi-stubgen(custom-new): ir.TupleType
// tvm-ffi-stubgen(custom-new): ir.FuncType
+// tvm-ffi-stubgen(custom-new): ir.PrimType
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
-// tvm-ffi-stubgen(import-object): crate::ir::location::UnknownLoc
+// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::dtype
// tvm-ffi-stubgen(import-object): tvm_ffi::Error
@@ -61,7 +62,7 @@
// tvm-ffi-stubgen(begin): import-section
use crate::ir::dtype;
use crate::ir::location::Location;
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
use crate::ir::support;
use crate::runtime::MemoryScope;
use std::ops::Deref;
@@ -491,7 +492,7 @@ impl PrimTypeObj {
impl PrimType {
/// Lossless complete-field allocation.
- pub fn new(loc: impl Into<Location>, dtype: impl Into<DLDataType>) -> Self
{
+ pub fn from_fields(loc: impl Into<Location>, dtype: impl Into<DLDataType>)
-> Self {
let obj = PrimTypeObj::new(loc.into(), dtype.into());
Self { base: ObjectArc::new(obj) }
}
@@ -716,7 +717,7 @@ impl PtrType {
let message = "PtrType element_type cannot be Type::Missing()";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
condition, message));
}
- Ok(Self::from_fields(UnknownLoc::new(), element_type, scope.clone()))
+ Ok(Self::from_fields(UNKNOWN_LOC, element_type, scope.clone()))
}
}
@@ -853,6 +854,6 @@ impl FuncType {
ret_type: impl Into<Type>,
) -> Self {
let arg_types =
arg_types.into_iter().map(Into::into).collect::<Array<Type>>();
- Self::from_fields(UnknownLoc::new(), arg_types, ret_type.into())
+ Self::from_fields(UNKNOWN_LOC, arg_types, ret_type.into())
}
}
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index 7f19af5df..c2a24c8fe 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -68,6 +68,7 @@
// tvm-ffi-stubgen(upcast): sym.SplitExpr -> PrimExpr
// tvm-ffi-stubgen(upcast): sym.SumExpr -> PrimExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ExitOnDrop;
diff --git a/rust/src/tirx/analysis/analysis.rs
b/rust/src/tirx/analysis/analysis.rs
index 656919db4..31ddb71d5 100644
--- a/rust/src/tirx/analysis/analysis.rs
+++ b/rust/src/tirx/analysis/analysis.rs
@@ -47,6 +47,7 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::Expr;
diff --git a/rust/src/tirx/function.rs b/rust/src/tirx/function.rs
index e63c3ed35..683f5e939 100644
--- a/rust/src/tirx/function.rs
+++ b/rust/src/tirx/function.rs
@@ -52,7 +52,7 @@
// tvm-ffi-stubgen(import-object): crate::ir::r#type::FuncType
// tvm-ffi-stubgen(import-object): crate::ir::r#type::MissingTypeObj
// tvm-ffi-stubgen(import-object): crate::ir::r#type::TupleType
-// tvm-ffi-stubgen(import-object): crate::ir::location::UnknownLoc
+// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(begin): import-section
use crate::ir::attrs::DictAttrs;
@@ -61,7 +61,7 @@ use crate::ir::expr::Var;
use crate::ir::function::BaseFunc;
use crate::ir::function::BaseFuncObj;
use crate::ir::location::Location;
-use crate::ir::location::UnknownLoc;
+use crate::ir::location::UNKNOWN_LOC;
use crate::ir::r#type::FuncType;
use crate::ir::r#type::MissingTypeObj;
use crate::ir::r#type::TupleType;
@@ -164,7 +164,7 @@ impl Function {
let params =
params.into_iter().map(Into::into).collect::<Array<Var>>();
let ret_type: Type = ret_type.into();
let ret_type = match ret_type.downcast_ref::<MissingTypeObj>() {
- Some(_) => TupleType::new(Vec::<Type>::new(),
UnknownLoc::new()).into(),
+ Some(_) => TupleType::new(Vec::<Type>::new(), UNKNOWN_LOC).into(),
None => ret_type,
};
let param_types = params.iter().map(|param|
param.ty.clone()).collect::<Vec<Type>>();
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index 7fe12752f..8c437ef8e 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -331,11 +331,10 @@ pub fn reinterpret_at(
/// Build a rational approximation of `erf(arg)` on `[-4, 4]` in `bits`-bit
floats, as the Metal and
/// WebGPU backends lower `erf` (C++ `tirx::fast_erf_float_expr(arg, bits)`).
pub fn fast_erf_float_expr(arg: impl IntoExpr<PrimExpr>, bits: u8) ->
Result<PrimExpr> {
- let dtype = PrimType::float(bits).dtype;
+ let ty = PrimType::float(bits);
// The C++ coefficients are `float` literals.
- let constant = |value: f32| -> ExprResult {
- FloatImm::from_dtype(dtype, f64::from(value), UNKNOWN_LOC).into()
- };
+ let constant =
+ |value: f32| -> ExprResult { FloatImm::new(&ty, f64::from(value),
UNKNOWN_LOC).into() };
// alpha_1, alpha_3, ..., alpha_13: the odd coefficients of the numerator.
#[allow(clippy::excessive_precision)]
let alpha: [f32; 7] = [
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index 720b8c43e..d5de159c1 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -63,6 +63,8 @@
// A tensor's shape is a list of `PrimExpr`s, as C++ `TensorTypeObj` holds it.
// tvm-ffi-stubgen(field): tirx.TensorType.shape -> Array<PrimExpr>
// tvm-ffi-stubgen(import-object): tvm_ffi::Array
+// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
diff --git a/rust/src/tirx/thread_scope.rs b/rust/src/tirx/thread_scope.rs
index 9c6956d82..7c71de7bb 100644
--- a/rust/src/tirx/thread_scope.rs
+++ b/rust/src/tirx/thread_scope.rs
@@ -39,6 +39,8 @@
// does.
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
// tvm-ffi-stubgen(import-object): tvm_ffi::String
+// tvm-ffi-stubgen(import-object): std::borrow::Borrow
+// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use std::borrow::Borrow;
diff --git a/rust/tests/ir/expr.rs b/rust/tests/ir/expr.rs
index f2c106ac2..d1c3a6acc 100644
--- a/rust/tests/ir/expr.rs
+++ b/rust/tests/ir/expr.rs
@@ -37,17 +37,17 @@ fn expressions() {
assert_eq!(GlobalVar::new("main", UNKNOWN_LOC).name_hint.as_str(), "main");
assert_eq!(StringImm::new("hi", UNKNOWN_LOC).value.as_str(), "hi");
assert_eq!(FloatImm::float32(1.5).value, 1.5);
- assert!(IntImm::from_dtype(dtype::uint(8), 256, UNKNOWN_LOC).is_err());
- assert_eq!(DataTypeImm::from_dtype(dtype::int(8), UNKNOWN_LOC).value,
dtype::int(8));
+ assert!(IntImm::new(PrimType::uint(8), 256, UNKNOWN_LOC).is_err());
+ assert_eq!(DataTypeImm::new(dtype::int(8), UNKNOWN_LOC).value,
dtype::int(8));
let i32_ty = PrimType::int(32);
- let constant = GenericConst::checked(&i32_ty, Any::from(i32_ty.clone()),
UNKNOWN_LOC).unwrap();
- assert!(GenericConst::checked(&i32_ty, Any::from(1i64),
UNKNOWN_LOC).is_err());
+ let constant = GenericConst::new(&i32_ty, Any::from(i32_ty.clone()),
UNKNOWN_LOC).unwrap();
+ assert!(GenericConst::new(&i32_ty, Any::from(1i64), UNKNOWN_LOC).is_err());
assert!(structural_equal(&constant.ty, &i32_ty));
let tuple = Tuple::new([PrimExpr::from(&x), int32(1)], UNKNOWN_LOC);
- let item = TupleGetItem::from_tuple(&tuple, 1, UNKNOWN_LOC).unwrap();
+ let item = TupleGetItem::new(&tuple, 1, UNKNOWN_LOC).unwrap();
assert!(structural_equal(&item.ty, &i32_ty));
- assert!(TupleGetItem::from_tuple(&tuple, 2, UNKNOWN_LOC).is_err());
+ assert!(TupleGetItem::new(&tuple, 2, UNKNOWN_LOC).is_err());
let subscripted =
Expr::from(&tuple).subscript([Subscript::Index(int32(0).into())],
UNKNOWN_LOC).unwrap();
assert!(Expr::try_from(subscripted).is_ok());
@@ -98,7 +98,12 @@ fn constants_agree_with_cpp() {
&[(&parse_dtype(name)).into(), (&value).into(),
(&Location::unknown()).into()],
);
let what = format!("FloatImm({name}, {value})");
- agree(&what, FloatImm::from_dtype(parse_dtype(name), value,
UNKNOWN_LOC), native);
+ agree(
+ &what,
+ PrimType::from_dtype(parse_dtype(name))
+ .and_then(|ty| FloatImm::new(ty, value, UNKNOWN_LOC)),
+ native,
+ );
};
for (name, value) in [
("float32", 1.5),
@@ -128,7 +133,12 @@ fn constants_agree_with_cpp() {
&[(&parse_dtype(name)).into(), (&value).into(),
(&Location::unknown()).into()],
);
let what = format!("IntImm({name}, {value})");
- agree(&what, IntImm::from_dtype(parse_dtype(name), value.clone(),
UNKNOWN_LOC), native);
+ agree(
+ &what,
+ PrimType::from_dtype(parse_dtype(name))
+ .and_then(|ty| IntImm::new(ty, value.clone(), UNKNOWN_LOC)),
+ native,
+ );
};
let big = |text: &str| -> BigInt { text.parse().unwrap() };
for (name, value) in [
@@ -156,7 +166,7 @@ fn constants_agree_with_cpp() {
let generic = |what: &str, ty: &Type, value: Any| {
let native =
native("ir.GenericConst", &[ty.into(), (&value).into(),
(&Location::unknown()).into()]);
- agree(what, GenericConst::checked(ty, value.clone(), UNKNOWN_LOC),
native);
+ agree(what, GenericConst::new(ty, value.clone(), UNKNOWN_LOC), native);
};
let i32_ty_as_type: Type = i32_ty.clone().into();
generic("GenericConst(object)", &i32_ty_as_type,
Any::from(i32_ty.clone()));
@@ -168,7 +178,7 @@ fn constants_agree_with_cpp() {
let int8 = dtype::int(8);
let native_dtype = native("ir.DataTypeImm", &[(&int8).into(),
(&Location::unknown()).into()]);
- agree("DataTypeImm", Ok(DataTypeImm::from_dtype(int8, UNKNOWN_LOC)),
native_dtype);
+ agree("DataTypeImm", Ok(DataTypeImm::new(int8, UNKNOWN_LOC)),
native_dtype);
let text = FfiString::from("hi");
let native_string = native("ir.StringImm", &[(&text).into(),
(&Location::unknown()).into()]);
agree("StringImm", Ok(StringImm::new("hi", UNKNOWN_LOC)), native_string);
@@ -265,7 +275,7 @@ fn expressions_agree_with_cpp() {
"ir.TupleGetItem",
&[(&value).into(), (&index).into(), (&Location::unknown()).into()],
);
- agree(what, TupleGetItem::from_tuple(&value, index, UNKNOWN_LOC),
native_item);
+ agree(what, TupleGetItem::new(&value, index, UNKNOWN_LOC),
native_item);
}
let a = tensor_var("A");
diff --git a/rust/tests/ir/main.rs b/rust/tests/ir/main.rs
index 3a2f75225..ce15f054c 100644
--- a/rust/tests/ir/main.rs
+++ b/rust/tests/ir/main.rs
@@ -49,7 +49,7 @@ fn parse_dtype(name: &str) -> DLDataType {
}
fn int32(value: i64) -> PrimExpr {
- IntImm::from_dtype(dtype::int(32), value, UNKNOWN_LOC).unwrap().into()
+ IntImm::new(PrimType::int(32), value, UNKNOWN_LOC).unwrap().into()
}
fn int32_var(name: &str) -> PrimVar {
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 111ae58d1..4d50c2ef9 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -672,7 +672,7 @@ fn static_init_blocks_define_operators_and_types() {
// The `RegionStmt` constructor checks a region of the Rust operator as it
checks one of
// `tirx.launch_thread`: the same outcome, and the same error but for the
operator's name.
let extent: Expr = int32(8).into();
- let wide: Expr = IntImm::from_dtype(parse_dtype("int64"), 8,
UNKNOWN_LOC).unwrap().into();
+ let wide: Expr = IntImm::new(PrimType::int(64), 8,
UNKNOWN_LOC).unwrap().into();
let flag: Expr = IntImm::bool(true).into();
let i = var("i", "int32");
let cases = vec![
diff --git a/rust/tests/ir/prim/op.rs b/rust/tests/ir/prim/op.rs
index 9d6af33bc..6b546de34 100644
--- a/rust/tests/ir/prim/op.rs
+++ b/rust/tests/ir/prim/op.rs
@@ -50,15 +50,17 @@ fn prim_var(ty: &str) -> PrimExpr {
}
fn imm(ty: &str, value: i64) -> PrimExpr {
- IntImm::from_dtype(parse_dtype(ty), value, UNKNOWN_LOC).unwrap().into()
+ IntImm::new(PrimType::from_dtype(parse_dtype(ty)).unwrap(), value,
UNKNOWN_LOC).unwrap().into()
}
fn big_imm(ty: &str, value: BigInt) -> PrimExpr {
- IntImm::from_dtype(parse_dtype(ty), value, UNKNOWN_LOC).unwrap().into()
+ IntImm::new(PrimType::from_dtype(parse_dtype(ty)).unwrap(), value,
UNKNOWN_LOC).unwrap().into()
}
fn fimm(ty: &str, value: f64) -> PrimExpr {
- FloatImm::from_dtype(parse_dtype(ty), value, UNKNOWN_LOC).unwrap().into()
+ FloatImm::new(PrimType::from_dtype(parse_dtype(ty)).unwrap(), value,
UNKNOWN_LOC)
+ .unwrap()
+ .into()
}
fn repr<T>(value: &T) -> String
diff --git a/rust/tests/ir/stmt.rs b/rust/tests/ir/stmt.rs
index e6a7f025b..942bf0e97 100644
--- a/rust/tests/ir/stmt.rs
+++ b/rust/tests/ir/stmt.rs
@@ -58,7 +58,7 @@ fn statements() {
assert_eq!(for_.kind, kind as i32);
assert_eq!(for_.body.seq.len(), 1);
let less = op("prim._OpLT", &x, int32(4));
- let while_ = While::checked(&less, &body, UNKNOWN_LOC).unwrap();
+ let while_ = While::new(&less, &body, UNKNOWN_LOC).unwrap();
assert_eq!(while_.body.seq.len(), 2);
let if_ = If::new(&less, &evaluate, Some(&Stmt::from(&ret)),
UNKNOWN_LOC).unwrap();
assert!(if_.else_case.is_some());
@@ -164,7 +164,7 @@ fn statements_agree_with_cpp() {
(
"For(uint8 literal)",
&i,
- IntImm::from_dtype(dtype::uint(8), 3, UNKNOWN_LOC).unwrap().into(),
+ IntImm::new(PrimType::uint(8), 3, UNKNOWN_LOC).unwrap().into(),
int32(8),
int32(1),
evaluate.clone(),
@@ -227,7 +227,7 @@ fn statements_agree_with_cpp() {
"ir.While",
&[(&condition).into(), (&evaluate).into(),
(&Location::unknown()).into()],
);
- agree(what, While::checked(&condition, &evaluate, UNKNOWN_LOC),
native_while);
+ agree(what, While::new(&condition, &evaluate, UNKNOWN_LOC),
native_while);
}
let a = tensor_var("A");
diff --git a/rust/tests/sym/main.rs b/rust/tests/sym/main.rs
index b6da0f31d..9025aad71 100644
--- a/rust/tests/sym/main.rs
+++ b/rust/tests/sym/main.rs
@@ -27,9 +27,7 @@ mod analyzer;
mod int_set;
mod iter_affine_map;
-use tvm_tirx::ir::{
- UNKNOWN_LOC, dtype, Expr, IntImm, Location, PrimExpr, PrimType, PrimVar,
Range, Var,
-};
+use tvm_tirx::ir::{UNKNOWN_LOC, Expr, IntImm, Location, PrimExpr, PrimType,
PrimVar, Range, Var};
use tvm_tirx::sym::{Analyzer, ConstraintContext};
use tvm_tirx::tvm_ffi::String as FfiString;
use tvm_tirx::tvm_ffi::object::ObjectRef;
@@ -37,7 +35,7 @@ use tvm_tirx::tvm_ffi::object::ObjectRef;
use common::global;
fn int32(value: i64) -> PrimExpr {
- IntImm::from_dtype(dtype::int(32), value, UNKNOWN_LOC).unwrap().into()
+ IntImm::new(PrimType::int(32), value, UNKNOWN_LOC).unwrap().into()
}
fn op(name: &str, a: impl Into<Expr>, b: impl Into<Expr>) -> PrimExpr {
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index 426e0f22a..8a58df347 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -40,8 +40,8 @@ mod transform;
use std::sync::atomic::{AtomicUsize, Ordering};
use tvm_tirx::ir::{
- Location, UNKNOWN_LOC, dtype, BaseFunc, DictAttrs, Expr, GlobalInfo,
GlobalVar, IRModule,
- IntImm, PrimExpr, PrimType, PrimVar, Range, SeqStmt, Stmt, Type, Var,
+ Location, UNKNOWN_LOC, BaseFunc, DictAttrs, Expr, GlobalInfo, GlobalVar,
IRModule, IntImm,
+ PrimExpr, PrimType, PrimVar, Range, SeqStmt, Stmt, Type, Var,
};
use tvm_tirx::runtime::MemoryScope;
use tvm_tirx::tirx::{Axis, Function, Iter, TileLayout};
@@ -50,7 +50,7 @@ use tvm_tirx::tvm_ffi::{Array, Map, ObjectRefCast, String as
FfiString};
use common::global;
fn int32(value: i64) -> PrimExpr {
- IntImm::from_dtype(dtype::int(32), value, UNKNOWN_LOC).unwrap().into()
+ IntImm::new(PrimType::int(32), value, UNKNOWN_LOC).unwrap().into()
}
fn var(name: &str, ty: impl Into<Type>) -> Var {
diff --git a/rust/tests/tirx/op.rs b/rust/tests/tirx/op.rs
index 8ae0b15e5..8629889be 100644
--- a/rust/tests/tirx/op.rs
+++ b/rust/tests/tirx/op.rs
@@ -76,7 +76,7 @@ fn logaddexp_agrees_with_cpp() {
(prim_var("float32x4").into(), prim_var("float32").into()),
(
prim_var("float32").into(),
- FloatImm::from_dtype(prim("float16").dtype, 0.5,
UNKNOWN_LOC).unwrap().into(),
+ FloatImm::new(prim("float16"), 0.5, UNKNOWN_LOC).unwrap().into(),
),
(prim_var("int32").into(), PrimVar::new("j", PrimType::int(32),
UNKNOWN_LOC).into()),
];
