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 085fc0b8c [Rust] Reorder items so each file reads top-down (#215)
085fc0b8c is described below
commit 085fc0b8cd22386a6192809e64ddf1cb545b9e6d
Author: Linzhang Li <[email protected]>
AuthorDate: Fri Nov 6 17:00:18 2026 -0500
[Rust] Reorder items so each file reads top-down (#215)
Reorder items within each file, as apache/tvm-ffi#946 does for tvm-ffi:
data structures first, each type with its impls, then the public API,
then helpers. Moves only, plus three section headers for the moved
helpers; review with `--color-moved`.
- `ir/prim/op.rs`: in each section the public functions come first, then
the helpers they call, in call order. `ExprResult` is followed by its
inherent impl, its conversions and the `operators` module, then `binary`
and `binary_literal`. `make_const_scalar`, `fold_cast` and its helpers,
the `fold_*` comparisons, `symbolic_infinities`, `fold_conditions`,
`fold_shift`, the math helpers and `unary_intrinsic` move below the
operators that use them.
- `ir/prim/expr.rs`: the operand checks (`Operands`, `binary_type`,
`bool_type`, `logical_type`) move from the top to the end, under their
section; `binary_new!` moves to the top, ahead of its uses.
- `ir/expr.rs`: `impl Expr` follows the generated `Expr`, before
`Subscript`; `check_float_imm` and `check_int_imm` move to the end,
under "What the constant constructors check".
- `ir/stmt.rs`: the statement types come first; `int_imm_value_fits`,
`region_body_params`, `object_address`, `flatten`, `is_no_op` and
`into_body` move to a Helpers section at the end.
- `ir/op_def.rs`: `OpDef`, `AttrKey`, `op_attr`, `CallEffectKind` and
`sig` come first; `complete`, `doc_text`, the hook types and the ABI
trampolines move to a Helpers section at the end.
- `sym/constraint.rs`: as a generated binding, `ConstraintContextObj`
and its layout assertions, then `ConstraintContext`, then their impls.
- `lib.rs`: the modules, `pub use tvm_ffi` and `exports` lead; the link
anchor and the scope guard follow.
- `ir/prim/const_fold.rs`: the arithmetic and comparison folds come
before the helpers they share.
- `ir/support.rs`: `not_a_prim_expr` moves to the end of "Primitive
types", below `prim_type` and `prim_dtype`, which use it.
- `tirx/tensor.rs`: `impl TensorType` follows the generated
`TensorType`, then `TensorVar` and its `sig::ArgType` impl, then
`tensor_is_scalar`; the private `alloc_alignment` moves to the end.
- `ir/dtype.rs`, `tirx/op.rs`: the private helpers (`scalar`;
`is_scalar_int64` and `storage_bytes`) move below the public functions
that use them.
---
rust/src/ir/dtype.rs | 10 +-
rust/src/ir/expr.rs | 361 +++++++++--------
rust/src/ir/op_def.rs | 179 ++++----
rust/src/ir/prim/const_fold.rs | 260 ++++++------
rust/src/ir/prim/expr.rs | 138 +++----
rust/src/ir/prim/op.rs | 898 ++++++++++++++++++++---------------------
rust/src/ir/stmt.rs | 185 ++++-----
rust/src/ir/support.rs | 18 +-
rust/src/lib.rs | 42 +-
rust/src/sym/constraint.rs | 98 ++---
rust/src/tirx/op.rs | 24 +-
rust/src/tirx/tensor.rs | 58 +--
12 files changed, 1140 insertions(+), 1131 deletions(-)
diff --git a/rust/src/ir/dtype.rs b/rust/src/ir/dtype.rs
index 2d8e28ecc..21b04361e 100644
--- a/rust/src/ir/dtype.rs
+++ b/rust/src/ir/dtype.rs
@@ -42,11 +42,6 @@
use tvm_ffi::DLDataType;
use tvm_ffi::DLDataTypeCode;
-/// The scalar type of `code` and `bits`.
-const fn scalar(code: DLDataTypeCode, bits: u8) -> DLDataType {
- DLDataType { code: code as u8, bits, lanes: 1 }
-}
-
/// Return the signed integer type of `bits` bits, such as `int32` for 32 (C++
/// `DataType::Int(bits)`).
pub const fn int(bits: u8) -> DLDataType {
@@ -85,3 +80,8 @@ pub const fn void() -> DLDataType {
pub const fn with_lanes(dtype: DLDataType, lanes: u16) -> DLDataType {
DLDataType { code: dtype.code, bits: dtype.bits, lanes }
}
+
+/// The scalar type of `code` and `bits`.
+const fn scalar(code: DLDataTypeCode, bits: u8) -> DLDataType {
+ DLDataType { code: code as u8, bits, lanes: 1 }
+}
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index 4ea215a3b..4063c7a96 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -202,6 +202,28 @@ unsafe impl UnsafeInit for ExprObj {
unsafe impl tvm_ffi::object::ImmutableObject for ExprObj {}
// tvm-ffi-stubgen(end)
+impl Expr {
+ /// Return `self[indices]`, as the subscript rule of its type realizes it:
a [`TupleGetItem`] of
+ /// a tuple, a [`TensorLoad`] or [`TensorRegion`] of a tensor
(`ir.SubscriptExprRealize`).
+ ///
+ /// # Errors
+ ///
+ /// Returns a `TypeError` if the type of the expression has no subscript
rule, or the error the
+ /// rule raises.
+ pub fn subscript(
+ &self,
+ indices: impl IntoIterator<Item = Subscript>,
+ loc: impl Into<Location>,
+ ) -> Result<Any> {
+ let indices = indices.into_iter().map(|index|
Any::from(&index)).collect::<Array<Any>>();
+ tvm_ffi::cached_global_func!("ir.SubscriptExprRealize").call_tuple((
+ self,
+ &indices,
+ &loc.into(),
+ ))
+ }
+}
+
/// An index of [`Expr::subscript`]: a primitive expression, or a
`start:stop:step` slice (C++
/// `SubscriptSlice`).
#[derive(Clone)]
@@ -226,28 +248,6 @@ impl From<&Subscript> for Any {
}
}
-impl Expr {
- /// Return `self[indices]`, as the subscript rule of its type realizes it:
a [`TupleGetItem`] of
- /// a tuple, a [`TensorLoad`] or [`TensorRegion`] of a tensor
(`ir.SubscriptExprRealize`).
- ///
- /// # Errors
- ///
- /// Returns a `TypeError` if the type of the expression has no subscript
rule, or the error the
- /// rule raises.
- pub fn subscript(
- &self,
- indices: impl IntoIterator<Item = Subscript>,
- loc: impl Into<Location>,
- ) -> Result<Any> {
- let indices = indices.into_iter().map(|index|
Any::from(&index)).collect::<Array<Any>>();
- tvm_ffi::cached_global_func!("ir.SubscriptExprRealize").call_tuple((
- self,
- &indices,
- &loc.into(),
- ))
- }
-}
-
//
---------------------------------------------------------------------------------------------
// Conversions into an expression
@@ -1260,91 +1260,6 @@ impl From<f64> for FloatImm {
}
}
-/// The checks of C++ `FloatImm(ty, value)`: a scalar floating-point type
whose range holds a
-/// finite `value`.
-fn check_float_imm(dtype: DLDataType, value: f64) -> Result<()> {
- use support::{check_failed, compare_failed, CppDouble};
- use tvm_ffi::VALUE_ERROR;
- const FLOAT4_E2M1FN: u8 = tvm_ffi::DLDataTypeCode::kDLFloat4_e2m1fn as u8;
- /// `ffi::DLExtDataTypeCode::kDLExtCustomBegin`.
- const CUSTOM_BEGIN: u8 = 129;
- if support::is_scalable_vector(dtype) ||
support::is_fixed_length_vector(dtype) {
- let condition = "!value_ty.IsScalableVector() &&
!value_ty.IsFixedLengthVector()";
- return Err(check_failed(VALUE_ERROR, condition, "FloatImm can only
take scalar."));
- }
- let (code, bits) = (dtype.code, dtype.bits);
- let float8 = support::float8_bound(code);
- let float6 = support::float6_bound(code);
- let supported = code == support::FLOAT
- || float8.is_some()
- || float6.is_some()
- || (code == support::BFLOAT && bits == 16)
- || (code == FLOAT4_E2M1FN && bits == 4)
- || code >= CUSTOM_BEGIN;
- let name = || support::dtype_name(dtype);
- if !supported {
- let condition = "value_ty.MatchesCode(DLDataTypeCode::kDLFloat, \
- DLDataTypeCode::kDLFloat8_e3m4,
DLDataTypeCode::kDLFloat8_e4m3, \
- DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
DLDataTypeCode::kDLFloat8_e4m3fn, \
- DLDataTypeCode::kDLFloat8_e4m3fnuz,
DLDataTypeCode::kDLFloat8_e5m2, \
- DLDataTypeCode::kDLFloat8_e5m2fnuz,
DLDataTypeCode::kDLFloat8_e8m0fnu, \
- DLDataTypeCode::kDLFloat6_e2m3fn,
DLDataTypeCode::kDLFloat6_e3m2fn) || \
-
value_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16) || \
-
value_ty.MatchesElementType(DLDataTypeCode::kDLFloat4_e2m1fn, 4) || \
- static_cast<int>(code) >= \
-
static_cast<int>(ffi::DLExtDataTypeCode::kDLExtCustomBegin)";
- let message = format!("FloatImm supports only float, but {} was
supplied.", name());
- return Err(check_failed(VALUE_ERROR, condition, &message));
- }
- if value.is_infinite() || value.is_nan() {
- return Ok(());
- }
- // The bound checked from below and how its message says so, and the bound
from above.
- let (low, high): ((&str, f64, &str), (&str, f64)) = if bits == 32 {
- (
- (
- "value >= std::numeric_limits<float>::lowest()",
- f32::MIN as f64,
- "exceeds minimum of",
- ),
- ("value <= std::numeric_limits<float>::max()", f32::MAX as f64),
- )
- } else if code == support::FLOAT && bits == 16 {
- (
- ("value >= -support::kMaxFloat16", -support::MAX_FLOAT16, "exceeds
minimum of"),
- ("value <= support::kMaxFloat16", support::MAX_FLOAT16),
- )
- } else if code == support::BFLOAT && bits == 16 {
- (
- ("value >= -support::kMaxBFloat16", -support::MAX_BFLOAT16,
"exceeds minimum of"),
- ("value <= support::kMaxBFloat16", support::MAX_BFLOAT16),
- )
- } else if let Some((bound, nonneg)) = float8 {
- let low = match nonneg {
- true => ("value >= 0", 0.0, "below zero for unsigned"),
- false => ("value >= -bound", -bound, "below minimum of"),
- };
- (low, ("value <= bound", bound))
- } else if let Some(bound) = float6 {
- (("value >= -bound", -bound, "below minimum of"), ("value <= bound",
bound))
- } else if code == FLOAT4_E2M1FN {
- (("value >= -bound", -6.0, "below minimum of"), ("value <= bound",
6.0))
- } else {
- return Ok(());
- };
- let failed = |condition: &str, bound: f64, what: &str| {
- let message = format!("Literal value {} {what} {}", CppDouble(value),
name());
- compare_failed(VALUE_ERROR, condition, CppDouble(value),
CppDouble(bound), &message)
- };
- if value < low.1 {
- return Err(failed(low.0, low.1, low.2));
- }
- if value > high.1 {
- return Err(failed(high.0, high.1, "exceeds maximum of"));
- }
- Ok(())
-}
-
// tvm-ffi-stubgen(begin): object/ir.IntImm
/// Complete: reflected fields fill [40, 56) exactly.
#[repr(C)]
@@ -1467,78 +1382,6 @@ impl From<bool> for IntImm {
}
}
-/// The checks of C++ `IntImm(ty, value)`: a scalar integer or boolean type
whose range holds
-/// `value`. The messages are built only when a check fails.
-fn check_int_imm(dtype: DLDataType, value: &BigInt) -> Result<()> {
- use support::{check_failed, compare_failed};
- use tvm_ffi::VALUE_ERROR;
- let name = || support::dtype_name(dtype);
- if support::is_scalable_vector(dtype) ||
support::is_fixed_length_vector(dtype) {
- let condition = "!value_ty.IsScalableVector() &&
!value_ty.IsFixedLengthVector()";
- let message = format!("IntImm can only take scalar, but {} was
supplied.", name());
- return Err(check_failed(VALUE_ERROR, condition, &message));
- }
- if !matches!(dtype.code, support::INT | support::UINT | support::BOOL) {
- let condition = "value_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt, \
- DLDataTypeCode::kDLBool)";
- let message =
- format!("IntImm supports only int or uint or bool type, but {} was
supplied.", name());
- return Err(check_failed(VALUE_ERROR, condition, &message));
- }
- let bits = dtype.bits as u32;
- if bits == 0 {
- let message = "IntImm requires a positive integer width";
- return Err(compare_failed(VALUE_ERROR, "bits > 0", bits, 0, message));
- }
- let zero = BigInt::from(0i64);
- if dtype.code == support::UINT {
- if *value < zero {
- let message =
- format!("Literal value {value} is negative for unsigned
integer type {}", name());
- return Err(compare_failed(VALUE_ERROR, "value >= 0", value, 0,
&message));
- }
- let message = || format!("Literal value {value} exceeds maximum of
{}", name());
- if bits <= 64 {
- if !value.to_u64().is_some_and(|small| bits == 64 || small < 1u64
<< bits) {
- let condition =
- "small.has_value() && (bits == 64 || *small < (uint64_t{1}
<< bits))";
- return Err(check_failed(VALUE_ERROR, condition, &message()));
- }
- } else {
- let limit = BigInt::from(1i64).try_shl(bits as i64)?;
- if *value >= limit {
- let condition = "value < ffi::BigInt(1) << bits";
- return Err(compare_failed(VALUE_ERROR, condition, value,
&limit, &message()));
- }
- }
- } else if bits == 1 || dtype.code == support::BOOL {
- if *value != zero && *value != BigInt::from(1i64) {
- let message = format!("{value} exceeds range of {}", name());
- return Err(check_failed(VALUE_ERROR, "value == 0 || value == 1",
&message));
- }
- } else {
- let message = || format!("Literal value {value} exceeds range of {}",
name());
- if bits <= 64 {
- let half = 1i128 << (bits - 1);
- let fits = value.to_i64().is_some_and(|small| {
- bits == 64 || ((small as i128) >= -half && (small as i128) <
half)
- });
- if !fits {
- let condition = "small.has_value() && (bits == 64 || (*small
>= -(int64_t{1} << \
- (bits - 1)) && *small < (int64_t{1} << (bits
- 1))))";
- return Err(check_failed(VALUE_ERROR, condition, &message()));
- }
- } else {
- let limit = BigInt::from(1i64).try_shl(bits as i64 - 1)?;
- if *value < -limit.clone() || *value >= limit {
- let condition = "value >= -limit && value < limit";
- return Err(check_failed(VALUE_ERROR, condition, &message()));
- }
- }
- }
- Ok(())
-}
-
// tvm-ffi-stubgen(begin): object/ir.GenericConst
/// Complete: reflected fields fill [40, 56) exactly.
#[repr(C)]
@@ -2314,3 +2157,163 @@
crate::def_expr_op_over_loads!(crate::sym::iter_affine_map::IterSplitExpr);
crate::def_expr_op_over_loads!(crate::sym::iter_affine_map::IterSumExpr);
crate::def_expr_op_over_loads!(crate::sym::analyzer::SplitExpr);
crate::def_expr_op_over_loads!(crate::sym::analyzer::SumExpr);
+
+//
---------------------------------------------------------------------------------------------
+// What the constant constructors check
+
+/// The checks of C++ `FloatImm(ty, value)`: a scalar floating-point type
whose range holds a
+/// finite `value`.
+fn check_float_imm(dtype: DLDataType, value: f64) -> Result<()> {
+ use support::{check_failed, compare_failed, CppDouble};
+ use tvm_ffi::VALUE_ERROR;
+ const FLOAT4_E2M1FN: u8 = tvm_ffi::DLDataTypeCode::kDLFloat4_e2m1fn as u8;
+ /// `ffi::DLExtDataTypeCode::kDLExtCustomBegin`.
+ const CUSTOM_BEGIN: u8 = 129;
+ if support::is_scalable_vector(dtype) ||
support::is_fixed_length_vector(dtype) {
+ let condition = "!value_ty.IsScalableVector() &&
!value_ty.IsFixedLengthVector()";
+ return Err(check_failed(VALUE_ERROR, condition, "FloatImm can only
take scalar."));
+ }
+ let (code, bits) = (dtype.code, dtype.bits);
+ let float8 = support::float8_bound(code);
+ let float6 = support::float6_bound(code);
+ let supported = code == support::FLOAT
+ || float8.is_some()
+ || float6.is_some()
+ || (code == support::BFLOAT && bits == 16)
+ || (code == FLOAT4_E2M1FN && bits == 4)
+ || code >= CUSTOM_BEGIN;
+ let name = || support::dtype_name(dtype);
+ if !supported {
+ let condition = "value_ty.MatchesCode(DLDataTypeCode::kDLFloat, \
+ DLDataTypeCode::kDLFloat8_e3m4,
DLDataTypeCode::kDLFloat8_e4m3, \
+ DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
DLDataTypeCode::kDLFloat8_e4m3fn, \
+ DLDataTypeCode::kDLFloat8_e4m3fnuz,
DLDataTypeCode::kDLFloat8_e5m2, \
+ DLDataTypeCode::kDLFloat8_e5m2fnuz,
DLDataTypeCode::kDLFloat8_e8m0fnu, \
+ DLDataTypeCode::kDLFloat6_e2m3fn,
DLDataTypeCode::kDLFloat6_e3m2fn) || \
+
value_ty.MatchesElementType(DLDataTypeCode::kDLBfloat, 16) || \
+
value_ty.MatchesElementType(DLDataTypeCode::kDLFloat4_e2m1fn, 4) || \
+ static_cast<int>(code) >= \
+
static_cast<int>(ffi::DLExtDataTypeCode::kDLExtCustomBegin)";
+ let message = format!("FloatImm supports only float, but {} was
supplied.", name());
+ return Err(check_failed(VALUE_ERROR, condition, &message));
+ }
+ if value.is_infinite() || value.is_nan() {
+ return Ok(());
+ }
+ // The bound checked from below and how its message says so, and the bound
from above.
+ let (low, high): ((&str, f64, &str), (&str, f64)) = if bits == 32 {
+ (
+ (
+ "value >= std::numeric_limits<float>::lowest()",
+ f32::MIN as f64,
+ "exceeds minimum of",
+ ),
+ ("value <= std::numeric_limits<float>::max()", f32::MAX as f64),
+ )
+ } else if code == support::FLOAT && bits == 16 {
+ (
+ ("value >= -support::kMaxFloat16", -support::MAX_FLOAT16, "exceeds
minimum of"),
+ ("value <= support::kMaxFloat16", support::MAX_FLOAT16),
+ )
+ } else if code == support::BFLOAT && bits == 16 {
+ (
+ ("value >= -support::kMaxBFloat16", -support::MAX_BFLOAT16,
"exceeds minimum of"),
+ ("value <= support::kMaxBFloat16", support::MAX_BFLOAT16),
+ )
+ } else if let Some((bound, nonneg)) = float8 {
+ let low = match nonneg {
+ true => ("value >= 0", 0.0, "below zero for unsigned"),
+ false => ("value >= -bound", -bound, "below minimum of"),
+ };
+ (low, ("value <= bound", bound))
+ } else if let Some(bound) = float6 {
+ (("value >= -bound", -bound, "below minimum of"), ("value <= bound",
bound))
+ } else if code == FLOAT4_E2M1FN {
+ (("value >= -bound", -6.0, "below minimum of"), ("value <= bound",
6.0))
+ } else {
+ return Ok(());
+ };
+ let failed = |condition: &str, bound: f64, what: &str| {
+ let message = format!("Literal value {} {what} {}", CppDouble(value),
name());
+ compare_failed(VALUE_ERROR, condition, CppDouble(value),
CppDouble(bound), &message)
+ };
+ if value < low.1 {
+ return Err(failed(low.0, low.1, low.2));
+ }
+ if value > high.1 {
+ return Err(failed(high.0, high.1, "exceeds maximum of"));
+ }
+ Ok(())
+}
+
+/// The checks of C++ `IntImm(ty, value)`: a scalar integer or boolean type
whose range holds
+/// `value`. The messages are built only when a check fails.
+fn check_int_imm(dtype: DLDataType, value: &BigInt) -> Result<()> {
+ use support::{check_failed, compare_failed};
+ use tvm_ffi::VALUE_ERROR;
+ let name = || support::dtype_name(dtype);
+ if support::is_scalable_vector(dtype) ||
support::is_fixed_length_vector(dtype) {
+ let condition = "!value_ty.IsScalableVector() &&
!value_ty.IsFixedLengthVector()";
+ let message = format!("IntImm can only take scalar, but {} was
supplied.", name());
+ return Err(check_failed(VALUE_ERROR, condition, &message));
+ }
+ if !matches!(dtype.code, support::INT | support::UINT | support::BOOL) {
+ let condition = "value_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt, \
+ DLDataTypeCode::kDLBool)";
+ let message =
+ format!("IntImm supports only int or uint or bool type, but {} was
supplied.", name());
+ return Err(check_failed(VALUE_ERROR, condition, &message));
+ }
+ let bits = dtype.bits as u32;
+ if bits == 0 {
+ let message = "IntImm requires a positive integer width";
+ return Err(compare_failed(VALUE_ERROR, "bits > 0", bits, 0, message));
+ }
+ let zero = BigInt::from(0i64);
+ if dtype.code == support::UINT {
+ if *value < zero {
+ let message =
+ format!("Literal value {value} is negative for unsigned
integer type {}", name());
+ return Err(compare_failed(VALUE_ERROR, "value >= 0", value, 0,
&message));
+ }
+ let message = || format!("Literal value {value} exceeds maximum of
{}", name());
+ if bits <= 64 {
+ if !value.to_u64().is_some_and(|small| bits == 64 || small < 1u64
<< bits) {
+ let condition =
+ "small.has_value() && (bits == 64 || *small < (uint64_t{1}
<< bits))";
+ return Err(check_failed(VALUE_ERROR, condition, &message()));
+ }
+ } else {
+ let limit = BigInt::from(1i64).try_shl(bits as i64)?;
+ if *value >= limit {
+ let condition = "value < ffi::BigInt(1) << bits";
+ return Err(compare_failed(VALUE_ERROR, condition, value,
&limit, &message()));
+ }
+ }
+ } else if bits == 1 || dtype.code == support::BOOL {
+ if *value != zero && *value != BigInt::from(1i64) {
+ let message = format!("{value} exceeds range of {}", name());
+ return Err(check_failed(VALUE_ERROR, "value == 0 || value == 1",
&message));
+ }
+ } else {
+ let message = || format!("Literal value {value} exceeds range of {}",
name());
+ if bits <= 64 {
+ let half = 1i128 << (bits - 1);
+ let fits = value.to_i64().is_some_and(|small| {
+ bits == 64 || ((small as i128) >= -half && (small as i128) <
half)
+ });
+ if !fits {
+ let condition = "small.has_value() && (bits == 64 || (*small
>= -(int64_t{1} << \
+ (bits - 1)) && *small < (int64_t{1} << (bits
- 1))))";
+ return Err(check_failed(VALUE_ERROR, condition, &message()));
+ }
+ } else {
+ let limit = BigInt::from(1i64).try_shl(bits as i64 - 1)?;
+ if *value < -limit.clone() || *value >= limit {
+ let condition = "value >= -limit && value < limit";
+ return Err(check_failed(VALUE_ERROR, condition, &message()));
+ }
+ }
+ }
+ Ok(())
+}
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 6fb4c8366..355b3a647 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -267,17 +267,6 @@ impl Drop for OpDef {
}
}
-/// Request the APIs of a defined operator in registered destinations through
`ir.expose_op`,
-/// after its signature and attributes are ready. An absent receiver is a
no-op; an installed
-/// receiver's error is returned.
-fn complete(state: Result<Op>) -> Result<Op> {
- let op = state?;
- if let Ok(expose) = tvm_ffi::Function::get_global("ir.expose_op") {
- expose.call_tuple((op.name.clone(),))?;
- }
- Ok(op)
-}
-
/// The key of an operator attribute, C++ `ffi::reflection::AttrKey<T>`: its
name and the type `T`
/// of its values.
///
@@ -356,83 +345,6 @@ impl From<CallEffectKind> for Any {
}
}
-/// The text of `doc` without the indentation its lines share and without the
blank lines around
-/// it. Text that is not indented is unchanged.
-fn doc_text(doc: &str) -> std::string::String {
- let indent = doc
- .lines()
- .filter(|line| !line.trim().is_empty())
- .map(|line| line.len() - line.trim_start().len())
- .min()
- .unwrap_or(0);
- let lines: Vec<&str> =
- doc.lines().map(|line|
line.get(indent..).unwrap_or("").trim_end()).collect();
- lines.join("\n").trim().to_string()
-}
-
-/// The C ABI of a native operator hook, `NativeFunctionView<R(const
CallNode*)>`: it returns
-/// the result, or the error, in an owned `TVMFFIAny`.
-type Hook = extern "C" fn(*const CallObj) -> TVMFFIAny;
-
-/// The C ABI of a native hook on a region statement,
`NativeFunctionView<R(const
-/// RegionStmtNode*)>`.
-type RegionHook = extern "C" fn(*const RegionStmtObj) -> TVMFFIAny;
-
-extern "C" fn validate_abi<F: Fn(&CallObj) -> Result<()>>(call: *const
CallObj) -> TVMFFIAny {
- call_hook(call, zero_sized::<F>())
-}
-
-extern "C" fn infer_type_abi<F: Fn(&CallObj) -> Result<Type>>(call: *const
CallObj) -> TVMFFIAny {
- call_hook(call, zero_sized::<F>())
-}
-
-extern "C" fn validate_signature<L: Checks>(call: *const CallObj) -> TVMFFIAny
{
- call_hook(call, sig::validate::<L>)
-}
-
-extern "C" fn region_body_params_abi<F: Fn(&CallObj) -> Result<Array<Var>>>(
- call: *const CallObj,
-) -> TVMFFIAny {
- call_hook(call, zero_sized::<F>())
-}
-
-extern "C" fn region_validate_abi<F: Fn(&RegionStmtObj) -> Result<()>>(
- region: *const RegionStmtObj,
-) -> TVMFFIAny {
- call_hook(region, zero_sized::<F>())
-}
-
-/// The only value of `F`, a function or a closure that captures nothing.
-fn zero_sized<F>() -> F {
- const {
- assert!(
- std::mem::size_of::<F>() == 0,
- "a native operator hook is a function or a closure that captures
nothing"
- )
- };
- // SAFETY: a zero-sized type has no bytes to initialize.
- unsafe { std::mem::zeroed() }
-}
-
-/// Run `hook` on the node C++ passes (the Call or region being typed or
checked) and return its
-/// result as the hook ABI does. A panic, which must not unwind into C++,
becomes an error.
-fn call_hook<N, R: Into<Any>>(node: *const N, hook: impl Fn(&N) -> Result<R>)
-> TVMFFIAny {
- let result = catch_unwind(AssertUnwindSafe(|| {
- // SAFETY: C++ passes the node being typed or checked, alive for the
call.
- match unsafe { node.as_ref() } {
- Some(node) => hook(node).map(Into::into),
- None => Err(Error::new(TYPE_ERROR, "Invalid Op Call", "")),
- }
- }));
- let value = match result {
- Ok(Ok(value)) => value,
- Ok(Err(error)) => Any::from(error),
- Err(payload) => Any::from(tvm_ffi::Error::from_panic(payload)),
- };
- // The C++ caller takes ownership of the value.
- Any::into_raw_ffi_any(value)
-}
-
/// Signature descriptors for [`OpDef::signature`], C++ `sig::arg<T>` and
friends.
///
/// [`arg`](sig::arg) and [`ty_arg`](sig::ty_arg) are fixed value and type
arguments,
@@ -874,3 +786,94 @@ pub mod sig {
Err(Error::new(TYPE_ERROR, &message, ""))
}
}
+
+//
---------------------------------------------------------------------------------------------
+// Helpers
+
+/// Request the APIs of a defined operator in registered destinations through
`ir.expose_op`,
+/// after its signature and attributes are ready. An absent receiver is a
no-op; an installed
+/// receiver's error is returned.
+fn complete(state: Result<Op>) -> Result<Op> {
+ let op = state?;
+ if let Ok(expose) = tvm_ffi::Function::get_global("ir.expose_op") {
+ expose.call_tuple((op.name.clone(),))?;
+ }
+ Ok(op)
+}
+
+/// The text of `doc` without the indentation its lines share and without the
blank lines around
+/// it. Text that is not indented is unchanged.
+fn doc_text(doc: &str) -> std::string::String {
+ let indent = doc
+ .lines()
+ .filter(|line| !line.trim().is_empty())
+ .map(|line| line.len() - line.trim_start().len())
+ .min()
+ .unwrap_or(0);
+ let lines: Vec<&str> =
+ doc.lines().map(|line|
line.get(indent..).unwrap_or("").trim_end()).collect();
+ lines.join("\n").trim().to_string()
+}
+
+/// The C ABI of a native operator hook, `NativeFunctionView<R(const
CallNode*)>`: it returns
+/// the result, or the error, in an owned `TVMFFIAny`.
+type Hook = extern "C" fn(*const CallObj) -> TVMFFIAny;
+
+/// The C ABI of a native hook on a region statement,
`NativeFunctionView<R(const
+/// RegionStmtNode*)>`.
+type RegionHook = extern "C" fn(*const RegionStmtObj) -> TVMFFIAny;
+
+extern "C" fn validate_abi<F: Fn(&CallObj) -> Result<()>>(call: *const
CallObj) -> TVMFFIAny {
+ call_hook(call, zero_sized::<F>())
+}
+
+extern "C" fn infer_type_abi<F: Fn(&CallObj) -> Result<Type>>(call: *const
CallObj) -> TVMFFIAny {
+ call_hook(call, zero_sized::<F>())
+}
+
+extern "C" fn validate_signature<L: Checks>(call: *const CallObj) -> TVMFFIAny
{
+ call_hook(call, sig::validate::<L>)
+}
+
+extern "C" fn region_body_params_abi<F: Fn(&CallObj) -> Result<Array<Var>>>(
+ call: *const CallObj,
+) -> TVMFFIAny {
+ call_hook(call, zero_sized::<F>())
+}
+
+extern "C" fn region_validate_abi<F: Fn(&RegionStmtObj) -> Result<()>>(
+ region: *const RegionStmtObj,
+) -> TVMFFIAny {
+ call_hook(region, zero_sized::<F>())
+}
+
+/// The only value of `F`, a function or a closure that captures nothing.
+fn zero_sized<F>() -> F {
+ const {
+ assert!(
+ std::mem::size_of::<F>() == 0,
+ "a native operator hook is a function or a closure that captures
nothing"
+ )
+ };
+ // SAFETY: a zero-sized type has no bytes to initialize.
+ unsafe { std::mem::zeroed() }
+}
+
+/// Run `hook` on the node C++ passes (the Call or region being typed or
checked) and return its
+/// result as the hook ABI does. A panic, which must not unwind into C++,
becomes an error.
+fn call_hook<N, R: Into<Any>>(node: *const N, hook: impl Fn(&N) -> Result<R>)
-> TVMFFIAny {
+ let result = catch_unwind(AssertUnwindSafe(|| {
+ // SAFETY: C++ passes the node being typed or checked, alive for the
call.
+ match unsafe { node.as_ref() } {
+ Some(node) => hook(node).map(Into::into),
+ None => Err(Error::new(TYPE_ERROR, "Invalid Op Call", "")),
+ }
+ }));
+ let value = match result {
+ Ok(Ok(value)) => value,
+ Ok(Err(error)) => Any::from(error),
+ Err(payload) => Any::from(tvm_ffi::Error::from_panic(payload)),
+ };
+ // The C++ caller takes ownership of the value.
+ Any::into_raw_ffi_any(value)
+}
diff --git a/rust/src/ir/prim/const_fold.rs b/rust/src/ir/prim/const_fold.rs
index 44fdb18c1..f2ce2234b 100644
--- a/rust/src/ir/prim/const_fold.rs
+++ b/rust/src/ir/prim/const_fold.rs
@@ -37,136 +37,6 @@ use crate::ir::support;
/// The folded expression, or `None` when there is nothing to fold.
pub(super) type Folded = Result<Option<PrimExpr>>;
-//
---------------------------------------------------------------------------------------------
-// Helpers
-
-/// C++ `IsIndexType`: a scalar `int32` or `int64`, which index arithmetic
folds eagerly.
-pub(super) fn is_index_type(dtype: DLDataType) -> bool {
- dtype.code == support::INT && (dtype.bits == 32 || dtype.bits == 64) &&
dtype.lanes == 1
-}
-
-/// C++ `IsRepresentable(value, dtype)`: whether the exact integer `value` is
a literal of the
-/// scalar type `dtype`, an integer in its range or a boolean 0 or 1 (as is an
`int1`).
-pub(super) fn is_representable(value: &BigInt, dtype: DLDataType) -> bool {
- if !support::is_scalar(dtype) || dtype.bits == 0 {
- return false;
- }
- let bits = u32::from(dtype.bits);
- if let Some(small) = value.to_i64() {
- return match dtype.code {
- support::UINT => small >= 0 && (bits >= 63 || (small as u64) <
(1u64 << bits)),
- support::BOOL => small == 0 || small == 1,
- support::INT if bits == 1 => small == 0 || small == 1,
- support::INT => {
- bits >= 64 || (-(1i64 << (bits - 1))..1i64 << (bits -
1)).contains(&small)
- }
- _ => false,
- };
- }
- match dtype.code {
- support::UINT => !value.is_negative() && *value < BigInt::from(1i64)
<< i64::from(bits),
- support::INT if bits > 1 => {
- let limit = BigInt::from(1i64) << i64::from(bits - 1);
- *value >= -limit.clone() && *value < limit
- }
- _ => false,
- }
-}
-
-/// C++ `IsInFloat32Range(value)`: a value that rounds to a finite float32, or
one that is not
-/// finite.
-pub(super) fn is_in_float32_range(value: f64) -> bool {
- !value.is_finite() ||
(f64::from(f32::MIN)..=f64::from(f32::MAX)).contains(&value)
-}
-
-/// The value of an integer constant.
-pub(super) fn int(expr: &PrimExpr) -> Option<&BigInt> {
- expr.downcast_ref::<IntImmObj>().map(|imm| &imm.value)
-}
-
-/// The value of a float constant.
-fn float(expr: &PrimExpr) -> Option<f64> {
- expr.downcast_ref::<FloatImmObj>().map(|imm| imm.value)
-}
-
-/// C++ `IntImm(dtype, value)`.
-fn int_imm(dtype: DLDataType, value: BigInt) -> Folded {
- 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
-/// the type.
-fn exact(dtype: DLDataType, value: BigInt) -> Folded {
- match is_representable(&value, dtype) {
- true => int_imm(dtype, value),
- false => Ok(None),
- }
-}
-
-/// C++ `FloatImm(dtype, value)`.
-fn float_imm(dtype: DLDataType, value: f64) -> Folded {
- Ok(Some(FloatImm::new(PrimType::from_dtype(dtype)?, value,
UNKNOWN_LOC)?.into()))
-}
-
-/// C++ `IntImm::Bool(value)`.
-fn boolean(value: bool) -> Folded {
- Ok(Some(IntImm::bool(value).into()))
-}
-
-/// A float32 operation on the operands rounded to float32, as C++ computes
`float` arithmetic, or
-/// `None` when the result is beyond the float32 range: C++ checks the result
in double precision
-/// before rounding can turn a finite overflow into infinity, and rounding
that result to float32
-/// gives the single-precision one (a double holds more than twice the digits
of a float).
-fn single(x: f64, y: f64, op: fn(f64, f64) -> f64) -> Option<f64> {
- let result = op(f64::from(x as f32), f64::from(y as f32));
- is_in_float32_range(result).then(|| f64::from(result as f32))
-}
-
-/// C++ `FloatImm(dtype, value)` of a folded value, or nothing to fold.
-fn float_imm_of(dtype: DLDataType, value: Option<f64>) -> Folded {
- match value {
- Some(value) => float_imm(dtype, value),
- None => Ok(None),
- }
-}
-
-/// The error of `TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero"`.
-fn check_int_divisor(divisor: &BigInt) -> Result<()> {
- match divisor.is_zero() {
- true => Err(support::compare_failed(
- INTERNAL_ERROR,
- "pb->value != 0",
- divisor,
- 0,
- "Divide by zero",
- )),
- false => Ok(()),
- }
-}
-
-/// The error of `TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero"`.
-fn check_float_divisor(divisor: f64) -> Result<()> {
- match divisor == 0.0 {
- true => Err(support::compare_failed(
- INTERNAL_ERROR,
- "fb->value != 0",
- support::CppDouble(divisor),
- 0,
- "Divide by zero",
- )),
- false => Ok(()),
- }
-}
-
-fn is_one(value: &BigInt) -> bool {
- value.to_i64() == Some(1)
-}
-
-/// The operand itself.
-fn keep(expr: &PrimExpr) -> Folded {
- Ok(Some(expr.clone()))
-}
-
//
---------------------------------------------------------------------------------------------
// Arithmetic
@@ -484,3 +354,133 @@ pub(super) fn not(a: &PrimExpr) -> Folded {
None => Ok(None),
}
}
+
+//
---------------------------------------------------------------------------------------------
+// Helpers
+
+/// C++ `IsIndexType`: a scalar `int32` or `int64`, which index arithmetic
folds eagerly.
+pub(super) fn is_index_type(dtype: DLDataType) -> bool {
+ dtype.code == support::INT && (dtype.bits == 32 || dtype.bits == 64) &&
dtype.lanes == 1
+}
+
+/// C++ `IsRepresentable(value, dtype)`: whether the exact integer `value` is
a literal of the
+/// scalar type `dtype`, an integer in its range or a boolean 0 or 1 (as is an
`int1`).
+pub(super) fn is_representable(value: &BigInt, dtype: DLDataType) -> bool {
+ if !support::is_scalar(dtype) || dtype.bits == 0 {
+ return false;
+ }
+ let bits = u32::from(dtype.bits);
+ if let Some(small) = value.to_i64() {
+ return match dtype.code {
+ support::UINT => small >= 0 && (bits >= 63 || (small as u64) <
(1u64 << bits)),
+ support::BOOL => small == 0 || small == 1,
+ support::INT if bits == 1 => small == 0 || small == 1,
+ support::INT => {
+ bits >= 64 || (-(1i64 << (bits - 1))..1i64 << (bits -
1)).contains(&small)
+ }
+ _ => false,
+ };
+ }
+ match dtype.code {
+ support::UINT => !value.is_negative() && *value < BigInt::from(1i64)
<< i64::from(bits),
+ support::INT if bits > 1 => {
+ let limit = BigInt::from(1i64) << i64::from(bits - 1);
+ *value >= -limit.clone() && *value < limit
+ }
+ _ => false,
+ }
+}
+
+/// C++ `IsInFloat32Range(value)`: a value that rounds to a finite float32, or
one that is not
+/// finite.
+pub(super) fn is_in_float32_range(value: f64) -> bool {
+ !value.is_finite() ||
(f64::from(f32::MIN)..=f64::from(f32::MAX)).contains(&value)
+}
+
+/// The value of an integer constant.
+pub(super) fn int(expr: &PrimExpr) -> Option<&BigInt> {
+ expr.downcast_ref::<IntImmObj>().map(|imm| &imm.value)
+}
+
+/// The value of a float constant.
+fn float(expr: &PrimExpr) -> Option<f64> {
+ expr.downcast_ref::<FloatImmObj>().map(|imm| imm.value)
+}
+
+/// C++ `IntImm(dtype, value)`.
+fn int_imm(dtype: DLDataType, value: BigInt) -> Folded {
+ 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
+/// the type.
+fn exact(dtype: DLDataType, value: BigInt) -> Folded {
+ match is_representable(&value, dtype) {
+ true => int_imm(dtype, value),
+ false => Ok(None),
+ }
+}
+
+/// C++ `FloatImm(dtype, value)`.
+fn float_imm(dtype: DLDataType, value: f64) -> Folded {
+ Ok(Some(FloatImm::new(PrimType::from_dtype(dtype)?, value,
UNKNOWN_LOC)?.into()))
+}
+
+/// C++ `IntImm::Bool(value)`.
+fn boolean(value: bool) -> Folded {
+ Ok(Some(IntImm::bool(value).into()))
+}
+
+/// A float32 operation on the operands rounded to float32, as C++ computes
`float` arithmetic, or
+/// `None` when the result is beyond the float32 range: C++ checks the result
in double precision
+/// before rounding can turn a finite overflow into infinity, and rounding
that result to float32
+/// gives the single-precision one (a double holds more than twice the digits
of a float).
+fn single(x: f64, y: f64, op: fn(f64, f64) -> f64) -> Option<f64> {
+ let result = op(f64::from(x as f32), f64::from(y as f32));
+ is_in_float32_range(result).then(|| f64::from(result as f32))
+}
+
+/// C++ `FloatImm(dtype, value)` of a folded value, or nothing to fold.
+fn float_imm_of(dtype: DLDataType, value: Option<f64>) -> Folded {
+ match value {
+ Some(value) => float_imm(dtype, value),
+ None => Ok(None),
+ }
+}
+
+/// The error of `TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero"`.
+fn check_int_divisor(divisor: &BigInt) -> Result<()> {
+ match divisor.is_zero() {
+ true => Err(support::compare_failed(
+ INTERNAL_ERROR,
+ "pb->value != 0",
+ divisor,
+ 0,
+ "Divide by zero",
+ )),
+ false => Ok(()),
+ }
+}
+
+/// The error of `TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero"`.
+fn check_float_divisor(divisor: f64) -> Result<()> {
+ match divisor == 0.0 {
+ true => Err(support::compare_failed(
+ INTERNAL_ERROR,
+ "fb->value != 0",
+ support::CppDouble(divisor),
+ 0,
+ "Divide by zero",
+ )),
+ false => Ok(()),
+ }
+}
+
+fn is_one(value: &BigInt) -> bool {
+ value.to_i64() == Some(1)
+}
+
+/// The operand itself.
+fn keep(expr: &PrimExpr) -> Folded {
+ Ok(Some(expr.clone()))
+}
diff --git a/rust/src/ir/prim/expr.rs b/rust/src/ir/prim/expr.rs
index 87c8dcb70..672a7b811 100644
--- a/rust/src/ir/prim/expr.rs
+++ b/rust/src/ir/prim/expr.rs
@@ -210,75 +210,6 @@ use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
// tvm-ffi-stubgen(end)
-//
---------------------------------------------------------------------------------------------
-// What the node constructors check
-
-/// What the constructor of a binary node checks of its operands, and the type
it gives the node
-/// (`TVM_DEFINE_BINOP_CONSTRUCTOR` and its siblings in `src/ir/prim/expr.cc`).
-#[derive(Clone, Copy)]
-enum Operands {
- /// `Add`, `Min`, ...: operands of one type, which is the node's.
- Arith,
- /// `LShift`, `BitwiseAnd`, ...: as `Arith`, of an integer type, or
boolean if allowed.
- Bitwise { allow_bool: bool },
- /// `EQ`, `LT`, ...: operands of one type; the node is boolean.
- Compare,
-}
-
-/// The type of the binary node `node` of `a` and `b`, after the checks of its
C++ constructor.
-fn binary_type(node: &str, operands: Operands, a: &PrimExpr, b: &PrimExpr) ->
Result<PrimType> {
- let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
- if a_dtype != b_dtype {
- let message = format!(
- "mismatched types. {} vs. {}\n",
- support::dtype_name(a_dtype),
- support::dtype_name(b_dtype)
- );
- let condition = "a_ty->dtype == b_ty->dtype";
- return Err(support::check_failed(tvm_ffi::TYPE_ERROR, condition,
&message));
- }
- match operands {
- Operands::Arith => Ok(a.ty()),
- Operands::Bitwise { allow_bool } => {
- let code = a_dtype.code;
- if matches!(code, support::INT | support::UINT) || (allow_bool &&
code == support::BOOL)
- {
- return Ok(a.ty());
- }
- let condition = format!(
- "a.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) || \
- ({allow_bool} && a.ty().MatchesCode(DLDataTypeCode::kDLBool))"
- );
- let operands = if allow_bool { "integer or boolean" } else {
"integer" };
- let message = format!("{node} requires {operands} operands");
- Err(support::check_failed(tvm_ffi::TYPE_ERROR, &condition,
&message))
- }
- Operands::Compare => bool_type(a_dtype),
- }
-}
-
-/// The boolean type with the lanes of `dtype` (C++
`PrimType(DLDataType{kDLBool, 8, lanes})`).
-fn bool_type(operand: DLDataType) -> Result<PrimType> {
- use crate::ir::dtype;
- PrimType::from_dtype(dtype::with_lanes(dtype::bool(), operand.lanes))
-}
-
-/// The type of the logical node of `a` and `b`, after the checks of the C++
`And` and `Or`
-/// constructors: boolean operands of one type.
-fn logical_type(a: &PrimExpr, b: &PrimExpr) -> Result<PrimType> {
- let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
- for (dtype, operand) in [(a_dtype, "a_ty"), (b_dtype, "b_ty")] {
- if dtype.code != support::BOOL {
- let condition =
format!("{operand}.MatchesCode(DLDataTypeCode::kDLBool)");
- return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
&condition, ""));
- }
- }
- if a_dtype != b_dtype {
- return Err(support::check_failed(tvm_ffi::TYPE_ERROR, "a_ty == b_ty",
"mismatched types"));
- }
- bool_type(a_dtype)
-}
-
/// `impl $node { pub fn new(a, b, loc) }` for each binary node, with the
checks of its C++
/// constructor.
macro_rules! binary_new {
@@ -2119,3 +2050,72 @@ impl Let {
Ok(Self::from_fields(loc.into(), ty, var, value, body))
}
}
+
+//
---------------------------------------------------------------------------------------------
+// What the node constructors check
+
+/// What the constructor of a binary node checks of its operands, and the type
it gives the node
+/// (`TVM_DEFINE_BINOP_CONSTRUCTOR` and its siblings in `src/ir/prim/expr.cc`).
+#[derive(Clone, Copy)]
+enum Operands {
+ /// `Add`, `Min`, ...: operands of one type, which is the node's.
+ Arith,
+ /// `LShift`, `BitwiseAnd`, ...: as `Arith`, of an integer type, or
boolean if allowed.
+ Bitwise { allow_bool: bool },
+ /// `EQ`, `LT`, ...: operands of one type; the node is boolean.
+ Compare,
+}
+
+/// The type of the binary node `node` of `a` and `b`, after the checks of its
C++ constructor.
+fn binary_type(node: &str, operands: Operands, a: &PrimExpr, b: &PrimExpr) ->
Result<PrimType> {
+ let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
+ if a_dtype != b_dtype {
+ let message = format!(
+ "mismatched types. {} vs. {}\n",
+ support::dtype_name(a_dtype),
+ support::dtype_name(b_dtype)
+ );
+ let condition = "a_ty->dtype == b_ty->dtype";
+ return Err(support::check_failed(tvm_ffi::TYPE_ERROR, condition,
&message));
+ }
+ match operands {
+ Operands::Arith => Ok(a.ty()),
+ Operands::Bitwise { allow_bool } => {
+ let code = a_dtype.code;
+ if matches!(code, support::INT | support::UINT) || (allow_bool &&
code == support::BOOL)
+ {
+ return Ok(a.ty());
+ }
+ let condition = format!(
+ "a.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) || \
+ ({allow_bool} && a.ty().MatchesCode(DLDataTypeCode::kDLBool))"
+ );
+ let operands = if allow_bool { "integer or boolean" } else {
"integer" };
+ let message = format!("{node} requires {operands} operands");
+ Err(support::check_failed(tvm_ffi::TYPE_ERROR, &condition,
&message))
+ }
+ Operands::Compare => bool_type(a_dtype),
+ }
+}
+
+/// The boolean type with the lanes of `dtype` (C++
`PrimType(DLDataType{kDLBool, 8, lanes})`).
+fn bool_type(operand: DLDataType) -> Result<PrimType> {
+ use crate::ir::dtype;
+ PrimType::from_dtype(dtype::with_lanes(dtype::bool(), operand.lanes))
+}
+
+/// The type of the logical node of `a` and `b`, after the checks of the C++
`And` and `Or`
+/// constructors: boolean operands of one type.
+fn logical_type(a: &PrimExpr, b: &PrimExpr) -> Result<PrimType> {
+ let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
+ for (dtype, operand) in [(a_dtype, "a_ty"), (b_dtype, "b_ty")] {
+ if dtype.code != support::BOOL {
+ let condition =
format!("{operand}.MatchesCode(DLDataTypeCode::kDLBool)");
+ return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
&condition, ""));
+ }
+ }
+ if a_dtype != b_dtype {
+ return Err(support::check_failed(tvm_ffi::TYPE_ERROR, "a_ty == b_ty",
"mismatched types"));
+ }
+ bool_type(a_dtype)
+}
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 8929859bd..c435f85f8 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -209,6 +209,16 @@ use tvm_ffi::VALUE_ERROR;
#[must_use = "an `ExprResult` holds the result; pass it on, or convert it with
`PrimExpr::try_from`"]
pub struct ExprResult(Result<PrimExpr>);
+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) => Ok(expr.dtype()),
+ Err(error) => Err(error.clone()),
+ }
+ }
+}
+
impl TryFrom<ExprResult> for Expr {
type Error = Error;
@@ -227,16 +237,6 @@ impl TryFrom<ExprResult> for PrimExpr {
}
}
-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) => Ok(expr.dtype()),
- Err(error) => Err(error.clone()),
- }
- }
-}
-
impl IntoExpr for ExprResult {
type Output = Expr;
@@ -288,29 +288,6 @@ impl From<&ExprResult> for ExprResult {
}
}
-/// `lhs <op> rhs`; a literal `rhs` is typed after `lhs`.
-fn binary(
- lhs: ExprResult,
- rhs: impl Operand,
- op: impl FnOnce(PrimExpr, PrimExpr) -> Result<PrimExpr>,
-) -> ExprResult {
- let rhs = typed_after(rhs.classify(), &lhs);
- match (lhs.0, rhs.0) {
- (Ok(a), Ok(b)) => ExprResult(op(a, b)),
- (Err(error), _) | (_, Err(error)) => ExprResult(Err(error)),
- }
-}
-
-/// `lhs <op> rhs` for a literal `lhs`, typed after `rhs`.
-fn binary_literal(
- lhs: impl Operand,
- rhs: ExprResult,
- op: impl FnOnce(PrimExpr, PrimExpr) -> Result<PrimExpr>,
-) -> ExprResult {
- let lhs = typed_after(lhs.classify(), &rhs);
- binary(lhs, rhs, op)
-}
-
mod operators {
//! The operator traits of [`ExprResult`], and of the expressions and
literals that start one.
@@ -409,6 +386,29 @@ mod operators {
unary_operators!(ExprResult, &ExprResult);
}
+/// `lhs <op> rhs`; a literal `rhs` is typed after `lhs`.
+fn binary(
+ lhs: ExprResult,
+ rhs: impl Operand,
+ op: impl FnOnce(PrimExpr, PrimExpr) -> Result<PrimExpr>,
+) -> ExprResult {
+ let rhs = typed_after(rhs.classify(), &lhs);
+ match (lhs.0, rhs.0) {
+ (Ok(a), Ok(b)) => ExprResult(op(a, b)),
+ (Err(error), _) | (_, Err(error)) => ExprResult(Err(error)),
+ }
+}
+
+/// `lhs <op> rhs` for a literal `lhs`, typed after `rhs`.
+fn binary_literal(
+ lhs: impl Operand,
+ rhs: ExprResult,
+ op: impl FnOnce(PrimExpr, PrimExpr) -> Result<PrimExpr>,
+) -> ExprResult {
+ let lhs = typed_after(lhs.classify(), &rhs);
+ binary(lhs, rhs, op)
+}
+
//
---------------------------------------------------------------------------------------------
// `PrimOps`: the operations Rust has no operator for
@@ -810,6 +810,45 @@ impl From<f32> for ConstValue {
}
}
+unlocated! {
+ /// Create the constant `value` of the type `ty`, broadcast over the lanes
of a vector type (C++
+ /// `prim::MakeConst(ty, value)`).
+ ///
+ /// # Errors
+ ///
+ /// Returns a `ValueError` if `value` is out of the range of `ty`, or an
`InternalError` if it
+ /// is negative for an unsigned `ty` or `ty` is not a type of constants,
such as `void`.
+ make_const(ty: impl Into<PrimType>, value: impl Into<ConstValue>) =>
make_const_at;
+}
+
+/// Build [`make_const`] at the location `loc`.
+pub fn make_const_at(
+ ty: impl Into<PrimType>,
+ value: impl Into<ConstValue>,
+ loc: impl Into<Location>,
+) -> Result<PrimExpr> {
+ let (dtype, value, loc) = (ty.into().dtype, value.into(), loc.into());
+ if !is_vector(dtype) {
+ return make_const_scalar(dtype, value, &loc);
+ }
+ let scalar = make_const_scalar(with_lanes(dtype, 1), value, &loc)?;
+ let lanes = match support::is_fixed_length_vector(dtype) {
+ true => IntImm::int32(support::lanes(dtype)?).into(),
+ false => super::vector_expr::vscale_lanes(vscale_factor(dtype))?,
+ };
+ Ok(Broadcast::new(scalar, lanes, loc)?.into())
+}
+
+/// Convert `value` into a primitive expression: a Rust `bool`, integer or
float becomes a constant
+/// of the default type, and a primitive expression is returned as is
(`prim.convert`).
+///
+/// # Errors
+///
+/// Returns a `TypeError` if `value` is not a primitive expression or a
literal.
+pub fn convert(value: impl Into<Any>) -> Result<PrimExpr> {
+
tvm_ffi::cached_global_func!("prim.convert").call_tuple((value.into(),))?.try_into()
+}
+
/// C++ `prim::MakeConstScalar(dtype, value, loc)`: an `IntImm` for an integer
or boolean type, a
/// `FloatImm` for a floating-point type.
pub(crate) fn make_const_scalar(
@@ -855,71 +894,151 @@ pub(crate) fn make_const_scalar(
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
+//
---------------------------------------------------------------------------------------------
+// Casts and type matching (C++ `cast` and `BinaryOpMatchTypes`)
+
unlocated! {
- /// Create the constant `value` of the type `ty`, broadcast over the lanes
of a vector type (C++
- /// `prim::MakeConst(ty, value)`).
+ /// Convert `value` to the type `ty`, folding a constant that keeps its
value and broadcasting a
+ /// scalar to a vector type (C++ `prim::cast(value, t)`, `prim._cast`).
+ ///
+ /// [`Cast::new`] builds the node without folding.
///
/// # Errors
///
- /// Returns a `ValueError` if `value` is out of the range of `ty`, or an
`InternalError` if it
- /// is negative for an unsigned `ty` or `ty` is not a type of constants,
such as `void`.
- make_const(ty: impl Into<PrimType>, value: impl Into<ConstValue>) =>
make_const_at;
+ /// Returns a `TypeError` if `value` has the void type, or an
`InternalError` if `value` and
+ /// `ty` are vectors of different lanes.
+ cast(value: impl IntoExpr<PrimExpr>, ty: impl Into<PrimType>) => cast_at;
}
-/// Build [`make_const`] at the location `loc`.
-pub fn make_const_at(
+/// Build [`cast`] at the location `loc`.
+pub fn cast_at(
+ value: impl IntoExpr<PrimExpr>,
ty: impl Into<PrimType>,
- value: impl Into<ConstValue>,
loc: impl Into<Location>,
) -> Result<PrimExpr> {
- let (dtype, value, loc) = (ty.into().dtype, value.into(), loc.into());
- if !is_vector(dtype) {
- return make_const_scalar(dtype, value, &loc);
+ let (value, ty, loc) = (value.into_owned_expr()?, ty.into(), loc.into());
+ let (from, to) = (value.dtype(), ty.dtype);
+ if from == to {
+ return Ok(value);
}
- let scalar = make_const_scalar(with_lanes(dtype, 1), value, &loc)?;
- let lanes = match support::is_fixed_length_vector(dtype) {
- true => IntImm::int32(support::lanes(dtype)?).into(),
- false => super::vector_expr::vscale_lanes(vscale_factor(dtype))?,
+ if from == void() {
+ let message = "Cannot cast an expression with the void sentinel type";
+ return Err(support::check_failed(TYPE_ERROR, "!value.ty().IsVoid()",
message));
+ }
+ if support::is_scalar(to) {
+ return cast_scalar(value, ty, &loc);
+ }
+ let element = PrimType::from_dtype(with_lanes(to, 1))?;
+ if !is_vector(from) {
+ let value = cast_at(value, element, loc.clone())?;
+ return Ok(Broadcast::new(value, lanes_of(to)?, loc)?.into());
+ }
+ if support::is_scalable_vector(from) != support::is_scalable_vector(to) {
+ let condition = "value.ty().IsScalableVector() ==
dtype.IsScalableVector()";
+ return Err(support::check_failed(INTERNAL_ERROR, condition, ""));
+ }
+ let lanes_match = match support::is_scalable_vector(from) {
+ true => vscale_factor(from) == vscale_factor(to),
+ false => support::lanes(from)? == support::lanes(to)?,
};
- Ok(Broadcast::new(scalar, lanes, loc)?.into())
+ if !lanes_match {
+ return Err(support::check_failed(INTERNAL_ERROR, "lanes_match", ""));
+ }
+ if let Some(broadcast) = value.downcast_ref::<BroadcastObj>() {
+ let value = cast_at(broadcast.value.clone(), element, loc.clone())?;
+ return Ok(Broadcast::new(value, broadcast.lanes.clone(), loc)?.into());
+ }
+ if let Some(ramp) = value.downcast_ref::<RampObj>() {
+ // Casting the base and the stride is valid only when every lane is
the same integer in
+ // both types: widening must not remove an overflow of the source, and
a floating-point
+ // cast need not be affine.
+ let integer = matches!(to.code, support::INT | support::UINT) &&
matches!(to.bits, 32 | 64);
+ if let (Some(base), Some(stride)) =
+ (const_fold::int(&ramp.base_), const_fold::int(&ramp.stride))
+ {
+ if support::is_fixed_length_vector(to) && integer {
+ let last = base + stride * i64::from(support::lanes(to)? - 1);
+ let source = support::prim_dtype(&ramp.base_)?;
+ let fits = |value: &BigInt|
const_fold::is_representable(value, element.dtype);
+ if fits(base)
+ && fits(stride)
+ && fits(&last)
+ && const_fold::is_representable(&last, source)
+ {
+ let base = cast_at(ramp.base_.clone(), element.clone(),
loc.clone())?;
+ let stride = cast_at(ramp.stride.clone(), element,
loc.clone())?;
+ return Ok(Ramp::new(base, stride, ramp.lanes.clone(),
loc)?.into());
+ }
+ }
+ }
+ }
+ Ok(Cast::new(ty, value, loc)?.into())
}
-/// Convert `value` into a primitive expression: a Rust `bool`, integer or
float becomes a constant
-/// of the default type, and a primitive expression is returned as is
(`prim.convert`).
+/// Bring `a` and `b` to one type, as every binary operator does first (C++
+/// `BinaryOpMatchTypes(lhs, rhs)`).
+///
+/// A scalar next to a vector is broadcast, and both are cast to the promoted
type: the wider float,
+/// a float over an integer, the wider integer, unsigned over signed.
///
/// # Errors
///
-/// Returns a `TypeError` if `value` is not a primitive expression or a
literal.
-pub fn convert(value: impl Into<Any>) -> Result<PrimExpr> {
-
tvm_ffi::cached_global_func!("prim.convert").call_tuple((value.into(),))?.try_into()
+/// Returns an `InternalError` if the lanes of the two differ or their types
cannot be promoted to
+/// one.
+pub fn match_types(
+ a: impl IntoExpr<PrimExpr>,
+ b: impl IntoExpr<PrimExpr>,
+) -> Result<(PrimExpr, PrimExpr)> {
+ match_types_at(a, b, UNKNOWN_LOC)
}
-//
---------------------------------------------------------------------------------------------
-// Casts and type matching (C++ `cast` and `BinaryOpMatchTypes`)
-
-/// The exponent C++ `frexp` gives a finite `value`, `e` with `value = m *
2^e` and
-/// `0.5 <= |m| < 1`, and 0 for zero.
-fn frexp_exponent(value: f64) -> i32 {
- match ((value.to_bits() >> 52) & 0x7ff) as i32 {
- 0 if value == 0.0 => 0,
- // A subnormal, scaled into the normal range.
- 0 => frexp_exponent(value * pow2(64)) - 64,
- biased => biased - 1022,
+/// Run [`match_types`] with the casts at the location `loc`.
+pub fn match_types_at(
+ a: impl IntoExpr<PrimExpr>,
+ b: impl IntoExpr<PrimExpr>,
+ loc: impl Into<Location>,
+) -> Result<(PrimExpr, PrimExpr)> {
+ let (a, b, loc) = (a.into_owned_expr()?, b.into_owned_expr()?, loc.into());
+ let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
+ if a_dtype == b_dtype {
+ return Ok((a, b));
+ }
+ let a = broadcast_to_match_lanes(a, b_dtype)?;
+ let b = broadcast_to_match_lanes(b, a.dtype())?;
+ let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
+ if support::is_scalable_vector(a_dtype) !=
support::is_scalable_vector(b_dtype) {
+ let condition = "lhs_ty.IsScalableVector() ==
rhs_ty.IsScalableVector()";
+ let message = "Can't match scalable and fixed length vectors";
+ return Err(support::check_failed(INTERNAL_ERROR, condition, message));
+ }
+ let lanes_match = match support::is_scalable_vector(a_dtype) {
+ true => vscale_factor(a_dtype) == vscale_factor(b_dtype),
+ false => support::lanes(a_dtype)? == support::lanes(b_dtype)?,
+ };
+ if !lanes_match {
+ let (lhs, rhs) = (support::dtype_name(a_dtype),
support::dtype_name(b_dtype));
+ let message = format!("Cannot match type {lhs} vs {rhs}");
+ return Err(support::check_failed(INTERNAL_ERROR, "lanes_match",
&message));
}
+ let promoted = promote_binary_op_type(a_dtype, b_dtype)?;
+ let a = match a_dtype == promoted {
+ true => a,
+ false => cast_at(a, PrimType::from_dtype(promoted)?, loc.clone())?,
+ };
+ let b = match b_dtype == promoted {
+ true => b,
+ false => cast_at(b, PrimType::from_dtype(promoted)?, loc)?,
+ };
+ Ok((a, b))
}
-/// `2^exponent` for an exponent of a normal double.
-fn pow2(exponent: i32) -> f64 {
- f64::from_bits(((exponent + 1023) as u64) << 52)
-}
-
-/// Whether the finite `value` of a float16 or bfloat16 `FloatImm` is a value
of its type: a
-/// `FloatImm` holds an unrounded double, which stands for the value the type
rounds it to.
-fn is_rounded_to_16_bits(value: f64, bfloat: bool) -> bool {
- let (digits, min_step) = if bfloat { (8, -133) } else { (11, -24) };
- let step = (frexp_exponent(value) - digits).max(min_step);
- let scaled = value * pow2(-step);
- scaled.trunc() == scaled
+/// C++ `cast` to a scalar type: a constant folded to `ty` keeping its
location (see
+/// [`fold_cast`]); anything else a [`Cast`].
+fn cast_scalar(value: PrimExpr, ty: PrimType, loc: &Location) ->
Result<PrimExpr> {
+ match fold_cast(&value, ty.dtype)? {
+ Some(folded) => Ok(folded),
+ None => Ok(Cast::new(ty, value, loc.clone())?.into()),
+ }
}
/// A constant cast to the scalar type `to` as C++ `cast` folds it: only when
the result is the
@@ -1005,91 +1124,29 @@ fn fold_cast(value: &PrimExpr, to: DLDataType) ->
Result<Option<PrimExpr>> {
Ok(None)
}
-/// C++ `cast` to a scalar type: a constant folded to `ty` keeping its
location (see
-/// [`fold_cast`]); anything else a [`Cast`].
-fn cast_scalar(value: PrimExpr, ty: PrimType, loc: &Location) ->
Result<PrimExpr> {
- match fold_cast(&value, ty.dtype)? {
- Some(folded) => Ok(folded),
- None => Ok(Cast::new(ty, value, loc.clone())?.into()),
- }
+/// Whether the finite `value` of a float16 or bfloat16 `FloatImm` is a value
of its type: a
+/// `FloatImm` holds an unrounded double, which stands for the value the type
rounds it to.
+fn is_rounded_to_16_bits(value: f64, bfloat: bool) -> bool {
+ let (digits, min_step) = if bfloat { (8, -133) } else { (11, -24) };
+ let step = (frexp_exponent(value) - digits).max(min_step);
+ let scaled = value * pow2(-step);
+ scaled.trunc() == scaled
}
-unlocated! {
- /// Convert `value` to the type `ty`, folding a constant that keeps its
value and broadcasting a
- /// scalar to a vector type (C++ `prim::cast(value, t)`, `prim._cast`).
- ///
- /// [`Cast::new`] builds the node without folding.
- ///
- /// # Errors
- ///
- /// Returns a `TypeError` if `value` has the void type, or an
`InternalError` if `value` and
- /// `ty` are vectors of different lanes.
- cast(value: impl IntoExpr<PrimExpr>, ty: impl Into<PrimType>) => cast_at;
+/// The exponent C++ `frexp` gives a finite `value`, `e` with `value = m *
2^e` and
+/// `0.5 <= |m| < 1`, and 0 for zero.
+fn frexp_exponent(value: f64) -> i32 {
+ match ((value.to_bits() >> 52) & 0x7ff) as i32 {
+ 0 if value == 0.0 => 0,
+ // A subnormal, scaled into the normal range.
+ 0 => frexp_exponent(value * pow2(64)) - 64,
+ biased => biased - 1022,
+ }
}
-/// Build [`cast`] at the location `loc`.
-pub fn cast_at(
- value: impl IntoExpr<PrimExpr>,
- ty: impl Into<PrimType>,
- loc: impl Into<Location>,
-) -> Result<PrimExpr> {
- let (value, ty, loc) = (value.into_owned_expr()?, ty.into(), loc.into());
- let (from, to) = (value.dtype(), ty.dtype);
- if from == to {
- return Ok(value);
- }
- if from == void() {
- let message = "Cannot cast an expression with the void sentinel type";
- return Err(support::check_failed(TYPE_ERROR, "!value.ty().IsVoid()",
message));
- }
- if support::is_scalar(to) {
- return cast_scalar(value, ty, &loc);
- }
- let element = PrimType::from_dtype(with_lanes(to, 1))?;
- if !is_vector(from) {
- let value = cast_at(value, element, loc.clone())?;
- return Ok(Broadcast::new(value, lanes_of(to)?, loc)?.into());
- }
- if support::is_scalable_vector(from) != support::is_scalable_vector(to) {
- let condition = "value.ty().IsScalableVector() ==
dtype.IsScalableVector()";
- return Err(support::check_failed(INTERNAL_ERROR, condition, ""));
- }
- let lanes_match = match support::is_scalable_vector(from) {
- true => vscale_factor(from) == vscale_factor(to),
- false => support::lanes(from)? == support::lanes(to)?,
- };
- if !lanes_match {
- return Err(support::check_failed(INTERNAL_ERROR, "lanes_match", ""));
- }
- if let Some(broadcast) = value.downcast_ref::<BroadcastObj>() {
- let value = cast_at(broadcast.value.clone(), element, loc.clone())?;
- return Ok(Broadcast::new(value, broadcast.lanes.clone(), loc)?.into());
- }
- if let Some(ramp) = value.downcast_ref::<RampObj>() {
- // Casting the base and the stride is valid only when every lane is
the same integer in
- // both types: widening must not remove an overflow of the source, and
a floating-point
- // cast need not be affine.
- let integer = matches!(to.code, support::INT | support::UINT) &&
matches!(to.bits, 32 | 64);
- if let (Some(base), Some(stride)) =
- (const_fold::int(&ramp.base_), const_fold::int(&ramp.stride))
- {
- if support::is_fixed_length_vector(to) && integer {
- let last = base + stride * i64::from(support::lanes(to)? - 1);
- let source = support::prim_dtype(&ramp.base_)?;
- let fits = |value: &BigInt|
const_fold::is_representable(value, element.dtype);
- if fits(base)
- && fits(stride)
- && fits(&last)
- && const_fold::is_representable(&last, source)
- {
- let base = cast_at(ramp.base_.clone(), element.clone(),
loc.clone())?;
- let stride = cast_at(ramp.stride.clone(), element,
loc.clone())?;
- return Ok(Ramp::new(base, stride, ramp.lanes.clone(),
loc)?.into());
- }
- }
- }
- }
- Ok(Cast::new(ty, value, loc)?.into())
+/// `2^exponent` for an exponent of a normal double.
+fn pow2(exponent: i32) -> f64 {
+ f64::from_bits(((exponent + 1023) as u64) << 52)
}
/// C++ `BroadcastToMatchLanes`: `value`, broadcast to the lanes of `other`
when it is a scalar
@@ -1145,63 +1202,6 @@ fn promote_binary_op_type(lhs: DLDataType, rhs:
DLDataType) -> Result<DLDataType
}
}
-/// Bring `a` and `b` to one type, as every binary operator does first (C++
-/// `BinaryOpMatchTypes(lhs, rhs)`).
-///
-/// A scalar next to a vector is broadcast, and both are cast to the promoted
type: the wider float,
-/// a float over an integer, the wider integer, unsigned over signed.
-///
-/// # Errors
-///
-/// Returns an `InternalError` if the lanes of the two differ or their types
cannot be promoted to
-/// one.
-pub fn match_types(
- a: impl IntoExpr<PrimExpr>,
- b: impl IntoExpr<PrimExpr>,
-) -> Result<(PrimExpr, PrimExpr)> {
- match_types_at(a, b, UNKNOWN_LOC)
-}
-
-/// Run [`match_types`] with the casts at the location `loc`.
-pub fn match_types_at(
- a: impl IntoExpr<PrimExpr>,
- b: impl IntoExpr<PrimExpr>,
- loc: impl Into<Location>,
-) -> Result<(PrimExpr, PrimExpr)> {
- let (a, b, loc) = (a.into_owned_expr()?, b.into_owned_expr()?, loc.into());
- let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
- if a_dtype == b_dtype {
- return Ok((a, b));
- }
- let a = broadcast_to_match_lanes(a, b_dtype)?;
- let b = broadcast_to_match_lanes(b, a.dtype())?;
- let (a_dtype, b_dtype) = (a.dtype(), b.dtype());
- if support::is_scalable_vector(a_dtype) !=
support::is_scalable_vector(b_dtype) {
- let condition = "lhs_ty.IsScalableVector() ==
rhs_ty.IsScalableVector()";
- let message = "Can't match scalable and fixed length vectors";
- return Err(support::check_failed(INTERNAL_ERROR, condition, message));
- }
- let lanes_match = match support::is_scalable_vector(a_dtype) {
- true => vscale_factor(a_dtype) == vscale_factor(b_dtype),
- false => support::lanes(a_dtype)? == support::lanes(b_dtype)?,
- };
- if !lanes_match {
- let (lhs, rhs) = (support::dtype_name(a_dtype),
support::dtype_name(b_dtype));
- let message = format!("Cannot match type {lhs} vs {rhs}");
- return Err(support::check_failed(INTERNAL_ERROR, "lanes_match",
&message));
- }
- let promoted = promote_binary_op_type(a_dtype, b_dtype)?;
- let a = match a_dtype == promoted {
- true => a,
- false => cast_at(a, PrimType::from_dtype(promoted)?, loc.clone())?,
- };
- let b = match b_dtype == promoted {
- true => b,
- false => cast_at(b, PrimType::from_dtype(promoted)?, loc)?,
- };
- Ok((a, b))
-}
-
//
---------------------------------------------------------------------------------------------
// Arithmetic (`src/ir/prim/op.cc`)
@@ -1264,26 +1264,6 @@ folded_binary!(
ne / ne_at: fold_ne => NE,
);
-fn fold_gt(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
- const_fold::compare(a, b, |x, y| x > y, |x, y| x > y)
-}
-
-fn fold_ge(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
- const_fold::compare(a, b, |x, y| x >= y, |x, y| x >= y)
-}
-
-fn fold_lt(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
- const_fold::compare(a, b, |x, y| x < y, |x, y| x < y)
-}
-
-fn fold_le(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
- const_fold::compare(a, b, |x, y| x <= y, |x, y| x <= y)
-}
-
-fn fold_ne(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
- const_fold::compare(a, b, |x, y| x != y, |x, y| x != y)
-}
-
unlocated! {
/// Build the boolean `a == b` (C++ `equal`, `prim._OpEQ`).
eq(a: impl Operand, b: impl Operand) => eq_at;
@@ -1411,16 +1391,6 @@ pub fn ceildiv_at(a: impl Operand, b: impl Operand, loc:
impl Into<Location>) ->
}
}
-/// `SymbolicLimits::pos_inf_` and `neg_inf_`, the opaque infinities of
integer analysis, which
-/// [`min`] and [`max`] recognize by identity; fetched once per thread.
-fn symbolic_infinities() -> (PrimExpr, PrimExpr) {
- use std::cell::OnceCell;
- thread_local! {
- static INFINITIES: OnceCell<(PrimExpr, PrimExpr)> = const {
OnceCell::new() };
- }
- support::shared(&INFINITIES, || Ok((sym::pos_inf()?, sym::neg_inf()?)))
-}
-
/// Build [`min`] at the location `loc`.
pub fn min_at(a: impl Operand, b: impl Operand, loc: impl Into<Location>) ->
Result<PrimExpr> {
let ((a, b), loc) = (operands(a, b)?, loc.into());
@@ -1467,6 +1437,36 @@ pub fn max_at(a: impl Operand, b: impl Operand, loc:
impl Into<Location>) -> Res
}
}
+fn fold_gt(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
+ const_fold::compare(a, b, |x, y| x > y, |x, y| x > y)
+}
+
+fn fold_ge(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
+ const_fold::compare(a, b, |x, y| x >= y, |x, y| x >= y)
+}
+
+fn fold_lt(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
+ const_fold::compare(a, b, |x, y| x < y, |x, y| x < y)
+}
+
+fn fold_le(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
+ const_fold::compare(a, b, |x, y| x <= y, |x, y| x <= y)
+}
+
+fn fold_ne(a: &PrimExpr, b: &PrimExpr) -> const_fold::Folded {
+ const_fold::compare(a, b, |x, y| x != y, |x, y| x != y)
+}
+
+/// `SymbolicLimits::pos_inf_` and `neg_inf_`, the opaque infinities of
integer analysis, which
+/// [`min`] and [`max`] recognize by identity; fetched once per thread.
+fn symbolic_infinities() -> (PrimExpr, PrimExpr) {
+ use std::cell::OnceCell;
+ thread_local! {
+ static INFINITIES: OnceCell<(PrimExpr, PrimExpr)> = const {
OnceCell::new() };
+ }
+ support::shared(&INFINITIES, || Ok((sym::pos_inf()?, sym::neg_inf()?)))
+}
+
//
---------------------------------------------------------------------------------------------
// Logic
@@ -1574,20 +1574,6 @@ pub fn any_at(
fold_conditions(conds, loc.into(), "any", logical_or_at)
}
-fn fold_conditions(
- conds: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
- loc: Location,
- name: &str,
- op: fn(PrimExpr, PrimExpr, Location) -> Result<PrimExpr>,
-) -> Result<PrimExpr> {
- 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, ""));
- };
- conds.try_fold(first?, |acc, cond| op(acc, cond?, loc.clone()))
-}
-
/// Build [`if_then_else`] at the location `loc`.
pub fn if_then_else_at(
cond: impl IntoExpr<PrimExpr>,
@@ -1609,46 +1595,23 @@ pub fn if_then_else_at(
call(Some(&ty), "prim.if_then_else", vec![cond, true_value, false_value],
loc)
}
-//
---------------------------------------------------------------------------------------------
-// Bitwise operators
-
-/// The folding of a shift of index types: a constant amount is checked
against the width, two
-/// constants fold to the value `shift` gives them, if any, and a shift by
zero is the operand
-/// itself (`TVM_PRIM_INDEX_CONST_PROPAGATION`).
-fn fold_shift(
- a: &PrimExpr,
- b: &PrimExpr,
- shift: impl FnOnce(&BigInt, &BigInt) -> Result<Option<BigInt>>,
- loc: &Location,
-) -> const_fold::Folded {
- let dtype = a.dtype();
- if !const_fold::is_index_type(dtype) ||
!const_fold::is_index_type(b.dtype()) {
- return Ok(None);
- }
- let (pa, pb) = (const_fold::int(a), const_fold::int(b));
- if let Some(y) = pb {
- if y.is_negative() || *y >= BigInt::from(i64::from(dtype.bits)) {
- let ty = a.ty();
- let condition = "pb->value >= 0 && pb->value < result_ty.bits()";
- let message = format!(
- "Shift amount must be non-negative and less than {} for type
{}",
- dtype.bits,
- support::repr(&ty)
- );
- return Err(support::check_failed(INTERNAL_ERROR, condition,
&message));
- }
- }
- if let (Some(x), Some(y)) = (pa, pb) {
- if let Some(value) = shift(x, y)? {
- return Ok(Some(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into()));
- }
- }
- if pb.is_some_and(BigInt::is_zero) {
- return Ok(Some(a.clone()));
- }
- Ok(None)
+fn fold_conditions(
+ conds: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
+ loc: Location,
+ name: &str,
+ op: fn(PrimExpr, PrimExpr, Location) -> Result<PrimExpr>,
+) -> Result<PrimExpr> {
+ 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, ""));
+ };
+ conds.try_fold(first?, |acc, cond| op(acc, cond?, loc.clone()))
}
+//
---------------------------------------------------------------------------------------------
+// Bitwise operators
+
unlocated! {
/// Build `a << b` of two integers (C++ `left_shift`, `prim.left_shift`).
///
@@ -1776,30 +1739,46 @@ pub fn bitwise_not_at(a: impl IntoExpr<PrimExpr>, loc:
impl Into<Location>) -> R
Ok(BitwiseNot::new(a, loc)?.into())
}
-//
---------------------------------------------------------------------------------------------
-// Math functions (`src/ir/prim/math.cc`)
-
-/// `floor`, `ceil`, `round`, `nearbyint` and `trunc`: an integer or a boolean
is itself, a float
-/// constant is folded with `fold`, keeping its location, and anything else is
a call.
-fn rounding(
- x: PrimExpr,
- op: &'static str,
- fold: fn(f64) -> f64,
- loc: Location,
-) -> Result<PrimExpr> {
- let dtype = x.dtype();
- if matches!(dtype.code, support::INT | support::UINT | support::BOOL) {
- return Ok(x);
+/// The folding of a shift of index types: a constant amount is checked
against the width, two
+/// constants fold to the value `shift` gives them, if any, and a shift by
zero is the operand
+/// itself (`TVM_PRIM_INDEX_CONST_PROPAGATION`).
+fn fold_shift(
+ a: &PrimExpr,
+ b: &PrimExpr,
+ shift: impl FnOnce(&BigInt, &BigInt) -> Result<Option<BigInt>>,
+ loc: &Location,
+) -> const_fold::Folded {
+ let dtype = a.dtype();
+ if !const_fold::is_index_type(dtype) ||
!const_fold::is_index_type(b.dtype()) {
+ return Ok(None);
}
- if let Some(imm) = x.downcast_ref::<FloatImmObj>() {
- return Ok(
- FloatImm::new(PrimType::from_dtype(dtype)?, fold(imm.value),
imm.loc.clone())?.into()
- );
+ let (pa, pb) = (const_fold::int(a), const_fold::int(b));
+ if let Some(y) = pb {
+ if y.is_negative() || *y >= BigInt::from(i64::from(dtype.bits)) {
+ let ty = a.ty();
+ let condition = "pb->value >= 0 && pb->value < result_ty.bits()";
+ let message = format!(
+ "Shift amount must be non-negative and less than {} for type
{}",
+ dtype.bits,
+ support::repr(&ty)
+ );
+ return Err(support::check_failed(INTERNAL_ERROR, condition,
&message));
+ }
}
- let ty = x.ty.clone();
- call(Some(&ty), op, vec![x], loc)
+ if let (Some(x), Some(y)) = (pa, pb) {
+ if let Some(value) = shift(x, y)? {
+ return Ok(Some(IntImm::new(PrimType::from_dtype(dtype)?, value,
loc.clone())?.into()));
+ }
+ }
+ if pb.is_some_and(BigInt::is_zero) {
+ return Ok(Some(a.clone()));
+ }
+ Ok(None)
}
+//
---------------------------------------------------------------------------------------------
+// Math functions (`src/ir/prim/math.cc`)
+
/// The unary math functions, each in a pair with its `_at` variant,
implemented by a function
/// of the expression and the location.
macro_rules! unary {
@@ -1865,94 +1844,6 @@ unary! {
};
}
-fn absolute(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
- let dtype = x.dtype();
- match dtype.code {
- support::INT => integer_abs(x, loc),
- support::FLOAT | support::BFLOAT => match
x.downcast_ref::<FloatImmObj>() {
- 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)
- }
- },
- support::UINT => Ok(x),
- _ => {
- let message = format!(
- "Data type {} not supported for absolute op. Skipping absolute
op...",
- support::repr(&x.ty)
- );
- Err(Error::new(INTERNAL_ERROR, &message, ""))
- }
- }
-}
-
-/// C++ `IntegerAbs(x, loc)` of a signed integer: a constant folds when its
absolute value is a
-/// value of its type; anything else is `select(x >= 0, x, -x)`.
-fn integer_abs(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
- let dtype = x.dtype();
- if let Some(imm) = x.downcast_ref::<IntImmObj>() {
- let value = match imm.value.is_negative() {
- true => -imm.value.clone(),
- false => imm.value.clone(),
- };
- if const_fold::is_representable(&value, dtype) {
- return Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
imm.loc.clone())?.into());
- }
- }
- let zero = make_const(PrimType::from_dtype(dtype)?, 0)?;
- let nonnegative = ge(x.clone(), zero)?;
- let negated = neg(x.clone())?;
- Ok(Select::new(nonnegative, x, negated, loc)?.into())
-}
-
-fn infinite(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
- let dtype = x.dtype();
- let result = PrimType::bool().with_lanes(support::lanes(dtype)? as u16)?;
- match dtype.code {
- support::INT | support::UINT => make_const_at(result, false, loc),
- support::FLOAT => {
- let infinity = infinity_at(x.ty(), loc.clone())?;
- let at_infinity = eq(abs_at(x.clone(), loc.clone())?, infinity)?;
- let not_nan = logical_not(isnan_at(x, loc)?)?;
- logical_and(at_infinity, not_nan)
- }
- _ => {
- let message = format!(
- "Data type {} not supported for finiteness ops. Skipping
it...",
- support::repr(&x.ty)
- );
- Err(Error::new(INTERNAL_ERROR, &message, ""))
- }
- }
-}
-
-fn finite(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
- let not_infinite = logical_not(isinf_at(x.clone(), loc.clone())?)?;
- let not_nan = logical_not(isnan_at(x, loc)?)?;
- logical_and(not_infinite, not_nan)
-}
-
-fn likely_hint(cond: PrimExpr, loc: Location) -> Result<PrimExpr> {
- if cond.downcast_ref::<IntImmObj>().is_some() {
- return Ok(cond);
- }
- let ty = cond.ty.clone();
- call(Some(&ty), "prim.likely", vec![cond], loc)
-}
-
-/// The check `TVM_FFI_ICHECK_EQ(dtype.lanes(), 1)` of the limits.
-fn check_scalar(dtype: DLDataType) -> Result<()> {
- let lanes = support::lanes(dtype)?;
- match lanes == 1 {
- true => Ok(()),
- false => Err(support::compare_failed(INTERNAL_ERROR, "dtype.lanes() ==
1", lanes, 1, "")),
- }
-}
-
unlocated! {
/// Create the largest value of the scalar type `ty` (C++ `max_value`,
`prim.max_value`).
///
@@ -2063,24 +1954,6 @@ pub fn vscale_at(loc: impl Into<Location>) ->
Result<PrimExpr> {
call(Some(&ty), "prim.vscale", vec![], loc.into())
}
-/// A function of two matched floats, called with the type of the first (`pow`
and `fmod`).
-fn float_binary(
- x: PrimExpr,
- y: PrimExpr,
- op: &'static str,
- what: &str,
- loc: Location,
-) -> Result<PrimExpr> {
- let (x, y) = match_types_at(x, y, loc.clone())?;
- if x.dtype().code != support::FLOAT {
- let condition = "x.ty().MatchesCode(DLDataTypeCode::kDLFloat)";
- let message = format!("{what} only applies to float");
- return Err(support::check_failed(INTERNAL_ERROR, condition, &message));
- }
- let ty = x.ty.clone();
- call(Some(&ty), op, vec![x, y], loc)
-}
-
unlocated! {
/// Build `x` to the power `y`, of two matched floats (C++ `pow`,
`prim._OpPow`).
///
@@ -2132,40 +2005,136 @@ pub fn fma_at(
call(Some(&ty), "prim.fma", vec![x, y.into_owned_expr()?,
z.into_owned_expr()?], loc.into())
}
-//
---------------------------------------------------------------------------------------------
-// Intrinsics (`TVM_DECLARE_INTRIN_UNARY` and friends in
`include/tvm/ir/prim/op.h`)
-
-/// A unary intrinsic `op`, a call typed as its argument. A bfloat16 argument
is computed in
-/// float32 and cast back; with `float_only`, any argument but a float or a
bfloat16 is a
-/// `TypeError` (C++ `CheckMathUnaryOpInputDType`).
-fn unary_intrinsic(
- op: &'static str,
+/// `floor`, `ceil`, `round`, `nearbyint` and `trunc`: an integer or a boolean
is itself, a float
+/// constant is folded with `fold`, keeping its location, and anything else is
a call.
+fn rounding(
x: PrimExpr,
- float_only: bool,
+ op: &'static str,
+ fold: fn(f64) -> f64,
loc: Location,
) -> Result<PrimExpr> {
let dtype = x.dtype();
- if float_only && dtype.code != support::FLOAT && !is_bfloat16(dtype) {
- let condition = "dtype.code() == DLDataTypeCode::kDLFloat || \
- dtype.MatchesElementType(DLDataTypeCode::kDLBfloat,
16)";
- let message =
- format!("{op} only supports floating-point inputs, but got {}",
support::repr(&x.ty));
- return Err(support::check_failed(TYPE_ERROR, condition, &message));
+ if matches!(dtype.code, support::INT | support::UINT | support::BOOL) {
+ return Ok(x);
}
- if !is_bfloat16(dtype) {
- let ty = x.ty.clone();
- return call(Some(&ty), op, vec![x], loc);
+ if let Some(imm) = x.downcast_ref::<FloatImmObj>() {
+ return Ok(
+ FloatImm::new(PrimType::from_dtype(dtype)?, fold(imm.value),
imm.loc.clone())?.into()
+ );
}
- let float32 = match support::is_scalable_vector(dtype) {
- true => support::scalable_vector(PrimType::float(32).dtype,
vscale_factor(dtype))?,
- false => with_lanes(PrimType::float(32).dtype, dtype.lanes),
- };
- let float32 = PrimType::from_dtype(float32)?;
- let x_float32: PrimExpr = Cast::new(float32.clone(), x,
loc.clone())?.into();
- let result = call(Some(&float32.into()), op, vec![x_float32],
loc.clone())?;
- Ok(Cast::new(PrimType::from_dtype(dtype)?, result, loc)?.into())
+ let ty = x.ty.clone();
+ call(Some(&ty), op, vec![x], loc)
+}
+
+fn absolute(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
+ let dtype = x.dtype();
+ match dtype.code {
+ support::INT => integer_abs(x, loc),
+ support::FLOAT | support::BFLOAT => match
x.downcast_ref::<FloatImmObj>() {
+ 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)
+ }
+ },
+ support::UINT => Ok(x),
+ _ => {
+ let message = format!(
+ "Data type {} not supported for absolute op. Skipping absolute
op...",
+ support::repr(&x.ty)
+ );
+ Err(Error::new(INTERNAL_ERROR, &message, ""))
+ }
+ }
+}
+
+/// C++ `IntegerAbs(x, loc)` of a signed integer: a constant folds when its
absolute value is a
+/// value of its type; anything else is `select(x >= 0, x, -x)`.
+fn integer_abs(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
+ let dtype = x.dtype();
+ if let Some(imm) = x.downcast_ref::<IntImmObj>() {
+ let value = match imm.value.is_negative() {
+ true => -imm.value.clone(),
+ false => imm.value.clone(),
+ };
+ if const_fold::is_representable(&value, dtype) {
+ return Ok(IntImm::new(PrimType::from_dtype(dtype)?, value,
imm.loc.clone())?.into());
+ }
+ }
+ let zero = make_const(PrimType::from_dtype(dtype)?, 0)?;
+ let nonnegative = ge(x.clone(), zero)?;
+ let negated = neg(x.clone())?;
+ Ok(Select::new(nonnegative, x, negated, loc)?.into())
}
+fn infinite(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
+ let dtype = x.dtype();
+ let result = PrimType::bool().with_lanes(support::lanes(dtype)? as u16)?;
+ match dtype.code {
+ support::INT | support::UINT => make_const_at(result, false, loc),
+ support::FLOAT => {
+ let infinity = infinity_at(x.ty(), loc.clone())?;
+ let at_infinity = eq(abs_at(x.clone(), loc.clone())?, infinity)?;
+ let not_nan = logical_not(isnan_at(x, loc)?)?;
+ logical_and(at_infinity, not_nan)
+ }
+ _ => {
+ let message = format!(
+ "Data type {} not supported for finiteness ops. Skipping
it...",
+ support::repr(&x.ty)
+ );
+ Err(Error::new(INTERNAL_ERROR, &message, ""))
+ }
+ }
+}
+
+fn finite(x: PrimExpr, loc: Location) -> Result<PrimExpr> {
+ let not_infinite = logical_not(isinf_at(x.clone(), loc.clone())?)?;
+ let not_nan = logical_not(isnan_at(x, loc)?)?;
+ logical_and(not_infinite, not_nan)
+}
+
+fn likely_hint(cond: PrimExpr, loc: Location) -> Result<PrimExpr> {
+ if cond.downcast_ref::<IntImmObj>().is_some() {
+ return Ok(cond);
+ }
+ let ty = cond.ty.clone();
+ call(Some(&ty), "prim.likely", vec![cond], loc)
+}
+
+/// The check `TVM_FFI_ICHECK_EQ(dtype.lanes(), 1)` of the limits.
+fn check_scalar(dtype: DLDataType) -> Result<()> {
+ let lanes = support::lanes(dtype)?;
+ match lanes == 1 {
+ true => Ok(()),
+ false => Err(support::compare_failed(INTERNAL_ERROR, "dtype.lanes() ==
1", lanes, 1, "")),
+ }
+}
+
+/// A function of two matched floats, called with the type of the first (`pow`
and `fmod`).
+fn float_binary(
+ x: PrimExpr,
+ y: PrimExpr,
+ op: &'static str,
+ what: &str,
+ loc: Location,
+) -> Result<PrimExpr> {
+ let (x, y) = match_types_at(x, y, loc.clone())?;
+ if x.dtype().code != support::FLOAT {
+ let condition = "x.ty().MatchesCode(DLDataTypeCode::kDLFloat)";
+ let message = format!("{what} only applies to float");
+ return Err(support::check_failed(INTERNAL_ERROR, condition, &message));
+ }
+ let ty = x.ty.clone();
+ call(Some(&ty), op, vec![x, y], loc)
+}
+
+//
---------------------------------------------------------------------------------------------
+// Intrinsics (`TVM_DECLARE_INTRIN_UNARY` and friends in
`include/tvm/ir/prim/op.h`)
+
/// The unary intrinsics; `float_only` marks `TVM_DECLARE_FLOAT_INTRIN_UNARY`.
macro_rules! unary_intrinsics {
($($(#[$doc:meta])* $name:ident / $at:ident => $op:literal,
$float_only:literal;)+) => {
@@ -2270,6 +2239,37 @@ binary_intrinsics! {
ldexp / ldexp_at => "prim.ldexp";
}
+/// A unary intrinsic `op`, a call typed as its argument. A bfloat16 argument
is computed in
+/// float32 and cast back; with `float_only`, any argument but a float or a
bfloat16 is a
+/// `TypeError` (C++ `CheckMathUnaryOpInputDType`).
+fn unary_intrinsic(
+ op: &'static str,
+ x: PrimExpr,
+ float_only: bool,
+ loc: Location,
+) -> Result<PrimExpr> {
+ let dtype = x.dtype();
+ if float_only && dtype.code != support::FLOAT && !is_bfloat16(dtype) {
+ let condition = "dtype.code() == DLDataTypeCode::kDLFloat || \
+ dtype.MatchesElementType(DLDataTypeCode::kDLBfloat,
16)";
+ let message =
+ format!("{op} only supports floating-point inputs, but got {}",
support::repr(&x.ty));
+ return Err(support::check_failed(TYPE_ERROR, condition, &message));
+ }
+ if !is_bfloat16(dtype) {
+ let ty = x.ty.clone();
+ return call(Some(&ty), op, vec![x], loc);
+ }
+ let float32 = match support::is_scalable_vector(dtype) {
+ true => support::scalable_vector(PrimType::float(32).dtype,
vscale_factor(dtype))?,
+ false => with_lanes(PrimType::float(32).dtype, dtype.lanes),
+ };
+ let float32 = PrimType::from_dtype(float32)?;
+ let x_float32: PrimExpr = Cast::new(float32.clone(), x,
loc.clone())?.into();
+ let result = call(Some(&float32.into()), op, vec![x_float32],
loc.clone())?;
+ Ok(Cast::new(PrimType::from_dtype(dtype)?, result, loc)?.into())
+}
+
//
---------------------------------------------------------------------------------------------
// Data type queries and operand checks the operators share
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index d2cc580b1..4ddcb5c05 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -572,17 +572,6 @@ impl For {
}
}
-/// Whether `value` is in the range of the integer type `dtype` (C++
`IntImmValueFits`).
-fn int_imm_value_fits(value: &BigInt, dtype: DLDataType) -> Result<bool> {
- let bits = dtype.bits as u32;
- if dtype.code == support::UINT {
- let limit = BigInt::from(1i64).try_shl(bits as i64)?;
- return Ok(!value.is_negative() && *value < limit);
- }
- let limit = BigInt::from(1i64).try_shl(bits as i64 - 1)?;
- Ok(*value >= -limit.clone() && *value < limit)
-}
-
// tvm-ffi-stubgen(begin): object/ir.While
/// Complete: reflected fields fill [32, 48) exactly.
#[repr(C)]
@@ -900,50 +889,6 @@ impl RegionStmt {
}
}
-/// The parameters the body of a region of `op` on `args` binds, which the
operator's
-/// `FRegionGetBodyParams` gives once its validator accepts the arguments
-/// (C++ `GetRegionBodyParams`).
-fn region_body_params(op: &Op, args: &Array<Expr>, attrs: &DictAttrs) ->
Result<Array<Var>> {
- let name = op.name.as_str();
- let Some(get_body_params) = support::NativeHook::of_op(op,
"FRegionGetBodyParams")? else {
- let message = format!(
- "{name} does not support region construction: FRegionGetBodyParams
is required"
- );
- return Err(support::check_failed(tvm_ffi::VALUE_ERROR,
"IsRegionOp(op)", &message));
- };
- let signature = Call::from_fields(
- UNKNOWN_LOC,
- MissingType::new(),
- op.clone(),
- args.clone(),
- Some(attrs.clone().into()),
- Array::from_vec(vec![]),
- );
- signature.validate()?;
- let params: Array<Var> = get_body_params.call_as(&signature)?;
- let mut definitions = std::collections::HashSet::new();
- for param in params.iter() {
- if param.ty.downcast_ref::<MissingTypeObj>().is_some() {
- let condition = "!param->ty.as<MissingType>().has_value()";
- let message = format!("FRegionGetBodyParams for {name} must return
typed variables");
- return Err(support::check_failed(tvm_ffi::VALUE_ERROR, condition,
&message));
- }
- if !definitions.insert(object_address(param)) {
- let condition = "definitions.insert(param.get()).second";
- let message =
- format!("FRegionGetBodyParams for {name} must return distinct
definitions");
- return Err(support::check_failed(tvm_ffi::VALUE_ERROR, condition,
&message));
- }
- }
- Ok(params)
-}
-
-/// The address of the object `value` refers to, which identifies a definition.
-fn object_address<T: ObjectRefCore>(value: &T) -> usize {
- // The pointer is only compared, never dereferenced.
- ObjectArc::as_ptr(ObjectRefCore::data(value)) as usize
-}
-
// tvm-ffi-stubgen(begin): object/ir.SeqStmt
/// Complete: reflected fields fill [32, 40) exactly.
#[repr(C)]
@@ -1002,42 +947,6 @@ impl SeqStmt {
}
}
-/// `seq` with each nested sequence replaced by its statements and the no-ops
removed.
-fn flatten(seq: Array<Stmt>) -> Array<Stmt> {
- let nested_or_no_op =
- |stmt: &Stmt| stmt.downcast_ref::<SeqStmtObj>().is_some() ||
is_no_op(stmt);
- if !seq.iter().any(nested_or_no_op) {
- return seq;
- }
- fn append(stmt: Stmt, flattened: &mut Vec<Stmt>) {
- match stmt.downcast_ref::<SeqStmtObj>() {
- Some(nested) => nested.seq.iter().for_each(|child|
append(child.clone(), flattened)),
- None if is_no_op(&stmt) => {}
- None => flattened.push(stmt),
- }
- }
- let mut flattened = Vec::new();
- seq.iter().for_each(|stmt| append(stmt.clone(), &mut flattened));
- Array::from_vec(flattened)
-}
-
-/// Whether `stmt` is `Evaluate(0)`, the no-op a sequence drops (C++
`IsSeqStmtNoOp`).
-fn is_no_op(stmt: &Stmt) -> bool {
- stmt.downcast_ref::<EvaluateObj>()
- .and_then(|evaluate| evaluate.value.downcast_ref::<IntImmObj>())
- .is_some_and(|imm| imm.value == BigInt::from(0i64))
-}
-
-/// A statement as a body, as C++ converts a `Stmt` to a `SeqStmt`: a sequence
as it is, any other
-/// statement as the sequence of it at its location (C++ `SeqStmt(stmt)`).
-fn into_body(stmt: Stmt) -> Result<SeqStmt> {
- if stmt.downcast_ref::<SeqStmtObj>().is_some() {
- return Ok(stmt.downcast().expect("a SeqStmt"));
- }
- let loc = stmt.loc.clone();
- Ok(SeqStmt::from_fields(loc, flatten(Array::from_vec(vec![stmt]))))
-}
-
// tvm-ffi-stubgen(begin): object/ir.If
/// Complete: reflected fields fill [32, 56) exactly.
#[repr(C)]
@@ -1326,3 +1235,97 @@ tvm_ffi::unsafe_impl_object_ref_upcast!(Continue =>
Stmt);
// tvm-ffi-stubgen(end)
// `ir.Continue` stores its location as given: the generated `Continue::new`
is the same.
+
+//
---------------------------------------------------------------------------------------------
+// Helpers
+
+/// Whether `value` is in the range of the integer type `dtype` (C++
`IntImmValueFits`).
+fn int_imm_value_fits(value: &BigInt, dtype: DLDataType) -> Result<bool> {
+ let bits = dtype.bits as u32;
+ if dtype.code == support::UINT {
+ let limit = BigInt::from(1i64).try_shl(bits as i64)?;
+ return Ok(!value.is_negative() && *value < limit);
+ }
+ let limit = BigInt::from(1i64).try_shl(bits as i64 - 1)?;
+ Ok(*value >= -limit.clone() && *value < limit)
+}
+
+/// The parameters the body of a region of `op` on `args` binds, which the
operator's
+/// `FRegionGetBodyParams` gives once its validator accepts the arguments
+/// (C++ `GetRegionBodyParams`).
+fn region_body_params(op: &Op, args: &Array<Expr>, attrs: &DictAttrs) ->
Result<Array<Var>> {
+ let name = op.name.as_str();
+ let Some(get_body_params) = support::NativeHook::of_op(op,
"FRegionGetBodyParams")? else {
+ let message = format!(
+ "{name} does not support region construction: FRegionGetBodyParams
is required"
+ );
+ return Err(support::check_failed(tvm_ffi::VALUE_ERROR,
"IsRegionOp(op)", &message));
+ };
+ let signature = Call::from_fields(
+ UNKNOWN_LOC,
+ MissingType::new(),
+ op.clone(),
+ args.clone(),
+ Some(attrs.clone().into()),
+ Array::from_vec(vec![]),
+ );
+ signature.validate()?;
+ let params: Array<Var> = get_body_params.call_as(&signature)?;
+ let mut definitions = std::collections::HashSet::new();
+ for param in params.iter() {
+ if param.ty.downcast_ref::<MissingTypeObj>().is_some() {
+ let condition = "!param->ty.as<MissingType>().has_value()";
+ let message = format!("FRegionGetBodyParams for {name} must return
typed variables");
+ return Err(support::check_failed(tvm_ffi::VALUE_ERROR, condition,
&message));
+ }
+ if !definitions.insert(object_address(param)) {
+ let condition = "definitions.insert(param.get()).second";
+ let message =
+ format!("FRegionGetBodyParams for {name} must return distinct
definitions");
+ return Err(support::check_failed(tvm_ffi::VALUE_ERROR, condition,
&message));
+ }
+ }
+ Ok(params)
+}
+
+/// The address of the object `value` refers to, which identifies a definition.
+fn object_address<T: ObjectRefCore>(value: &T) -> usize {
+ // The pointer is only compared, never dereferenced.
+ ObjectArc::as_ptr(ObjectRefCore::data(value)) as usize
+}
+
+/// `seq` with each nested sequence replaced by its statements and the no-ops
removed.
+fn flatten(seq: Array<Stmt>) -> Array<Stmt> {
+ let nested_or_no_op =
+ |stmt: &Stmt| stmt.downcast_ref::<SeqStmtObj>().is_some() ||
is_no_op(stmt);
+ if !seq.iter().any(nested_or_no_op) {
+ return seq;
+ }
+ fn append(stmt: Stmt, flattened: &mut Vec<Stmt>) {
+ match stmt.downcast_ref::<SeqStmtObj>() {
+ Some(nested) => nested.seq.iter().for_each(|child|
append(child.clone(), flattened)),
+ None if is_no_op(&stmt) => {}
+ None => flattened.push(stmt),
+ }
+ }
+ let mut flattened = Vec::new();
+ seq.iter().for_each(|stmt| append(stmt.clone(), &mut flattened));
+ Array::from_vec(flattened)
+}
+
+/// Whether `stmt` is `Evaluate(0)`, the no-op a sequence drops (C++
`IsSeqStmtNoOp`).
+fn is_no_op(stmt: &Stmt) -> bool {
+ stmt.downcast_ref::<EvaluateObj>()
+ .and_then(|evaluate| evaluate.value.downcast_ref::<IntImmObj>())
+ .is_some_and(|imm| imm.value == BigInt::from(0i64))
+}
+
+/// A statement as a body, as C++ converts a `Stmt` to a `SeqStmt`: a sequence
as it is, any other
+/// statement as the sequence of it at its location (C++ `SeqStmt(stmt)`).
+fn into_body(stmt: Stmt) -> Result<SeqStmt> {
+ if stmt.downcast_ref::<SeqStmtObj>().is_some() {
+ return Ok(stmt.downcast().expect("a SeqStmt"));
+ }
+ let loc = stmt.loc.clone();
+ Ok(SeqStmt::from_fields(loc, flatten(Array::from_vec(vec![stmt]))))
+}
diff --git a/rust/src/ir/support.rs b/rust/src/ir/support.rs
index 1e2e0e928..38d418d34 100644
--- a/rust/src/ir/support.rs
+++ b/rust/src/ir/support.rs
@@ -230,15 +230,6 @@ pub(crate) fn prim_dtype(expr: &(impl AsRef<Expr> +
?Sized)) -> Result<DLDataTyp
}
}
-fn not_a_prim_expr(expr: &Expr) -> Error {
- let message = format!(
- "Expected `ir.PrimExpr` but got `{}[ty={}]`",
- type_key(AnyView::from(expr).type_index()),
- type_key(AnyView::from(&expr.ty).type_index()),
- );
- Error::new(TYPE_ERROR, &message, "")
-}
-
/// The lanes of a dtype as C++ decodes them: negative for a scalable vector.
pub(crate) fn encoded_lanes(dtype: DLDataType) -> i16 {
dtype.lanes as i16
@@ -281,6 +272,15 @@ pub(crate) fn scalable_vector(dtype: DLDataType, factor:
i32) -> Result<DLDataTy
Ok(DLDataType { code: dtype.code, bits: dtype.bits, lanes: (-factor) as
u16 })
}
+fn not_a_prim_expr(expr: &Expr) -> Error {
+ let message = format!(
+ "Expected `ir.PrimExpr` but got `{}[ty={}]`",
+ type_key(AnyView::from(expr).type_index()),
+ type_key(AnyView::from(&expr.ty).type_index()),
+ );
+ Error::new(TYPE_ERROR, &message, "")
+}
+
//
---------------------------------------------------------------------------------------------
// Sequences
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 356e3ac5e..a371a2502 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -160,6 +160,27 @@
//! uv run --no-sync cargo test --manifest-path rust/Cargo.toml
//! ```
+pub mod backend;
+pub mod codegen;
+pub mod ir;
+mod r#macro;
+pub mod runtime;
+pub mod script;
+pub mod sym;
+pub mod target;
+pub mod tirx;
+
+pub use tvm_ffi;
+
+// The errors of the static init blocks of every library built on the crate,
which
+// `import tvm_tirx` raises for `libtvm_tirx_rust`:
`__tvm_ffi_tirx_register_ir`, the TVM FFI
+// function `tirx_register_ir` of the library as `tvm_ffi.load_module` loads
it. The module keeps
+// the symbol, which `#[no_mangle]` exports from wherever it is, out of the
crate's documentation.
+#[doc(hidden)]
+pub mod exports {
+ tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir,
tvm_ffi::run_static_init_blocks);
+}
+
// The crate reaches `libtvm_tirx_compiler` through the FFI registry, so a
linker that drops
// unreferenced libraries (`--as-needed`) would drop the library the build
script links, and its
// registrations would never run. Each TIRx shared library exports a
`<library>_link_anchor`
@@ -201,24 +222,3 @@ impl Drop for ExitOnDrop {
let _ = self.0.call_packed(&[]);
}
}
-
-pub mod backend;
-pub mod codegen;
-pub mod ir;
-mod r#macro;
-pub mod runtime;
-pub mod script;
-pub mod sym;
-pub mod target;
-pub mod tirx;
-
-pub use tvm_ffi;
-
-// The errors of the static init blocks of every library built on the crate,
which
-// `import tvm_tirx` raises for `libtvm_tirx_rust`:
`__tvm_ffi_tirx_register_ir`, the TVM FFI
-// function `tirx_register_ir` of the library as `tvm_ffi.load_module` loads
it. The module keeps
-// the symbol, which `#[no_mangle]` exports from wherever it is, out of the
crate's documentation.
-#[doc(hidden)]
-pub mod exports {
- tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir,
tvm_ffi::run_static_init_blocks);
-}
diff --git a/rust/src/sym/constraint.rs b/rust/src/sym/constraint.rs
index 687c55121..2b659af70 100644
--- a/rust/src/sym/constraint.rs
+++ b/rust/src/sym/constraint.rs
@@ -286,6 +286,16 @@ pub struct ConstraintContextObj {
constraints: UnsafeCell<Array<Constraint>>,
}
+// These verify the Rust mirror against the matching x86_64 header layout, not
+// an arbitrary loaded binary's hidden fields. Stub generation cannot certify
them.
+const _: () = {
+ assert!(size_of::<ConstraintContextObj>() == 48);
+ assert!(align_of::<ConstraintContextObj>() == 8);
+ assert!(offset_of!(ConstraintContextObj, parent) == 24);
+ assert!(offset_of!(ConstraintContextObj, parent_prefix_size) == 32);
+ assert!(offset_of!(ConstraintContextObj, constraints) == 40);
+};
+
/// A context of constraints: the premises in effect at a point of a
traversal, ancestors' first.
///
/// Clones share the context's identity and later appends; use
@@ -296,46 +306,6 @@ pub struct ConstraintContext {
base: ObjectArc<ConstraintContextObj>,
}
-impl ConstraintContext {
- /// Create an empty root context, with the complete native-compatible
layout (C++
- /// `ConstraintContext()`).
- pub fn new() -> Self {
- Self::allocate(None, 0)
- }
-
- /// Create an empty child that captures the first `parent_prefix_size`
local entries of
- /// `parent`.
- ///
- /// Earlier ancestor prefixes are inherited unchanged; later appends to
the parent are excluded.
- ///
- /// # Errors
- ///
- /// Returns an error if `parent_prefix_size` is negative or larger than
the parent's local size.
- pub fn with_parent_prefix_size(parent: Self, parent_prefix_size: i64) ->
Result<Self> {
- if parent_prefix_size < 0 || parent_prefix_size > parent.local_size() {
- return Err(Error::new(
- VALUE_ERROR,
- "Constraint context parent local prefix is out of bounds",
- "",
- ));
- }
- Ok(Self::allocate(Some(parent.base), parent_prefix_size))
- }
-
- fn allocate(parent: Option<ObjectArc<ConstraintContextObj>>, prefix: i64)
-> Self {
- Self {
- base: ObjectArc::new(ConstraintContextObj {
- base: Object::new(),
- parent,
- parent_prefix_size: prefix,
- // Start with an empty Rust array. Native Append owns all
growth;
- // never install a nonempty Rust allocation in this mutable
field.
- constraints: UnsafeCell::new(Array::default()),
- }),
- }
- }
-}
-
impl ConstraintContextObj {
/// Append `constraint` to the local entries
(`sym.ConstraintContextAppend`).
///
@@ -397,12 +367,42 @@ impl ConstraintContextObj {
}
}
-// These verify the Rust mirror against the matching x86_64 header layout, not
-// an arbitrary loaded binary's hidden fields. Stub generation cannot certify
them.
-const _: () = {
- assert!(size_of::<ConstraintContextObj>() == 48);
- assert!(align_of::<ConstraintContextObj>() == 8);
- assert!(offset_of!(ConstraintContextObj, parent) == 24);
- assert!(offset_of!(ConstraintContextObj, parent_prefix_size) == 32);
- assert!(offset_of!(ConstraintContextObj, constraints) == 40);
-};
+impl ConstraintContext {
+ /// Create an empty root context, with the complete native-compatible
layout (C++
+ /// `ConstraintContext()`).
+ pub fn new() -> Self {
+ Self::allocate(None, 0)
+ }
+
+ /// Create an empty child that captures the first `parent_prefix_size`
local entries of
+ /// `parent`.
+ ///
+ /// Earlier ancestor prefixes are inherited unchanged; later appends to
the parent are excluded.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if `parent_prefix_size` is negative or larger than
the parent's local size.
+ pub fn with_parent_prefix_size(parent: Self, parent_prefix_size: i64) ->
Result<Self> {
+ if parent_prefix_size < 0 || parent_prefix_size > parent.local_size() {
+ return Err(Error::new(
+ VALUE_ERROR,
+ "Constraint context parent local prefix is out of bounds",
+ "",
+ ));
+ }
+ Ok(Self::allocate(Some(parent.base), parent_prefix_size))
+ }
+
+ fn allocate(parent: Option<ObjectArc<ConstraintContextObj>>, prefix: i64)
-> Self {
+ Self {
+ base: ObjectArc::new(ConstraintContextObj {
+ base: Object::new(),
+ parent,
+ parent_prefix_size: prefix,
+ // Start with an empty Rust array. Native Append owns all
growth;
+ // never install a nonempty Rust allocation in this mutable
field.
+ constraints: UnsafeCell::new(Array::default()),
+ }),
+ }
+ }
+}
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index 24a890618..6a236ebb2 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -198,18 +198,6 @@ pub fn logaddexp_at(
log(add(exp(a)?, exp(b)?)?)
}
-/// Whether `dtype` is a scalar 64-bit integer, which a pointer reinterprets
as.
-fn is_scalar_int64(dtype: DLDataType) -> bool {
- support::is_scalar(dtype)
- && dtype.bits == 64
- && matches!(dtype.code, support::INT | support::UINT)
-}
-
-/// C++ `PrimType::StorageBytes()`.
-fn storage_bytes(dtype: DLDataType) -> Result<u64> {
- Ok((u64::from(dtype.bits) * support::lanes(dtype)? as u64).div_ceil(8))
-}
-
/// Reinterpret `value` bit for bit as the primitive or pointer type `ty`
/// (C++ `tirx::reinterpret(target_ty, value, loc)`, `tirx.reinterpret`).
///
@@ -375,3 +363,15 @@ pub fn gpu_thread_return_at(loc: impl Into<Location>) ->
Result<PrimExpr> {
let loc: Location = loc.into();
tvm_ffi::cached_global_func!("tirx.gpu_thread_return").call_tuple((&loc,))?.try_into()
}
+
+/// Whether `dtype` is a scalar 64-bit integer, which a pointer reinterprets
as.
+fn is_scalar_int64(dtype: DLDataType) -> bool {
+ support::is_scalar(dtype)
+ && dtype.bits == 64
+ && matches!(dtype.code, support::INT | support::UINT)
+}
+
+/// C++ `PrimType::StorageBytes()`.
+fn storage_bytes(dtype: DLDataType) -> Result<u64> {
+ Ok((u64::from(dtype.bits) * support::lanes(dtype)? as u64).div_ceil(8))
+}
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index 994682266..e076249ef 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -163,24 +163,6 @@ unsafe impl tvm_ffi::object::ImmutableObject for
TensorTypeObj {}
tvm_ffi::unsafe_impl_object_ref_upcast!(TensorType => Type);
// tvm-ffi-stubgen(end)
-/// The runtime's allocation alignment (C++ `runtime::kAllocAlignment`, which
a build may set),
-/// read from the default the reflection registers for
`TensorType::data_alignment`.
-fn alloc_alignment() -> Result<i32> {
- static ALLOC_ALIGNMENT: std::sync::OnceLock<i32> =
std::sync::OnceLock::new();
- if let Some(alignment) = ALLOC_ALIGNMENT.get() {
- return Ok(*alignment);
- }
- let Some(default) = support::field_default(TensorTypeObj::type_index(),
"data_alignment")
- else {
- let message = "tirx.TensorType registers no default data_alignment";
- return Err(tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, message, ""));
- };
- let alignment = i32::try_from(i64::try_from(default)?).map_err(|_| {
- tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, "the default
data_alignment exceeds i32", "")
- })?;
- Ok(*ALLOC_ALIGNMENT.get_or_init(|| alignment))
-}
-
impl TensorType {
/// Create the type of a tensor of `shape` and `dtype` in `scope`, aligned
to `data_alignment`
/// bytes (the runtime's allocation alignment when 0 or less), with
`layout` if given (C++
@@ -220,6 +202,20 @@ impl TensorType {
}
}
+/// A signature descriptor ([`sig::ArgType`]) of a `Var`, not a subclass,
whose type is a
+/// [`TensorType`] (C++ `tirx::TensorVar`).
+pub struct TensorVar;
+
+impl sig::ArgType for TensorVar {
+ fn check(value: &TVMFFIAny) -> bool {
+ sig::var_of_type::<TensorTypeObj>(value)
+ }
+
+ fn type_str() -> std::string::String {
+ crate::ir::expr::VarObj::TYPE_KEY.into()
+ }
+}
+
/// Return whether the tensor variable `tensor` is a scalar: one element, the
default alignment and
/// a trivial layout; `alloc_or_decl` is true for `alloc_scalar`, false for
`decl_scalar`
/// (`tirx.TensorIsScalar`).
@@ -233,16 +229,20 @@ pub fn tensor_is_scalar(tensor: impl IntoObjectRef<Var>,
alloc_or_decl: bool) ->
.try_into()
}
-/// A signature descriptor ([`sig::ArgType`]) of a `Var`, not a subclass,
whose type is a
-/// [`TensorType`] (C++ `tirx::TensorVar`).
-pub struct TensorVar;
-
-impl sig::ArgType for TensorVar {
- fn check(value: &TVMFFIAny) -> bool {
- sig::var_of_type::<TensorTypeObj>(value)
- }
-
- fn type_str() -> std::string::String {
- crate::ir::expr::VarObj::TYPE_KEY.into()
+/// The runtime's allocation alignment (C++ `runtime::kAllocAlignment`, which
a build may set),
+/// read from the default the reflection registers for
`TensorType::data_alignment`.
+fn alloc_alignment() -> Result<i32> {
+ static ALLOC_ALIGNMENT: std::sync::OnceLock<i32> =
std::sync::OnceLock::new();
+ if let Some(alignment) = ALLOC_ALIGNMENT.get() {
+ return Ok(*alignment);
}
+ let Some(default) = support::field_default(TensorTypeObj::type_index(),
"data_alignment")
+ else {
+ let message = "tirx.TensorType registers no default data_alignment";
+ return Err(tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, message, ""));
+ };
+ let alignment = i32::try_from(i64::try_from(default)?).map_err(|_| {
+ tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, "the default
data_alignment exceeds i32", "")
+ })?;
+ Ok(*ALLOC_ALIGNMENT.get_or_init(|| alignment))
}
