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 f17fc6ecb [Rust] Adopt tvm-ffi EnumRef, NativeFunction, Display,
type_key and field_default (#222)
f17fc6ecb is described below
commit f17fc6ecb29aa7ba1a17e54e804932952ed9f74d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Nov 7 08:16:26 2026 -0500
[Rust] Adopt tvm-ffi EnumRef, NativeFunction, Display, type_key and
field_default (#222)
Adopt the tvm-ffi and drop the TIRx helpers they replace.
- Bump `3rdparty/tvm-ffi` to apache/tvm-ffi@f3096d8d and regenerate the
bindings: `MemoryScope`, `ThreadScope` and `Axis` mirror `ffi.Enum` as
`EnumObj` and implement `EnumRef`; reflection traits come from
`tvm_ffi::reflection`.
- Enums: values are looked up with `EnumRef::get` and registered with
`EnumDef`. `MemoryScope::register` is removed
(`EnumDef::<MemoryScope>::new`), `ThreadScope::get` is now the `EnumRef`
lookup, and `Axis::get`, which registers an unknown name, is renamed
`Axis::get_or_register`.
- Native hooks: `reflection::NativeFunction` builds `OpDef`'s hooks from
Rust functions and calls the hooks the registries hold, replacing
`support::NativeHook`, `support::call_with_refs` and `OpDef`'s C ABI
adapters. The hook setters take `Copy` functions, and a panic in a hook
now aborts instead of becoming an error.
- Error messages print values with `Display` and name types with
`reflection::type_key`, replacing `support::repr` and
`support::type_key`.
- The TensorLoad and TensorStore type attributes are looked up once
through a `LazyLock`, replacing `support::TypeAttr`.
- The default `TensorType` alignment is read with
`reflection::field_default`, replacing `support::field_default`.
- Drop the tensor scalar check (`TensorTypeObj::IsScalar`,
`TensorVar::IsScalar`, `tirx.TensorIsScalar`, Python `is_scalar`, Rust
`tensor_is_scalar`) and inline it into ForceNarrowIndexToInt32, its one
caller.
---
3rdparty/tvm-ffi | 2 +-
include/tvm/tirx/expr.h | 7 -
include/tvm/tirx/type.h | 3 -
python/tvm_tirx/tirx/tensor.py | 17 --
rust/src/ir/attrs.rs | 4 +-
rust/src/ir/expr.rs | 52 +++---
rust/src/ir/function.rs | 4 +-
rust/src/ir/global_info.rs | 4 +-
rust/src/ir/location.rs | 4 +-
rust/src/ir/op_def.rs | 112 +++----------
rust/src/ir/prim/expr.rs | 3 +-
rust/src/ir/prim/op.rs | 34 ++--
rust/src/ir/prim/vector_expr.rs | 4 +-
rust/src/ir/stmt.rs | 65 ++++----
rust/src/ir/support.rs | 202 ++----------------------
rust/src/ir/transform/core.rs | 2 +-
rust/src/ir/type.rs | 4 +-
rust/src/ir/type_def.rs | 7 +-
rust/src/runtime.rs | 38 ++---
rust/src/sym/constraint.rs | 4 +-
rust/src/sym/int_set.rs | 2 +-
rust/src/sym/iter_affine_map.rs | 2 +-
rust/src/target.rs | 2 +-
rust/src/tirx/layout.rs | 33 ++--
rust/src/tirx/op.rs | 35 ++--
rust/src/tirx/structural.rs | 2 +-
rust/src/tirx/tensor.rs | 36 +----
rust/src/tirx/thread_scope.rs | 38 ++---
rust/src/tirx/tile/dispatch.rs | 4 +-
rust/src/tirx/tile/instruction.rs | 2 +-
rust/src/tirx/transform/function_pass.rs | 2 +-
rust/tests/ir/op_def.rs | 31 +---
rust/tests/ir/type.rs | 5 +-
rust/tests/runtime/main.rs | 5 +-
rust/tests/tirx/layout.rs | 4 +-
rust/tests/tirx/main.rs | 5 +-
rust/tests/tirx/tensor.rs | 6 +-
rust/tests/tirx/thread_scope.rs | 1 +
rust/tests/tirx/tile/dispatch.rs | 2 +-
src/tirx/ir/tensor.cc | 7 -
src/tirx/ir/type.cc | 8 -
src/tirx/transform/force_narrow_index_to_i32.cc | 12 +-
42 files changed, 241 insertions(+), 575 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index 7bedac88a..f3096d8db 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit 7bedac88ac4f2737d089244bae4d7be9d3599140
+Subproject commit f3096d8db5341b1af3786e360ef07d06f8631966
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index 1891960a0..5f192ce8b 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -71,13 +71,6 @@ class TensorVar : public Var {
/*! \brief Return the storage scope associated with this tensor. */
runtime::MemoryScope scope() const { return ty()->scope; }
- /*!
- * \brief Return true if the tensor is a scalar.
- * \param alloc_or_decl Whether to consider alloc_scalar and decl_scalar as
scalar. True for
- * alloc_scalar, False for decl_scalar.
- */
- TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
-
explicit TensorVar(ffi::ObjectPtr<VarObj> node) :
TensorVar(Var(std::move(node))) {}
explicit TensorVar(ffi::UnsafeInit tag) : Var(tag) {}
TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TensorVar);
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index 52ba65f5f..22ffc96b6 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -107,9 +107,6 @@ class TensorTypeObj : public TypeObj {
/*! \return type of the physical pointer projected by tensor_data_ptr. */
PtrType DataPtrType() const { return PtrType(dtype, scope); }
- /*! \brief Whether this type supports scalar tensor syntax. */
- TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
-
/*! \brief Return the constant element count, or nullopt for symbolic
extents or overflow. */
TVM_TIRX_DLL ffi::Optional<int64_t> ConstantAllocationSize() const;
diff --git a/python/tvm_tirx/tirx/tensor.py b/python/tvm_tirx/tirx/tensor.py
index 28db73cf5..ce238c1d8 100644
--- a/python/tvm_tirx/tirx/tensor.py
+++ b/python/tvm_tirx/tirx/tensor.py
@@ -149,7 +149,6 @@ class TensorType(Type):
at(loc, self)
__expr_methods__ = (
- "is_scalar",
"ptr_to",
"view",
"permute",
@@ -201,22 +200,6 @@ class TensorType(Type):
_check_tensor_receiver(self, expr)
return self.scope
- def is_scalar(self, expr, alloc_or_decl=True):
- """Check if the tensor is a scalar.
-
- Parameters
- ----------
- alloc_or_decl : bool, optional
- Whether to consider alloc_scalar and decl_scalar as scalar. True
for alloc_scalar,
- False for decl_scalar.
-
- Returns
- -------
- bool: True if the tensor is a scalar, False otherwise.
- """
- _check_tensor_receiver(self, expr)
- return _ffi_api.TensorIsScalar(expr, alloc_or_decl)
-
def ptr_to(self, expr, indices):
"""Get the pointer to the tensor at the given indices (logical
indices).
diff --git a/rust/src/ir/attrs.rs b/rust/src/ir/attrs.rs
index a8531cdbc..eedf5aaba 100644
--- a/rust/src/ir/attrs.rs
+++ b/rust/src/ir/attrs.rs
@@ -31,9 +31,9 @@ use tvm_ffi::Any;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
-use tvm_ffi::UnsafeInit;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
// tvm-ffi-stubgen(end)
/// The base of every attribute object, the metadata attached to a call, a
function or a module.
diff --git a/rust/src/ir/expr.rs b/rust/src/ir/expr.rs
index f697bb813..3494be7f8 100644
--- a/rust/src/ir/expr.rs
+++ b/rust/src/ir/expr.rs
@@ -59,6 +59,9 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
+// tvm-ffi-stubgen(import-object): tvm_ffi::reflection::NativeFunction
+// tvm-ffi-stubgen(import-object): tvm_ffi::reflection::TypeAttrColumn
+// tvm-ffi-stubgen(import-object): std::sync::LazyLock
// tvm-ffi-stubgen(import-object): tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex
// tvm-ffi-stubgen(import-object): crate::ir::location::UNKNOWN_LOC
// tvm-ffi-stubgen(import-object): crate::ir::op::OpObj
@@ -98,6 +101,7 @@ use crate::ir::support;
use std::borrow::Borrow;
use std::ops::Deref;
use std::ptr::addr_of_mut;
+use std::sync::LazyLock;
use tvm_ffi::Any;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
@@ -108,10 +112,12 @@ use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
-use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
-use tvm_ffi::UnsafeInit;
use tvm_ffi::object::IntoObjectRef;
+use tvm_ffi::reflection::NativeFunction;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::TypeAttrColumn;
+use tvm_ffi::reflection::UnsafeInit;
use tvm_ffi::reflection::drop_if_defined;
use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
// tvm-ffi-stubgen(end)
@@ -595,14 +601,15 @@ fn scalar_int(ty: &PrimType) -> bool {
/// What a typed expression's conversion error names a value it rejects: an
expression with its
/// type, `ir.StringImm[ty=ir.StringType]`, or the type key of another value.
fn typed_expr_mismatch(value: AnyView<'_>) -> std::string::String {
+ let key = |index: i32|
tvm_ffi::reflection::type_key(index).unwrap_or_default();
let Some(expr) = value.downcast_ref::<ExprObj>() else {
- return support::type_key(value.type_index());
+ return key(value.type_index()).to_string();
};
let ty = match expr.ty.downcast_ref::<PrimTypeObj>() {
Some(prim) => support::dtype_name(prim.dtype),
- None => support::type_key(AnyView::from(&expr.ty).type_index()),
+ None => key(AnyView::from(&expr.ty).type_index()).to_string(),
};
- format!("{}[ty={ty}]", support::type_key(value.type_index()))
+ format!("{}[ty={ty}]", key(value.type_index()))
}
//
---------------------------------------------------------------------------------------------
@@ -982,7 +989,7 @@ impl TupleGetItem {
if index < 0 {
let message = format!(
"Index out of bounds: Tuple {} cannot be accessed with
negative index {index}",
- support::repr(&tuple_value)
+ tuple_value
);
return Err(support::compare_failed(
tvm_ffi::INDEX_ERROR,
@@ -999,7 +1006,7 @@ impl TupleGetItem {
let message = format!(
"Index out of bounds: Tuple {} is of size {size}, and
cannot be accessed \
with index {index}",
- support::repr(&tuple_value)
+ tuple_value
);
let condition = "index < tuple_type->fields.size()";
let error = support::compare_failed(
@@ -1099,25 +1106,30 @@ impl TensorLoad {
indices: impl IntoIterator<Item = impl IntoExpr<PrimExpr>>,
loc: impl Into<Location>,
) -> Result<Self> {
- static INFER_TYPE: support::TypeAttr =
- support::TypeAttr::new("__tvm_tensor_load_infer_type__");
+ // The native library registers the column when it loads, so one
lookup suffices.
+ static INFER_TYPE: LazyLock<Option<TypeAttrColumn>> =
+ LazyLock::new(||
TypeAttrColumn::new("__tvm_tensor_load_infer_type__"));
let (source, loc) = (source.into_owned_expr()?, loc.into());
let indices = indices
.into_iter()
.map(IntoExpr::<PrimExpr>::into_owned_expr)
.collect::<Result<Array<PrimExpr>>>()?;
let ty_index = AnyView::from(&source.ty).type_index();
- let Some(infer) = INFER_TYPE.get(ty_index) else {
- let message =
- format!("Type {} does not support TensorLoad",
support::type_key(ty_index));
+ let Some(infer) = INFER_TYPE.and_then(|column| column.get(ty_index))
else {
+ let message = format!(
+ "Type {} does not support TensorLoad",
+ tvm_ffi::reflection::type_key(ty_index).unwrap_or_default()
+ );
return Err(support::check_failed(tvm_ffi::TYPE_ERROR,
"infer.has_value()", &message));
};
// The hook types a load that has no type yet. Such a load is no
`TensorLoad`, which is
- // typed by a `PrimType`, so it is made as an `Expr`.
+ // typed by a `PrimType`, so only its object is made.
let base_ = ExprObj::new(loc.clone(), MissingType::new().into());
let untyped = TensorLoadObj { base_, source: source.clone(), indices:
indices.clone() };
- let untyped = Expr::from(Self { data_: ObjectArc::new(untyped) });
- let ty: PrimType =
support::NativeHook::from_any(infer)?.call_as(&untyped)?;
+ // SAFETY: the column holds `PrimType(const TensorLoadObj*)` hooks.
+ let infer: NativeFunction<(&TensorLoadObj,), PrimType> =
+ unsafe { NativeFunction::from_any(infer)? };
+ let ty = infer.call((&*ObjectArc::new(untyped),))?;
Ok(Self::from_fields(loc, ty, source, indices))
}
}
@@ -1801,7 +1813,10 @@ impl Call {
if op.validator().type_index() == TVMFFITypeIndex::kTVMFFINone as i32 {
return Ok(());
}
- support::NativeHook::from_any(op.validator().clone())?.call_void(self)
+ // SAFETY: an operator validator is a `void(const CallObj*)` hook.
+ let validate: NativeFunction<(&CallObj,)> =
+ unsafe { NativeFunction::from_any(op.validator().clone())? };
+ validate.call((&**self,))
}
/// Return the result type the operator infers from the current inputs,
without changing the
@@ -1821,8 +1836,9 @@ impl Call {
if let Some(fixed) = support::op_attr(op, "TFixedReturnType")? {
return fixed.try_into();
}
- match support::NativeHook::of_op(op, "FInferType")? {
- Some(infer) => infer.call_as(self),
+ // SAFETY: `FInferType` holds `Type(const CallObj*)` hooks.
+ match unsafe { support::op_hook::<(&CallObj,), Type>(op,
"FInferType")? } {
+ Some(infer) => infer.call((&**self,)),
None => Ok(MissingType::new().into()),
}
}
diff --git a/rust/src/ir/function.rs b/rust/src/ir/function.rs
index 4e191fa7c..bf4c79e89 100644
--- a/rust/src/ir/function.rs
+++ b/rust/src/ir/function.rs
@@ -43,9 +43,9 @@ use tvm_ffi::Map;
use tvm_ffi::ObjectArc;
use tvm_ffi::RValueRef;
use tvm_ffi::Result;
-use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
-use tvm_ffi::UnsafeInit;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
use tvm_ffi::reflection::drop_if_defined;
// tvm-ffi-stubgen(end)
diff --git a/rust/src/ir/global_info.rs b/rust/src/ir/global_info.rs
index bd30b9032..85d10c816 100644
--- a/rust/src/ir/global_info.rs
+++ b/rust/src/ir/global_info.rs
@@ -28,8 +28,8 @@
use std::ptr::addr_of_mut;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ShallowCopy;
-use tvm_ffi::UnsafeInit;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
// tvm-ffi-stubgen(end)
/// The base of the module-level information an [`IRModule`](super::IRModule)
keeps by name in
diff --git a/rust/src/ir/location.rs b/rust/src/ir/location.rs
index 48d0b8b12..2ef082ca2 100644
--- a/rust/src/ir/location.rs
+++ b/rust/src/ir/location.rs
@@ -39,9 +39,9 @@ use std::ptr::addr_of_mut;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::Result;
-use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
-use tvm_ffi::UnsafeInit;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
// tvm-ffi-stubgen(end)
/// Where an IR node comes from, for diagnostics: an [`UnknownLoc`], a
[`SourceLoc`] or a
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index e649e189a..544711ce5 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -22,11 +22,10 @@
//! [`AttrKey`], such as the shared ones in [`op_attr`].
use std::any::TypeId;
-use std::ffi::c_void;
use std::marker::PhantomData;
-use std::panic::{catch_unwind, AssertUnwindSafe};
-use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex};
+use tvm_ffi::reflection::NativeFunction;
+use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
use tvm_ffi::{Any, AnyView, Array, Error, ObjectRefCore, Result, String,
TYPE_ERROR, VALUE_ERROR};
use super::attrs::Attrs;
@@ -34,7 +33,6 @@ use super::expr::{CallObj, Expr, Var};
use super::op::{Op, OpObj};
use super::r#type::Type;
use super::stmt::RegionStmtObj;
-use super::support;
use sig::{Checks, Specs};
/// A chain of steps that registers a new operator.
@@ -44,6 +42,9 @@ use sig::{Checks, Specs};
/// returns it. Register an operator before building IR with it, as a
/// [`tvm_ffi::static_init_block!`] does.
///
+/// A hook, such as the one [`OpDef::set_infer_type`] sets, reports failure
through its `Result`; a
+/// panic in it aborts the process, since it would unwind into native code.
+///
/// # Examples
///
/// ```
@@ -108,8 +109,8 @@ impl OpDef {
let args: Vec<AnyView> = lists.iter().map(AnyView::from).collect();
let def = self.invoke("_set_signature", &args).then(|op| {
if !op.validator_is_custom() {
- let hook: Hook = validate_signature::<S::Checks>;
- op.write_validator(Any::from(hook as *mut c_void), false);
+ let hook: NativeFunction<(&CallObj,)> =
sig::validate::<S::Checks>.into();
+ op.write_validator(hook.into(), false);
}
Ok(())
});
@@ -168,11 +169,10 @@ impl OpDef {
/// precedence. A second one fails the chain with a `ValueError`.
pub fn set_infer_type<F>(self, infer: F) -> Self
where
- F: Fn(&CallObj) -> Result<Type> + 'static,
+ F: Fn(&CallObj) -> Result<Type> + Copy + 'static,
{
- let _ = infer;
- let hook: Hook = infer_type_abi::<F>;
- self.attr("FInferType", Any::from(hook as *mut c_void), false)
+ let hook: NativeFunction<(&CallObj,), Type> = infer.into();
+ self.attr("FInferType", hook.into(), false)
}
/// Set a custom validator of the operator, `validate`, a function or a
closure that captures
@@ -182,16 +182,15 @@ impl OpDef {
/// one fails the chain with a `ValueError`.
pub fn set_validator<F>(self, validate: F) -> Self
where
- F: Fn(&CallObj) -> Result<()> + 'static,
+ F: Fn(&CallObj) -> Result<()> + Copy + 'static,
{
- let _ = validate;
+ let hook: NativeFunction<(&CallObj,)> = validate.into();
self.then(|op| {
if op.validator_is_custom() {
let message = format!("Validator of {} is already registered",
op.name.as_str());
return Err(Error::new(VALUE_ERROR, &message, ""));
}
- let hook: Hook = validate_abi::<F>;
- op.write_validator(Any::from(hook as *mut c_void), true);
+ op.write_validator(hook.into(), true);
Ok(())
})
}
@@ -204,11 +203,10 @@ impl OpDef {
/// fails the chain with a `ValueError`.
pub fn set_region_body_params<F>(self, body_params: F) -> Self
where
- F: Fn(&CallObj) -> Result<Array<Var>> + 'static,
+ F: Fn(&CallObj) -> Result<Array<Var>> + Copy + 'static,
{
- let _ = body_params;
- let hook: Hook = region_body_params_abi::<F>;
- self.attr("FRegionGetBodyParams", Any::from(hook as *mut c_void),
false)
+ let hook: NativeFunction<(&CallObj,), Array<Var>> = body_params.into();
+ self.attr("FRegionGetBodyParams", hook.into(), false)
}
/// Set the check of the operator's regions, its `FRegionValidate`, to
`validate`, a function or
@@ -218,11 +216,10 @@ impl OpDef {
/// the chain with a `ValueError`.
pub fn set_region_validate<F>(self, validate: F) -> Self
where
- F: Fn(&RegionStmtObj) -> Result<()> + 'static,
+ F: Fn(&RegionStmtObj) -> Result<()> + Copy + 'static,
{
- let _ = validate;
- let hook: RegionHook = region_validate_abi::<F>;
- self.attr("FRegionValidate", Any::from(hook as *mut c_void), false)
+ let hook: NativeFunction<(&RegionStmtObj,)> = validate.into();
+ self.attr("FRegionValidate", hook.into(), false)
}
fn attr(self, key: &str, value: Any, replace: bool) -> Self {
@@ -367,7 +364,7 @@ pub mod sig {
///
/// By default, the description is the type key of `value`.
fn mismatch(value: AnyView<'_>) -> std::string::String {
- support::type_key(value.type_index())
+
tvm_ffi::reflection::type_key(value.type_index()).unwrap_or_default().to_string()
}
}
@@ -584,7 +581,11 @@ pub mod sig {
let actual = match &call.attrs {
Some(attrs) if
attrs.downcast_ref::<T::ContainerType>().is_some() => return Ok(()),
Some(attrs) => {
- format!("`{}`",
support::type_key(AnyView::from(attrs).type_index()))
+ format!(
+ "`{}`",
+
tvm_ffi::reflection::type_key(AnyView::from(attrs).type_index())
+ .unwrap_or_default()
+ )
}
None => "None".into(),
};
@@ -804,66 +805,3 @@ fn doc_text(doc: &str) -> std::string::String {
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/expr.rs b/rust/src/ir/prim/expr.rs
index 10393cd63..d6f4249ca 100644
--- a/rust/src/ir/prim/expr.rs
+++ b/rust/src/ir/prim/expr.rs
@@ -2210,7 +2210,8 @@ impl Let {
let Some(var_ty) = var.ty.downcast_ref::<PrimTypeObj>() else {
let message = format!(
"Cannot treat type `{}` as type `ir.PrimType`",
- support::type_key(AnyView::from(&var.ty).type_index())
+
tvm_ffi::reflection::type_key(AnyView::from(&var.ty).type_index())
+ .unwrap_or_default()
);
return Err(tvm_ffi::Error::new(tvm_ffi::TYPE_ERROR, &message, ""));
};
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 967011270..a896f4d92 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -849,7 +849,7 @@ pub(crate) fn make_const_scalar(
};
return Ok(FloatImm::new(ty, value, loc.clone())?.into());
}
- let message = format!("cannot make const for type {}", support::repr(&ty));
+ let message = format!("cannot make const for type {ty}");
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
@@ -1529,8 +1529,7 @@ pub fn logical_not_at(a: impl IntoExpr<PrimExpr>, loc:
impl Into<Location>) -> R
check_code(&a, "a", &BOOLEAN, BOOL_CODES, || {
format!(
"Expected boolean argument for ! operator (logical NOT), but
received {} of type {}",
- support::repr(&a),
- support::repr(&a.ty)
+ a, a.ty
)
})?;
match const_fold::not(&a)? {
@@ -1711,8 +1710,7 @@ pub fn bitwise_not_at(a: impl IntoExpr<PrimExpr>, loc:
impl Into<Location>) -> R
format!(
"Expected integer or boolean argument for ~ operator (bitwise
NOT), but received {} \
of type {}",
- support::repr(&a),
- support::repr(&a.ty)
+ a, a.ty
)
})?;
Ok(BitwiseNot::new(a, loc)?.into())
@@ -1738,7 +1736,7 @@ fn fold_shift(
let message = format!(
"Shift amount must be non-negative and less than {} for type
{}",
ty.bits(),
- support::repr(ty)
+ ty
);
return Err(support::check_failed(INTERNAL_ERROR, condition,
&message));
}
@@ -1864,7 +1862,7 @@ pub fn max_value_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result
_ => match small_float_bound(code) {
Some((max, _)) => float(max),
None => {
- let message = format!("Cannot decide max_value for type{}",
support::repr(&ty));
+ let message = format!("Cannot decide max_value for type{ty}");
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
},
@@ -1891,7 +1889,7 @@ pub fn min_value_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result
Some((_, true)) => float(0.0),
Some((max, false)) => float(-max),
None => {
- let message = format!("Cannot decide min_value for type{}",
support::repr(&ty));
+ let message = format!("Cannot decide min_value for type{ty}");
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
},
@@ -1906,7 +1904,7 @@ pub fn infinity_at(ty: impl Into<PrimType>, loc: impl
Into<Location>) -> Result<
(support::FLOAT, 64) => Ok(FloatImm::new(ty, f64::INFINITY,
loc)?.into()),
(support::FLOAT, 32 | 16) => Ok(FloatImm::new(ty,
f64::from(f32::INFINITY), loc)?.into()),
_ => {
- let message = format!("Cannot decide infinity for type {}",
support::repr(&ty));
+ let message = format!("Cannot decide infinity for type {ty}");
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
}
@@ -2004,7 +2002,7 @@ fn absolute(x: PrimExpr, loc: Location) ->
Result<PrimExpr> {
_ => {
let message = format!(
"Data type {} not supported for absolute op. Skipping absolute
op...",
- support::repr(&x.ty)
+ x.ty
);
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
@@ -2042,10 +2040,8 @@ fn infinite(x: PrimExpr, loc: Location) ->
Result<PrimExpr> {
logical_and(at_infinity, not_nan)
}
_ => {
- let message = format!(
- "Data type {} not supported for finiteness ops. Skipping
it...",
- support::repr(&x.ty)
- );
+ let message =
+ format!("Data type {} not supported for finiteness ops.
Skipping it...", x.ty);
Err(Error::new(INTERNAL_ERROR, &message, ""))
}
}
@@ -2211,8 +2207,7 @@ fn unary_intrinsic(
if float_only && !ty.matches_code(&[DLDataTypeCode::kDLFloat]) &&
!is_bfloat16(&ty) {
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));
+ let message = format!("{op} only supports floating-point inputs, but
got {}", x.ty);
return Err(support::check_failed(TYPE_ERROR, condition, &message));
}
if !is_bfloat16(&ty) {
@@ -2313,8 +2308,8 @@ const BOOL_CODES: &str = "DLDataTypeCode::kDLBool";
/// `TVM_FFI_ICHECK(x.ty().MatchesCode(kDLInt, kDLUInt)) << x` of the integer
divisions, on both
/// operands.
fn check_integers(a: &PrimExpr, b: &PrimExpr) -> Result<()> {
- check_code(a, "a", &INTEGER, INTEGER_CODES, || support::repr(a))?;
- check_code(b, "b", &INTEGER, INTEGER_CODES, || support::repr(b))
+ check_code(a, "a", &INTEGER, INTEGER_CODES, || a.to_string())?;
+ check_code(b, "b", &INTEGER, INTEGER_CODES, || b.to_string())
}
/// The operand checks of an operator `symbol` (`what` names it) that takes
`kind` arguments:
@@ -2332,8 +2327,7 @@ fn check_operands(
format!(
"Expected {kind} argument as {which} of {symbol} operator
({what}), but received \
{} of type {}",
- support::repr(operand),
- support::repr(&operand.ty)
+ operand, operand.ty
)
})?;
}
diff --git a/rust/src/ir/prim/vector_expr.rs b/rust/src/ir/prim/vector_expr.rs
index 1339ddbaf..500f9e926 100644
--- a/rust/src/ir/prim/vector_expr.rs
+++ b/rust/src/ir/prim/vector_expr.rs
@@ -209,7 +209,7 @@ impl Ramp {
);
let message = format!(
"Ramp {field} must be scalar int32, int64, uint32, or
uint64, but got {}",
- support::repr(ty)
+ ty
);
return Err(support::check_failed(tvm_ffi::TYPE_ERROR,
&condition, &message));
}
@@ -449,7 +449,7 @@ fn vector_lanes(element: &PrimType, lanes: PrimExpr) ->
Result<(PrimType, PrimEx
return Ok((ty, IntImm::int32(count).into()));
}
let Some(factor) = vscale_factor(&lanes)? else {
- let message = format!("Invalid expression for scalable lanes {}",
support::repr(&lanes));
+ let message = format!("Invalid expression for scalable lanes {lanes}");
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
"vscale_factor", &message));
};
let ty = PrimType::try_from(support::scalable_vector(element.dtype,
factor)?)?;
diff --git a/rust/src/ir/stmt.rs b/rust/src/ir/stmt.rs
index 363f4f89b..9d288955b 100644
--- a/rust/src/ir/stmt.rs
+++ b/rust/src/ir/stmt.rs
@@ -46,7 +46,10 @@
// tvm-ffi-stubgen(import-object): tvm_ffi::DLDataTypeCode
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectRefCore
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
-// tvm-ffi-stubgen(import-object): tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex
+// tvm-ffi-stubgen(import-object): tvm_ffi::reflection::NativeFunction
+// tvm-ffi-stubgen(import-object): tvm_ffi::reflection::TypeAttrColumn
+// tvm-ffi-stubgen(import-object): std::sync::LazyLock
+// tvm-ffi-stubgen(import-object): crate::ir::expr::CallObj
// tvm-ffi-stubgen(import-object): crate::ir::expr::Call
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntImm
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntImmObj
@@ -73,6 +76,7 @@
// tvm-ffi-stubgen(begin): import-section
use crate::ir::attrs::DictAttrs;
use crate::ir::expr::Call;
+use crate::ir::expr::CallObj;
use crate::ir::expr::Expr;
use crate::ir::expr::IntImm;
use crate::ir::expr::IntImmObj;
@@ -90,6 +94,7 @@ use crate::ir::r#type::PrimType;
use crate::ir::support;
use std::ops::Deref;
use std::ptr::addr_of_mut;
+use std::sync::LazyLock;
use tvm_ffi::Any;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
@@ -101,12 +106,13 @@ use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
-use tvm_ffi::ShallowCopy;
use tvm_ffi::String;
-use tvm_ffi::UnsafeInit;
use tvm_ffi::VALUE_ERROR;
+use tvm_ffi::reflection::NativeFunction;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::TypeAttrColumn;
+use tvm_ffi::reflection::UnsafeInit;
use tvm_ffi::reflection::drop_if_defined;
-use tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
// tvm-ffi-stubgen(end)
/// The base of every statement: its source location `loc`.
@@ -376,8 +382,7 @@ impl AssertStmt {
if !condition_ty.matches_code(&[DLDataTypeCode::kDLBool]) {
let message = format!(
"AssertStmt should have boolean condition, but received {}
with dtype {}",
- support::repr(&condition),
- support::repr(condition_ty)
+ condition, condition_ty
);
let check = "condition_ty.MatchesCode(DLDataTypeCode::kDLBool)";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR, check,
&message));
@@ -551,8 +556,7 @@ impl For {
let message = format!(
"TIR For nodes require a scalar integer as the {field},
but received {} with \
dtype {}",
- support::repr(expr),
- support::repr(ty)
+ expr, ty
);
let condition = "dtype.IsScalar() &&
(dtype.MatchesCode(DLDataTypeCode::kDLUInt, \
DLDataTypeCode::kDLInt))";
@@ -569,8 +573,7 @@ impl For {
if !int_imm_value_fits(&imm.value, &var_ty)? {
let message = format!(
"Literal value {} is not representable in the loop
variable's dtype ({})",
- imm.value,
- support::repr(&var_ty)
+ imm.value, var_ty
);
let condition = "IntImmValueFits(a->value, var_ty)";
return Err(support::check_failed(
@@ -584,8 +587,7 @@ impl For {
if ty.bits() > var_ty.bits() {
let message = format!(
" Loop variable's dtype ({}) is narrower than that of
`begin` or `end` ({})",
- support::repr(&var_ty),
- support::repr(&ty)
+ var_ty, ty
);
let condition = "e_ty.bits() <= var_ty.bits()";
return Err(support::check_failed(tvm_ffi::INTERNAL_ERROR,
condition, &message));
@@ -597,7 +599,7 @@ impl For {
if ty.dtype == var_ty.dtype {
return Ok(());
}
- let message = format!("{} vs {}", support::repr(&var_ty),
support::repr(ty));
+ let message = format!("{var_ty} vs {ty}");
Err(support::check_failed(tvm_ffi::INTERNAL_ERROR, condition,
&message))
};
let (begin, end) = (promote(begin)?, promote(end)?);
@@ -775,28 +777,31 @@ impl TensorStore {
value: impl IntoExpr<PrimExpr>,
loc: impl Into<Location>,
) -> Result<Self> {
- static VALIDATE: support::TypeAttr =
- support::TypeAttr::new("__tvm_tensor_store_validate__");
+ // The native library registers the column when it loads, so one
lookup suffices.
+ static VALIDATE: LazyLock<Option<TypeAttrColumn>> =
+ LazyLock::new(||
TypeAttrColumn::new("__tvm_tensor_store_validate__"));
let (dest, value) = (dest.into_owned_expr()?,
value.into_owned_expr()?);
let indices = indices
.into_iter()
.map(IntoExpr::into_owned_expr)
.collect::<Result<Array<PrimExpr>>>()?;
let ty_index = AnyView::from(&dest.ty).type_index();
- let Some(validate) = VALIDATE.get(ty_index) else {
- let message =
- format!("Type {} does not support TensorStore",
support::type_key(ty_index));
+ let Some(validate) = VALIDATE.and_then(|column| column.get(ty_index))
else {
+ let message = format!(
+ "Type {} does not support TensorStore",
+ tvm_ffi::reflection::type_key(ty_index).unwrap_or_default()
+ );
return Err(support::check_failed(
tvm_ffi::TYPE_ERROR,
"validate.has_value()",
&message,
));
};
- let result = support::call_with_refs(validate, &dest, &indices,
&value)?;
- if result.type_index() != TVMFFITypeIndex::kTVMFFINone as i32 {
- let message = "Packed hook result type mismatch";
- return Err(tvm_ffi::Error::new(tvm_ffi::TYPE_ERROR, message, ""));
- }
+ // SAFETY: the column holds
+ // `void(const Expr&, const Array<PrimExpr>&, const PrimExpr&)` hooks.
+ let validate: NativeFunction<(&Expr, &Array<PrimExpr>, &PrimExpr)> =
+ unsafe { NativeFunction::from_any(validate)? };
+ validate.call((&dest, &indices, &value))?;
Ok(Self::from_fields(loc.into(), dest, indices, value))
}
}
@@ -945,8 +950,11 @@ impl RegionStmt {
}
}
let region = Self::from_fields(loc.into(), op, args, body_params,
attrs, body, result_vars);
- if let Some(validate) = support::NativeHook::of_op(®ion.op,
"FRegionValidate")? {
- validate.call_void(®ion)?;
+ // SAFETY: `FRegionValidate` holds `void(const RegionStmtObj*)` hooks.
+ let validate =
+ unsafe { support::op_hook::<(&RegionStmtObj,), ()>(®ion.op,
"FRegionValidate")? };
+ if let Some(validate) = validate {
+ validate.call((&*region,))?;
}
Ok(region)
}
@@ -1362,7 +1370,10 @@ fn int_imm_value_fits(value: &BigInt, ty: &PrimType) ->
Result<bool> {
/// `FRegionGetBodyParams` gives once its validator accepts the arguments.
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 {
+ // SAFETY: `FRegionGetBodyParams` holds `Array<Var>(const CallObj*)` hooks.
+ let get_body_params =
+ unsafe { support::op_hook::<(&CallObj,), Array<Var>>(op,
"FRegionGetBodyParams")? };
+ let Some(get_body_params) = get_body_params else {
let message = format!(
"{name} does not support region construction: FRegionGetBodyParams
is required"
);
@@ -1377,7 +1388,7 @@ fn region_body_params(op: &Op, args: &Array<Expr>, attrs:
&DictAttrs) -> Result<
Array::from_vec(vec![]),
);
signature.validate()?;
- let params: Array<Var> = get_body_params.call_as(&signature)?;
+ let params = get_body_params.call((&*signature,))?;
let mut definitions = std::collections::HashSet::new();
for param in params.iter() {
if param.ty.downcast_ref::<MissingTypeObj>().is_some() {
diff --git a/rust/src/ir/support.rs b/rust/src/ir/support.rs
index 6860d139e..25d143f5a 100644
--- a/rust/src/ir/support.rs
+++ b/rust/src/ir/support.rs
@@ -23,19 +23,14 @@
use std::cell::{OnceCell, RefCell};
use std::collections::HashMap;
-use std::ffi::c_void;
use std::fmt::Display;
use std::rc::Rc;
-use std::sync::OnceLock;
use std::thread::LocalKey;
-use tvm_ffi::tvm_ffi_sys::{
- TVMFFIAny, TVMFFIFieldFlagBitMask, TVMFFIGetTypeInfo, TVMFFIObject,
TVMFFITypeIndex,
-};
use tvm_ffi::{
- Any, AnyView, DLDataType, DLDataTypeCode, DLDataTypeExt, Error, ErrorKind,
Function, ObjectArc,
- ObjectRefCore, Result, TypeAttrColumn, INTERNAL_ERROR, TYPE_ERROR,
+ Any, DLDataType, DLDataTypeCode, DLDataTypeExt, Error, ErrorKind, Result,
INTERNAL_ERROR,
};
+use tvm_ffi::reflection::NativeFunction;
use super::{Op, OpAttrMap, OpObj, PrimType};
@@ -108,55 +103,11 @@ pub(crate) fn compare_failed(
Error::new(kind, &format!("Check failed: {condition} ({lhs} vs. {rhs}) :
{message}"), "")
}
-/// A value as error messages print it (`ffi.ReprPrint`).
-pub(crate) fn repr<T>(value: &T) -> std::string::String
-where
- for<'a> AnyView<'a>: From<&'a T>,
-{
- tvm_ffi::cached_global_func!("ffi.ReprPrint")
- .call_packed(&[AnyView::from(value)])
- .and_then(tvm_ffi::String::try_from)
- .map(|text| text.as_str().to_string())
- .unwrap_or_default()
-}
-
/// A dtype as error messages print it, such as `int32`.
pub(crate) fn dtype_name(dtype: DLDataType) -> std::string::String {
dtype.to_string()
}
-/// The registered key of a type index.
-pub(crate) fn type_key(type_index: i32) -> std::string::String {
- // SAFETY: the type index of a live value is registered, and its
information is static.
- unsafe {
- let info = TVMFFIGetTypeInfo(type_index);
- (*info).type_key.as_str().to_string()
- }
-}
-
-/// The default value the reflection of the type of `type_index` registers for
its own field
-/// `name`, or `None` when it registers none or a factory.
-pub(crate) fn field_default(type_index: i32, name: &str) -> Option<Any> {
- use TVMFFIFieldFlagBitMask::*;
- // SAFETY: the type information of a registered type is static, and its
`num_fields` fields
- // are initialized.
- let field = unsafe {
- let info = &*TVMFFIGetTypeInfo(type_index);
- let fields = std::slice::from_raw_parts(info.fields, info.num_fields
as usize);
- fields.iter().find(|field| field.name.as_str() == name)?
- };
- let has_default = field.flags & kTVMFFIFieldFlagBitMaskHasDefault as i64
!= 0;
- let from_factory = field.flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory
as i64 != 0;
- if !has_default || from_factory {
- return None;
- }
- // SAFETY: the registry owns the default value; it is copied, never
released.
- let value = std::mem::ManuallyDrop::new(unsafe {
- Any::from_raw_ffi_any(field.default_value_or_factory)
- });
- Some((*value).clone())
-}
-
/// A double as C++ `std::ostream` prints it by default: `%g`, six significant
digits.
#[derive(Clone, Copy)]
pub(crate) struct CppDouble(pub f64);
@@ -240,120 +191,17 @@ pub(crate) fn scalable_vector(dtype: DLDataType, factor:
i32) -> Result<DLDataTy
//
---------------------------------------------------------------------------------------------
// Registry hooks
-/// A C++ `NativeFunctionView<R(const Node*)>` hook that a registry holds: the
C ABI adapter its
-/// opaque pointer names, or a packed function a frontend registered in its
place.
-pub(crate) enum NativeHook {
- Adapter(extern "C" fn(*const c_void) -> TVMFFIAny),
- Packed(Function),
-}
-
-impl NativeHook {
- /// The hook `value` holds.
- pub(crate) fn from_any(value: Any) -> Result<Self> {
- if value.type_index() != TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 {
- return Ok(Self::Packed(value.try_into()?));
- }
- let address: *mut c_void = value.try_into()?;
- // SAFETY: the opaque pointer of a `NativeFunctionView` is its C ABI
adapter.
- let adapter = unsafe {
- std::mem::transmute::<*mut c_void, extern "C" fn(*const c_void) ->
TVMFFIAny>(address)
- };
- Ok(Self::Adapter(adapter))
- }
-
- /// The hook `op` registers in the attribute column `column`, if any.
- pub(crate) fn of_op(op: &OpObj, column: &'static str) ->
Result<Option<Self>> {
- op_attr(op, column)?.map(Self::from_any).transpose()
- }
-
- /// Call the hook on `node`: its result, or the error it reports.
- pub(crate) fn call<T: ObjectRefCore>(&self, node: &T) -> Result<Any>
- where
- for<'a> AnyView<'a>: From<&'a T>,
- {
- let result = match self {
- // SAFETY: `node` keeps the object alive for the call, and the
adapter returns an
- // owned value.
- Self::Adapter(adapter) => unsafe {
- let node = ObjectArc::as_ptr(ObjectRefCore::data(node)) as
*const c_void;
- Any::from_raw_ffi_any(adapter(node))
- },
- Self::Packed(function) =>
function.call_packed(&[AnyView::from(node)])?,
- };
- expected(result)
- }
-
- /// [`NativeHook::call`] of a hook that returns an `R`; a packed hook's
result of another type
- /// is a `TypeError`.
- pub(crate) fn call_as<T: ObjectRefCore, R: TryFrom<Any>>(&self, node: &T)
-> Result<R>
- where
- for<'a> AnyView<'a>: From<&'a T>,
- {
- R::try_from(self.call(node)?).map_err(|_| hook_result_mismatch())
- }
-
- /// [`NativeHook::call`] of a hook that returns nothing.
- pub(crate) fn call_void<T: ObjectRefCore>(&self, node: &T) -> Result<()>
- where
- for<'a> AnyView<'a>: From<&'a T>,
- {
- match self.call(node)?.type_index() == TVMFFITypeIndex::kTVMFFINone as
i32 {
- true => Ok(()),
- false => Err(hook_result_mismatch()),
- }
- }
-}
-
-/// Call a hook that takes three object references, such as a C++
-/// `NativeFunctionView<void(const Expr&, const Array<PrimExpr>&, const
PrimExpr&)>`, and return its
-/// result or the error it reports. The C ABI adapter takes each reference as
the address of an
-/// object pointer; a packed function takes the values.
-pub(crate) fn call_with_refs<A, B, C>(hook: Any, a: &A, b: &B, c: &C) ->
Result<Any>
-where
- A: ObjectRefCore,
- B: ObjectRefCore,
- C: ObjectRefCore,
- for<'x> AnyView<'x>: From<&'x A> + From<&'x B> + From<&'x C>,
-{
- if hook.type_index() != TVMFFITypeIndex::kTVMFFIOpaquePtr as i32 {
- let function: Function = hook.try_into()?;
- return expected(function.call_packed(&[
- AnyView::from(a),
- AnyView::from(b),
- AnyView::from(c),
- ])?);
- }
- type Adapter = extern "C" fn(
- *const *const TVMFFIObject,
- *const *const TVMFFIObject,
- *const *const TVMFFIObject,
- ) -> TVMFFIAny;
- fn object<T: ObjectRefCore>(value: &T) -> *const TVMFFIObject {
- // The pointer is only read while `value` keeps the object alive.
- ObjectArc::as_ptr(ObjectRefCore::data(value)) as *const TVMFFIObject
- }
- let address: *mut c_void = hook.try_into()?;
- let (a, b, c) = (object(a), object(b), object(c));
- // SAFETY: the opaque pointer of a `NativeFunctionView` is its C ABI
adapter; a C++ object
- // reference is the address of its single object pointer, which outlives
the call; and the
- // adapter returns an owned value.
- let result = unsafe {
- let adapter = std::mem::transmute::<*mut c_void, Adapter>(address);
- Any::from_raw_ffi_any(adapter(&a, &b, &c))
- };
- expected(result)
-}
-
-/// The value of an `Expected` result, or its error.
-fn expected(result: Any) -> Result<Any> {
- match result.type_index() == TVMFFITypeIndex::kTVMFFIError as i32 {
- true => Err(result.try_into()?),
- false => Ok(result),
- }
-}
-
-fn hook_result_mismatch() -> Error {
- Error::new(TYPE_ERROR, "Packed hook result type mismatch", "")
+/// The hook `op` registers in the attribute column `column`, if any.
+///
+/// # Safety
+///
+/// A native function in the column must take `Args` and return `R`, as its
declaration says.
+pub(crate) unsafe fn op_hook<Args, R>(
+ op: &OpObj,
+ column: &'static str,
+) -> Result<Option<NativeFunction<Args, R>>> {
+ // SAFETY: the caller vouches for the signature.
+ op_attr(op, column)?.map(|hook| unsafe { NativeFunction::from_any(hook)
}).transpose()
}
thread_local! {
@@ -390,27 +238,3 @@ pub(crate) fn cached_op(name: &'static str) -> Result<Op> {
OPS.with(|ops| ops.borrow_mut().insert(name, op.clone()));
Ok(op)
}
-
-/// A type attribute column, looked up once it exists.
-pub(crate) struct TypeAttr {
- name: &'static str,
- column: OnceLock<TypeAttrColumn>,
-}
-
-impl TypeAttr {
- pub(crate) const fn new(name: &'static str) -> Self {
- Self { name, column: OnceLock::new() }
- }
-
- /// The value the type of `type_index` registers in this column, if any.
- pub(crate) fn get(&self, type_index: i32) -> Option<Any> {
- let column = match self.column.get() {
- Some(column) => *column,
- None => {
- let column = TypeAttrColumn::new(self.name)?;
- *self.column.get_or_init(|| column)
- }
- };
- column.get(type_index)
- }
-}
diff --git a/rust/src/ir/transform/core.rs b/rust/src/ir/transform/core.rs
index 88245c3b9..57a3357d0 100644
--- a/rust/src/ir/transform/core.rs
+++ b/rust/src/ir/transform/core.rs
@@ -36,13 +36,13 @@ use crate::ir::support;
use std::ops::Deref;
use tvm_ffi::Any;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// A transformation of an [`IRModule`], with a name and an optimization level
([`PassInfo`]).
diff --git a/rust/src/ir/type.rs b/rust/src/ir/type.rs
index 1ab2cc3cf..5d03a591e 100644
--- a/rust/src/ir/type.rs
+++ b/rust/src/ir/type.rs
@@ -61,9 +61,9 @@ use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
-use tvm_ffi::ShallowCopy;
-use tvm_ffi::UnsafeInit;
use tvm_ffi::VALUE_ERROR;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
use tvm_ffi::reflection::drop_if_defined;
// tvm-ffi-stubgen(end)
diff --git a/rust/src/ir/type_def.rs b/rust/src/ir/type_def.rs
index d73efee64..a589a4fd7 100644
--- a/rust/src/ir/type_def.rs
+++ b/rust/src/ir/type_def.rs
@@ -22,8 +22,6 @@
use tvm_ffi::{Function, ObjectCore, ObjectRefCore, Result, String};
-use super::support;
-
/// Expose the IR type of the reference `R`, an `rs_object` type, to Python.
///
/// Call it in a [`tvm_ffi::static_init_block!`], so that the type is
registered when its program or
@@ -71,9 +69,10 @@ use super::support;
/// }
/// ```
pub fn register_type<R: ObjectRefCore>() -> Result<()> {
- let type_index = R::ContainerType::type_index();
+ // The key of the registered type, which `type_index` registers first.
+ let key =
tvm_ffi::reflection::type_key(R::ContainerType::type_index()).unwrap_or_default();
if let Ok(expose) = Function::get_global("ir.expose_type") {
-
expose.call_tuple((String::from(support::type_key(type_index).as_str()),))?;
+ expose.call_tuple((String::from(key),))?;
}
Ok(())
}
diff --git a/rust/src/runtime.rs b/rust/src/runtime.rs
index 68f3b78e5..d008993e0 100644
--- a/rust/src/runtime.rs
+++ b/rust/src/runtime.rs
@@ -36,28 +36,20 @@
use tvm_ffi::Any;
use tvm_ffi::DLDataType;
use tvm_ffi::DLDevice;
-use tvm_ffi::Object;
+use tvm_ffi::EnumObj;
use tvm_ffi::ObjectArc;
use tvm_ffi::Result;
use tvm_ffi::Shape;
use tvm_ffi::String;
use tvm_ffi::Tensor;
-
-/// Header-only stand-in for the builtin `ffi.Enum`; it only carries the
ancestor depth.
-#[allow(dead_code)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "ffi.Enum"]
-struct FfiEnumObj {
- base_: Object,
-}
// tvm-ffi-stubgen(end)
// `ir.PtrType.storage_scope` refers to this type.
/// A registered memory space, such as `global` or `shared`, that a pointer or
tensor lives in.
///
/// Registration order has no meaning, and a backend must support a scope
before it can allocate
-/// storage in it; [`MemoryScope::register`] adds one.
+/// storage in it; [`EnumDef::new`](tvm_ffi::reflection::EnumDef::new) adds
one, and
+/// [`EnumRef::get`](tvm_ffi::EnumRef::get) looks one up by name.
// tvm-ffi-stubgen(begin): object/runtime.MemoryScope
#[allow(missing_docs)]
#[repr(C)]
@@ -68,31 +60,25 @@ pub struct MemoryScope {
/// The object behind [`MemoryScope`].
///
-/// Opaque: parent 'ffi.Enum' is opaque (no-mirror). Fields are read through
the C ABI getters.
+/// Complete: reflected fields fill [48, 48) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "runtime.MemoryScope"]
#[type_final]
pub struct MemoryScopeObj {
- base_: FfiEnumObj,
+ base_: EnumObj,
}
+const _: () = {
+ assert!(::core::mem::size_of::<MemoryScopeObj>() == 48);
+ assert!(::core::mem::align_of::<MemoryScopeObj>() == 8);
+};
+
unsafe impl tvm_ffi::object::ImmutableObject for MemoryScopeObj {}
-// tvm-ffi-stubgen(end)
-impl MemoryScope {
- /// Register the memory scope `name` (`runtime.MemoryScopeRegister`).
- ///
- /// # Errors
- ///
- /// Returns an error if `name` is empty, contains a null character, or is
already registered.
- pub fn register(name: &str) -> Result<Self> {
- tvm_ffi::cached_global_func!("runtime.MemoryScopeRegister")
- .call_tuple((String::from(name),))?
- .try_into()
- }
-}
+impl tvm_ffi::EnumRef for MemoryScope {}
+// tvm-ffi-stubgen(end)
/// Allocate an uninitialized tensor on `device`, in the memory scope
`mem_scope` or the global one
/// (`runtime.TensorAllocWithScope`).
diff --git a/rust/src/sym/constraint.rs b/rust/src/sym/constraint.rs
index 4bcff25dd..3ebe83c17 100644
--- a/rust/src/sym/constraint.rs
+++ b/rust/src/sym/constraint.rs
@@ -41,8 +41,8 @@ use std::ops::Deref;
use std::ptr::addr_of_mut;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
-use tvm_ffi::ShallowCopy;
-use tvm_ffi::UnsafeInit;
+use tvm_ffi::reflection::ShallowCopy;
+use tvm_ffi::reflection::UnsafeInit;
// tvm-ffi-stubgen(end)
/// The base of every fact a [`ConstraintContext`] holds; it owns no analyzer
state.
diff --git a/rust/src/sym/int_set.rs b/rust/src/sym/int_set.rs
index 014296323..570e6e770 100644
--- a/rust/src/sym/int_set.rs
+++ b/rust/src/sym/int_set.rs
@@ -46,12 +46,12 @@ use crate::sym::analyzer::Analyzer;
use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of every set of integers, one dimension each, that bounds the
values an expression
diff --git a/rust/src/sym/iter_affine_map.rs b/rust/src/sym/iter_affine_map.rs
index 278140833..e3b8ceff7 100644
--- a/rust/src/sym/iter_affine_map.rs
+++ b/rust/src/sym/iter_affine_map.rs
@@ -61,13 +61,13 @@ use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of the iterator map forms, [`IterSplitExpr`] and [`IterSumExpr`].
They hold the result
diff --git a/rust/src/target.rs b/rust/src/target.rs
index 5ccd51896..f396ebdfe 100644
--- a/rust/src/target.rs
+++ b/rust/src/target.rs
@@ -33,7 +33,6 @@
use crate::ir::module::IRModule;
use tvm_ffi::Any;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Map;
use tvm_ffi::Module;
use tvm_ffi::Object;
@@ -42,6 +41,7 @@ use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
use tvm_ffi::object::ObjectRef;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// A compilation target: its `kind`, the `keys` that classify it, such as
`cpu`, the `attrs` of
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index ba8362e09..827ed534b 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -48,22 +48,14 @@ use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::Any;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
+use tvm_ffi::EnumObj;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
use tvm_ffi::String;
-
-/// Header-only stand-in for the builtin `ffi.Enum`; it only carries the
ancestor depth.
-#[allow(dead_code)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "ffi.Enum"]
-struct FfiEnumObj {
- base_: Object,
-}
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The base of every layout: a map from the logical indices of a tensor to a
coordinate on each
@@ -82,7 +74,7 @@ struct FfiEnumObj {
/// use tvm_tirx::tirx::{Axis, Iter, Layout, TileLayout};
/// use tvm_tirx::tvm_ffi::Map;
///
-/// let m = Axis::get("m");
+/// let m = Axis::get_or_register("m");
/// let shard = [Iter::new(8, 8, m.clone()), Iter::new(8, 1, m)];
/// let layout = Layout::from(&TileLayout::new(shard, Vec::<Iter>::new(),
&Map::new()));
/// assert_eq!(script(&layout.size(None)?)?, "64");
@@ -398,8 +390,8 @@ unsafe impl tvm_ffi::object::ImmutableObject for IterObj {}
// The native `tirx.Iter` stores its fields as given, as the generated
`Iter::new` does.
/// A named coordinate a layout maps to: the memory axis `m`, or a thread
axis, such as `laneid`,
-/// that indexes the threads of a scope. Each name is one registered axis,
which [`Axis::get`]
-/// looks up.
+/// that indexes the threads of a scope. Each name is one registered axis,
which
+/// [`Axis::get_or_register`] looks up, registering it on first use.
// tvm-ffi-stubgen(begin): object/tirx.Axis
#[allow(missing_docs)]
#[repr(C)]
@@ -410,23 +402,30 @@ pub struct Axis {
/// The object behind [`Axis`].
///
-/// Opaque: parent 'ffi.Enum' is opaque (no-mirror). Fields are read through
the C ABI getters.
+/// Complete: reflected fields fill [48, 48) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "tirx.Axis"]
#[type_final]
pub struct AxisObj {
- base_: FfiEnumObj,
+ base_: EnumObj,
}
+const _: () = {
+ assert!(::core::mem::size_of::<AxisObj>() == 48);
+ assert!(::core::mem::align_of::<AxisObj>() == 8);
+};
+
unsafe impl tvm_ffi::object::ImmutableObject for AxisObj {}
+
+impl tvm_ffi::EnumRef for Axis {}
// tvm-ffi-stubgen(end)
impl Axis {
/// Look up the axis named `name`, such as the memory axis `m`,
registering it on first use
- /// (`tirx.AxisGet`).
- pub fn get(name: &str) -> Self {
+ /// (`tirx.AxisGet`). [`EnumRef::get`](tvm_ffi::EnumRef::get) only looks
up.
+ pub fn get_or_register(name: &str) -> Self {
support::infallible(
tvm_ffi::cached_global_func!("tirx.AxisGet")
.call_tuple((String::from(name),))
diff --git a/rust/src/tirx/op.rs b/rust/src/tirx/op.rs
index b5262acd5..4a0c02205 100644
--- a/rust/src/tirx/op.rs
+++ b/rust/src/tirx/op.rs
@@ -183,7 +183,7 @@ pub fn logaddexp_at(
for (operand, side) in [(&a, "a"), (&b, "b")] {
if !operand.prim_ty().matches_code(&[DLDataTypeCode::kDLFloat]) {
let condition =
format!("{side}.ty().MatchesCode(DLDataTypeCode::kDLFloat)");
- let message = support::repr(operand);
+ let message = operand.to_string();
return Err(support::check_failed(INTERNAL_ERROR, &condition,
&message));
}
}
@@ -229,10 +229,7 @@ pub fn reinterpret_at(
};
if value.downcast_ref::<StringImmObj>().is_some() {
if !target_is_pointer {
- let message = format!(
- "String reinterpret requires a pointer target, but got {}",
- support::repr(&ty)
- );
+ let message = format!("String reinterpret requires a pointer
target, but got {ty}");
return type_error("target_ty.as<PtrTypeObj>()", message);
}
} else if let Some(target) = ty.downcast_ref::<PrimType>() {
@@ -252,38 +249,28 @@ pub fn reinterpret_at(
((value_dtype.MatchesCode(DLDataTypeCode::kDLFloat4_e2m1fn) || \
target_prim.MatchesCode(DLDataTypeCode::kDLFloat4_e2m1fn)) && \
value_dtype.StorageBytes() ==
target_prim.StorageBytes())";
- let message = format!(
- "Reinterpret requires size match {} vs {}",
- support::repr(&ty),
- support::repr(&value.ty)
- );
+ let message = format!("Reinterpret requires size match {}
vs {}", ty, value.ty);
return Err(support::check_failed(INTERNAL_ERROR,
condition, &message));
}
}
} else {
if !source_is_pointer {
- let message = format!(
- "Reinterpret source must be PrimType or PtrType, but got
{}",
- support::repr(&value.ty)
- );
+ let message =
+ format!("Reinterpret source must be PrimType or PtrType,
but got {}", value.ty);
return type_error("value->ty.as<PtrTypeObj>()", message);
}
if !is_scalar_int64(target) {
let condition = "target_dtype.value().IsScalar() &&
target_dtype.value().bits() == 64 && \
target_dtype.value().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)";
let message = format!(
- "Pointer reinterpret requires a scalar 64-bit integer
target, but got {}",
- support::repr(&ty)
+ "Pointer reinterpret requires a scalar 64-bit integer
target, but got {ty}"
);
return type_error(condition, message);
}
}
} else {
if !target_is_pointer {
- let message = format!(
- "Reinterpret target must be PrimType or PtrType, but got {}",
- support::repr(&ty)
- );
+ let message = format!("Reinterpret target must be PrimType or
PtrType, but got {ty}");
return type_error("target_ty.as<PtrTypeObj>()", message);
}
if let Some(source) = value.ty.downcast_ref::<PrimType>() {
@@ -292,15 +279,13 @@ pub fn reinterpret_at(
source_dtype.value().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)";
let message = format!(
"Pointer reinterpret requires a scalar 64-bit integer
source, but got {}",
- support::repr(&value.ty)
+ value.ty
);
return type_error(condition, message);
}
} else if !source_is_pointer {
- let message = format!(
- "Reinterpret source must be PrimType or PtrType, but got {}",
- support::repr(&value.ty)
- );
+ let message =
+ format!("Reinterpret source must be PrimType or PtrType, but
got {}", value.ty);
return type_error("value->ty.as<PtrTypeObj>()", message);
}
}
diff --git a/rust/src/tirx/structural.rs b/rust/src/tirx/structural.rs
index bfc48d411..b452bc262 100644
--- a/rust/src/tirx/structural.rs
+++ b/rust/src/tirx/structural.rs
@@ -25,11 +25,11 @@
// tvm-ffi-stubgen(begin): import-section
use crate::sym::constraint::ConstraintContext;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
+use tvm_ffi::reflection::FieldGetter;
/// Header-only stand-in for the builtin `ffi.StructuralMutator`; it only
carries the ancestor depth.
#[allow(dead_code)]
diff --git a/rust/src/tirx/tensor.rs b/rust/src/tirx/tensor.rs
index d28bc7516..763f27786 100644
--- a/rust/src/tirx/tensor.rs
+++ b/rust/src/tirx/tensor.rs
@@ -17,9 +17,9 @@
* under the License.
*/
-//! Tensors. A tensor is a [`Var`] whose type is a [`TensorType`], which fixes
its shape, element
-//! type, memory scope, alignment and [`Layout`]; loads and stores subscript
the variable, and
-//! [`TensorVar`] matches one in an operator signature.
+//! Tensors. A tensor is a [`Var`](crate::ir::Var) whose type is a
[`TensorType`], which fixes its
+//! shape, element type, memory scope, alignment and [`Layout`]; loads and
stores subscript the
+//! variable, and [`TensorVar`] matches one in an operator signature.
// tvm-ffi-stubgen(missing-docs): allow
// `custom-new` reserves `new` for a hand-written constructor (defaults); the
generated
@@ -28,38 +28,30 @@
// Names the hand-written wrappers use: the import section keeps them even
when no generated block
// does.
-// tvm-ffi-stubgen(import-object): crate::ir::expr::Var
// tvm-ffi-stubgen(import-object): crate::sym::analyzer::Analyzer
// tvm-ffi-stubgen(import-object): tvm_ffi::AnyView
// tvm-ffi-stubgen(import-object): tvm_ffi::ObjectCore
// tvm-ffi-stubgen(import-object): tvm_ffi::Result
-// tvm-ffi-stubgen(import-object): crate::ir::support
// tvm-ffi-stubgen(import-object): crate::ir::sig
// tvm-ffi-stubgen(import-object): crate::ir::expr::IntoExpr
// tvm-ffi-stubgen(import-object): crate::ir::expr::PrimExpr
// A tensor's shape is a list of `PrimExpr`s, as C++ `TensorTypeObj` holds it.
// tvm-ffi-stubgen(field): tirx.TensorType.shape -> Array<PrimExpr>
// tvm-ffi-stubgen(import-object): tvm_ffi::Array
-// tvm-ffi-stubgen(import-object): std::borrow::Borrow
-// tvm-ffi-stubgen(import-object): tvm_ffi::IntoObjectRef
// tvm-ffi-stubgen(begin): import-section
use crate::ir::expr::IntoExpr;
use crate::ir::expr::PrimExpr;
-use crate::ir::expr::Var;
use crate::ir::location::Location;
use crate::ir::r#type::PrimType;
use crate::ir::r#type::Type;
use crate::ir::r#type::TypeObj;
use crate::ir::sig;
-use crate::ir::support;
use crate::runtime::MemoryScope;
use crate::sym::analyzer::Analyzer;
use crate::tirx::layout::Layout;
-use std::borrow::Borrow;
use std::ops::Deref;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
-use tvm_ffi::IntoObjectRef;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::Result;
@@ -78,8 +70,9 @@ use tvm_ffi::Result;
/// use tvm_tirx::runtime::MemoryScope;
/// use tvm_tirx::script::script;
/// use tvm_tirx::tirx::TensorType;
+/// use tvm_tirx::tvm_ffi::reflection::EnumDef;
///
-/// let scratch = MemoryScope::register("my_ext.scratch")?;
+/// let scratch = EnumDef::<MemoryScope>::new("my_ext.scratch")?.value();
/// let ty = TensorType::new([16, 16], PrimType::float(32), &scratch, 0, None,
UNKNOWN_LOC)?;
/// assert_eq!(script(&ty)?, r#"T.Tensor((16, 16), T.float32,
scope="my_ext.scratch")"#);
///
@@ -209,8 +202,8 @@ impl TensorType {
}
}
-/// A signature descriptor ([`sig::ArgType`]) that matches a [`Var`] whose
type is a
-/// [`TensorType`]; it is not a type of its own.
+/// A signature descriptor ([`sig::ArgType`]) that matches a
[`Var`](crate::ir::Var) whose type is
+/// a [`TensorType`]; it is not a type of its own.
pub struct TensorVar;
impl sig::ArgType for TensorVar {
@@ -223,19 +216,6 @@ impl sig::ArgType for TensorVar {
}
}
-/// 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`).
-///
-/// # Errors
-///
-/// Returns an error if the type of `tensor` is not a [`TensorType`].
-pub fn tensor_is_scalar(tensor: impl IntoObjectRef<Var>, alloc_or_decl: bool)
-> Result<bool> {
- tvm_ffi::cached_global_func!("tirx.TensorIsScalar")
- .call_tuple((tensor.into_object_ref().borrow(), alloc_or_decl))?
- .try_into()
-}
-
/// The runtime's allocation alignment, which a build may configure, read from
the default the
/// reflection registers for `TensorType::data_alignment`.
fn alloc_alignment() -> Result<i32> {
@@ -243,7 +223,7 @@ fn alloc_alignment() -> Result<i32> {
if let Some(alignment) = ALLOC_ALIGNMENT.get() {
return Ok(*alignment);
}
- let Some(default) = support::field_default(TensorTypeObj::type_index(),
"data_alignment")
+ let Some(default) =
tvm_ffi::reflection::field_default::<TensorTypeObj>("data_alignment")
else {
let message = "tirx.TensorType registers no default data_alignment";
return Err(tvm_ffi::Error::new(tvm_ffi::INTERNAL_ERROR, message, ""));
diff --git a/rust/src/tirx/thread_scope.rs b/rust/src/tirx/thread_scope.rs
index 5d0d69fd0..ea69dd257 100644
--- a/rust/src/tirx/thread_scope.rs
+++ b/rust/src/tirx/thread_scope.rs
@@ -30,31 +30,23 @@
// tvm-ffi-stubgen(begin): import-section
use std::borrow::Borrow;
+use tvm_ffi::EnumObj;
use tvm_ffi::IntoObjectRef;
-use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
use tvm_ffi::Result;
use tvm_ffi::String;
-
-/// Header-only stand-in for the builtin `ffi.Enum`; it only carries the
ancestor depth.
-#[allow(dead_code)]
-#[repr(C)]
-#[derive(tvm_ffi::derive::Object)]
-#[type_key = "ffi.Enum"]
-struct FfiEnumObj {
- base_: Object,
-}
// tvm-ffi-stubgen(end)
/// A level of a backend's thread hierarchy, named `backend.level`, such as
`cuda.warp`.
///
-/// Backends register their scopes; [`ThreadScope::get`] only looks one up.
Within a backend,
+/// Backends register their scopes; [`EnumRef::get`] only looks one up. Within
a backend,
/// [`rank`](ThreadScope::rank) orders the levels, a wider scope ranking lower.
///
/// # Examples
///
/// ```
/// use tvm_tirx::tirx::ThreadScope;
+/// use tvm_tirx::tvm_ffi::EnumRef;
///
/// let warp = ThreadScope::get("cuda.warp")?;
/// assert_eq!((warp.backend()?, warp.local_name()?), ("cuda".into(),
"warp".into()));
@@ -63,6 +55,8 @@ struct FfiEnumObj {
/// assert!(cta.is_higher_than(&warp)?);
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
+///
+/// [`EnumRef::get`]: tvm_ffi::EnumRef::get
// tvm-ffi-stubgen(begin): object/tirx.ThreadScope
#[allow(missing_docs)]
#[repr(C)]
@@ -73,31 +67,27 @@ pub struct ThreadScope {
/// The object behind [`ThreadScope`].
///
-/// Opaque: parent 'ffi.Enum' is opaque (no-mirror). Fields are read through
the C ABI getters.
+/// Complete: reflected fields fill [48, 48) exactly.
#[allow(missing_docs)]
#[repr(C)]
#[derive(tvm_ffi::derive::Object)]
#[type_key = "tirx.ThreadScope"]
#[type_final]
pub struct ThreadScopeObj {
- base_: FfiEnumObj,
+ base_: EnumObj,
}
+const _: () = {
+ assert!(::core::mem::size_of::<ThreadScopeObj>() == 48);
+ assert!(::core::mem::align_of::<ThreadScopeObj>() == 8);
+};
+
unsafe impl tvm_ffi::object::ImmutableObject for ThreadScopeObj {}
+
+impl tvm_ffi::EnumRef for ThreadScope {}
// tvm-ffi-stubgen(end)
impl ThreadScope {
- /// Look up the thread scope named `name`, such as `cuda.warp`
(`tirx.ThreadScopeGet`).
- ///
- /// # Errors
- ///
- /// Returns an error if no thread scope is registered as `name`.
- pub fn get(name: &str) -> Result<Self> {
- tvm_ffi::cached_global_func!("tirx.ThreadScopeGet")
- .call_tuple((String::from(name),))?
- .try_into()
- }
-
/// Return the backend that registered the scope, such as `cuda`
(`tirx.ThreadScopeGetBackend`).
pub fn backend(&self) -> Result<String> {
tvm_ffi::cached_global_func!("tirx.ThreadScopeGetBackend").call_tuple((self,))?.try_into()
diff --git a/rust/src/tirx/tile/dispatch.rs b/rust/src/tirx/tile/dispatch.rs
index ab05b394c..06322e2dd 100644
--- a/rust/src/tirx/tile/dispatch.rs
+++ b/rust/src/tirx/tile/dispatch.rs
@@ -47,7 +47,6 @@ use crate::tirx::thread_scope::ThreadScope;
use tvm_ffi::Any;
use tvm_ffi::AnyView;
use tvm_ffi::Array;
-use tvm_ffi::FieldGetter;
use tvm_ffi::Map;
use tvm_ffi::Object;
use tvm_ffi::ObjectArc;
@@ -56,6 +55,7 @@ use tvm_ffi::ObjectRefCore;
use tvm_ffi::Result;
use tvm_ffi::String;
use tvm_ffi::object::ObjectRef;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// The context a tile instruction call is lowered in: the `target`, the
`thread_scope` of the
@@ -74,7 +74,7 @@ use tvm_ffi::object::ObjectRef;
/// use tvm_tirx::target::Target;
/// use tvm_tirx::tirx::tile::DispatchContext;
/// use tvm_tirx::tirx::ThreadScope;
-/// use tvm_tirx::tvm_ffi::Map;
+/// use tvm_tirx::tvm_ffi::{EnumRef, Map};
///
/// let tx = PrimVar::new("tx", PrimType::int(32), UNKNOWN_LOC);
/// let ctx = DispatchContext::new(
diff --git a/rust/src/tirx/tile/instruction.rs
b/rust/src/tirx/tile/instruction.rs
index ab939455c..ff644684a 100644
--- a/rust/src/tirx/tile/instruction.rs
+++ b/rust/src/tirx/tile/instruction.rs
@@ -232,7 +232,7 @@ pub fn as_tensor_region(value: impl IntoExpr) ->
Result<TensorRegion> {
/// use tvm_tirx::ir::{Attrs, Call, OpDef, PrimType, Type, UNKNOWN_LOC};
/// use tvm_tirx::tirx::tile::{configure_tensor_instruction, ScopeAttrs,
ScopeAttrsObj};
/// use tvm_tirx::tirx::ThreadScope;
-/// use tvm_tirx::tvm_ffi::{ensure, ObjectRefCore, Result, VALUE_ERROR};
+/// use tvm_tirx::tvm_ffi::{ensure, EnumRef, ObjectRefCore, Result,
VALUE_ERROR};
///
/// let op = OpDef::new("my_ext.tile_fill", "Fill a tile.")
/// .add_arg("value", "The fill value.")
diff --git a/rust/src/tirx/transform/function_pass.rs
b/rust/src/tirx/transform/function_pass.rs
index 839837b53..6b9e51a4d 100644
--- a/rust/src/tirx/transform/function_pass.rs
+++ b/rust/src/tirx/transform/function_pass.rs
@@ -38,11 +38,11 @@ use crate::ir::transform::core::PassInfo;
use crate::ir::transform::core::PassObj;
use crate::tirx::function::Function;
use std::ops::Deref;
-use tvm_ffi::FieldGetter;
use tvm_ffi::ObjectArc;
use tvm_ffi::ObjectCore;
use tvm_ffi::RValueRef;
use tvm_ffi::Result;
+use tvm_ffi::reflection::FieldGetter;
// tvm-ffi-stubgen(end)
/// A pass that replaces each TIRx function of a module by what a function
returns for it, or
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 93290469d..0a65e206a 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -39,8 +39,8 @@ use tvm_tirx::tvm_ffi::static_init::{StaticInitBlock,
STATIC_INIT_BLOCKS};
use tvm_tirx::tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
use tvm_tirx::tvm_ffi::{
self, rs_object, run_static_init_blocks, Any, AnyView, Array, Error,
ErrorKind, Function, Map,
- ObjectCore, ObjectRefCore, Result, String as FfiString, ATTRIBUTE_ERROR,
INTERNAL_ERROR,
- TYPE_ERROR, VALUE_ERROR,
+ ObjectCore, ObjectRefCore, Result, String as FfiString, ATTRIBUTE_ERROR,
TYPE_ERROR,
+ VALUE_ERROR,
};
use super::{int32, parse_dtype, tensor_var};
@@ -100,10 +100,6 @@ fn only_variables(call: &CallObj) -> Result<()> {
}
}
-fn faulty(_call: &CallObj) -> Result<()> {
- panic!("validator bug")
-}
-
#[test]
fn ops_registered_from_rust_are_used_natively() {
let _serial = serial();
@@ -156,15 +152,12 @@ fn ops_registered_from_rust_are_used_natively() {
assert_eq!(INFERRED.load(Ordering::Relaxed), 3);
validate(&identity).unwrap();
- // An error a Rust hook returns reaches the native caller with its kind
and message, and a
- // panic becomes an error instead of unwinding into C++.
+ // An error a Rust hook returns reaches the native caller with its kind
and message.
let untyped = OpDef::new("rust_test.untyped", "")
.set_infer_type(|_: &CallObj| Err(Error::new(TYPE_ERROR, "no result
type", "")))
.op()
.unwrap();
assert_error(call(&untyped, vec![]), TYPE_ERROR, "no result type");
- let faulty = OpDef::new("rust_test.faulty",
"").set_validator(faulty).op().unwrap();
- assert_error(validate(&call(&faulty, vec![]).unwrap()), INTERNAL_ERROR,
"validator bug");
// A generated validator rejects what the C++ one generated from the same
signature does,
// with the same message.
@@ -444,9 +437,6 @@ tvm_ffi::static_init_block! {
.set_region_body_params(launch_body_params)
.set_region_validate(no_results_or_attributes);
- OpDef::new("rust_test.panicking_region", "")
- .set_region_body_params(|_: &CallObj| -> Result<Array<Var>> {
panic!("hook bug") });
-
register_type::<BlockType>()?;
}
@@ -706,19 +696,4 @@ fn static_init_blocks_define_operators_and_types() {
let cpp = region(&native, tag, extent, params, attrs, results);
assert_eq!(rust, cpp, "{what}");
}
-
- // A panic in a hook becomes an error instead of unwinding into C++.
- let panicking = Op::get("rust_test.panicking_region").unwrap();
- let body = Evaluate::new(0, UNKNOWN_LOC);
- let empty = DictAttrs::new(Map::new());
- let built = RegionStmt::new(
- panicking,
- Vec::<Expr>::new(),
- Vec::<Var>::new(),
- empty,
- body,
- Vec::<Var>::new(),
- UNKNOWN_LOC,
- );
- assert!(built.is_err(), "a panicking hook fails the constructor");
}
diff --git a/rust/tests/ir/type.rs b/rust/tests/ir/type.rs
index 87f756776..7c733cf2c 100644
--- a/rust/tests/ir/type.rs
+++ b/rust/tests/ir/type.rs
@@ -28,10 +28,11 @@ use tvm_tirx::ir::{
OpaqueType, PrimType, PtrType, StringType, TensorRegionType, TupleType,
Type, TypeObj, Var,
VarObj,
};
+use tvm_tirx::tvm_ffi::reflection::{FieldGetter, ShallowCopy};
use tvm_tirx::tvm_ffi::{rs_object, structural_hash, structural_map};
use tvm_tirx::tvm_ffi::{
- Any, AnyView, Array, DLDataType, DLDataTypeExt, FieldGetter, ObjectCore,
ObjectRefCore,
- ShallowCopy, String as FfiString, WalkOrder,
+ Any, AnyView, Array, DLDataType, DLDataTypeExt, ObjectCore, ObjectRefCore,
String as FfiString,
+ WalkOrder,
};
use super::{scope};
diff --git a/rust/tests/runtime/main.rs b/rust/tests/runtime/main.rs
index 7b424072f..57acb6f3c 100644
--- a/rust/tests/runtime/main.rs
+++ b/rust/tests/runtime/main.rs
@@ -25,6 +25,7 @@ mod common;
use tvm_tirx::runtime::{self, MemoryScope};
use tvm_tirx::ir::dtype;
+use tvm_tirx::tvm_ffi::reflection::EnumDef;
use tvm_tirx::tvm_ffi::{DLDevice, DLDeviceType, TypeIndex};
#[test]
@@ -42,8 +43,8 @@ fn tensors() {
#[test]
fn memory_scopes() {
- let scope = MemoryScope::register("rust_test.scratch").unwrap();
- let again = MemoryScope::register("rust_test.scratch");
+ let scope =
EnumDef::<MemoryScope>::new("rust_test.scratch").unwrap().value();
+ let again = EnumDef::<MemoryScope>::new("rust_test.scratch");
assert!(again.is_err());
// The CPU allocator supports the global scope only.
let cpu = DLDevice::new(DLDeviceType::kDLCPU, 0);
diff --git a/rust/tests/tirx/layout.rs b/rust/tests/tirx/layout.rs
index aea687855..1469babae 100644
--- a/rust/tests/tirx/layout.rs
+++ b/rust/tests/tirx/layout.rs
@@ -30,7 +30,7 @@ use crate::common::{global, serial};
#[test]
fn memory_axes() {
let _serial = serial();
- let m = Axis::get("m");
+ let m = Axis::get_or_register("m");
assert!(m.is_memory().unwrap() && !m.is_thread().unwrap());
assert!(m.backend().unwrap().is_none());
assert!(m.get_scope().unwrap().is_none());
@@ -73,7 +73,7 @@ fn layouts() {
fn thread_layout_scope() {
let _serial = serial();
global("tirx.cuda.RegisterLayoutAxes").call_tuple(()).unwrap();
- let laneid = Axis::get("laneid");
+ let laneid = Axis::get_or_register("laneid");
let lanes =
TileLayout::new([Iter::new(int32(32), int32(1), laneid)],
Vec::<Iter>::new(), &Map::new());
let (inner, outer) = lanes.get_scope().unwrap().unwrap();
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index 5d5a43546..85b996172 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -45,6 +45,7 @@ use tvm_tirx::ir::{
};
use tvm_tirx::runtime::MemoryScope;
use tvm_tirx::tirx::{Axis, Function, Iter, TileLayout};
+use tvm_tirx::tvm_ffi::reflection::EnumDef;
use tvm_tirx::tvm_ffi::{Array, Map, ObjectRefCore, String as FfiString};
use common::global;
@@ -85,12 +86,12 @@ fn as_i64(expr: &impl ObjectRefCore) -> i64 {
fn memory_scope() -> MemoryScope {
static COUNT: AtomicUsize = AtomicUsize::new(0);
let index = COUNT.fetch_add(1, Ordering::Relaxed);
- MemoryScope::register(&format!("rust.tests.tirx.{index}")).unwrap()
+
EnumDef::<MemoryScope>::new(&format!("rust.tests.tirx.{index}")).unwrap().value()
}
/// The row-major layout of an 8 x 8 tile.
fn tile_8x8() -> TileLayout {
- let m = Axis::get("m");
+ let m = Axis::get_or_register("m");
TileLayout::new(
[Iter::new(int32(8), int32(8), m.clone()), Iter::new(int32(8),
int32(1), m)],
Vec::<Iter>::new(),
diff --git a/rust/tests/tirx/tensor.rs b/rust/tests/tirx/tensor.rs
index 69e501008..e24ab8a24 100644
--- a/rust/tests/tirx/tensor.rs
+++ b/rust/tests/tirx/tensor.rs
@@ -20,10 +20,10 @@
//! Tests of `tirx::tensor`: tensor types.
use tvm_tirx::ir::{UNKNOWN_LOC, UnknownLoc, Expr, PrimExpr, PrimType,
StringImm};
-use tvm_tirx::tirx::{tensor_is_scalar, Layout, TensorType};
+use tvm_tirx::tirx::{Layout, TensorType};
use tvm_tirx::tvm_ffi::Array;
-use super::{int32, memory_scope, tile_8x8, var};
+use super::{int32, memory_scope, tile_8x8};
use crate::common::{agree, agree_kind, native, serial};
#[test]
@@ -45,8 +45,6 @@ fn tensor_types() {
)
.unwrap();
assert!(laid_out.has_default_layout(None).unwrap());
- assert!(tensor_is_scalar(var("A", &ty), true).unwrap());
- assert!(!tensor_is_scalar(var("B", &laid_out), true).unwrap());
}
#[test]
diff --git a/rust/tests/tirx/thread_scope.rs b/rust/tests/tirx/thread_scope.rs
index d57467390..5b0791794 100644
--- a/rust/tests/tirx/thread_scope.rs
+++ b/rust/tests/tirx/thread_scope.rs
@@ -20,6 +20,7 @@
//! Tests of `tirx::thread_scope`.
use tvm_tirx::tirx::ThreadScope;
+use tvm_tirx::tvm_ffi::EnumRef;
use crate::common::serial;
diff --git a/rust/tests/tirx/tile/dispatch.rs b/rust/tests/tirx/tile/dispatch.rs
index 0414c99c9..59e9bf2b7 100644
--- a/rust/tests/tirx/tile/dispatch.rs
+++ b/rust/tests/tirx/tile/dispatch.rs
@@ -23,7 +23,7 @@ use tvm_tirx::ir::{UnknownLoc, UNKNOWN_LOC, PrimExpr,
PrimType, Stmt, Var};
use tvm_tirx::target::Target;
use tvm_tirx::tirx::tile::DispatchContext;
use tvm_tirx::tirx::{TensorType, ThreadScope};
-use tvm_tirx::tvm_ffi::{Any, Map, ObjectRefCore, String as FfiString};
+use tvm_tirx::tvm_ffi::{Any, EnumRef, Map, ObjectRefCore, String as FfiString};
use crate::common::{global, serial, structural_equal};
use crate::{int32, memory_scope, prim_var, var};
diff --git a/src/tirx/ir/tensor.cc b/src/tirx/ir/tensor.cc
index f8b200210..a20df28c9 100644
--- a/src/tirx/ir/tensor.cc
+++ b/src/tirx/ir/tensor.cc
@@ -418,12 +418,5 @@ ffi::Array<PrimExpr>
TensorTypeObj::ElemOffset(ffi::Array<PrimExpr> input_indice
TensorVar::TensorVar(ffi::String name, TensorType type, Location loc)
: Var(Var(std::move(name), std::move(type), std::move(loc))) {}
-bool TensorVar::IsScalar(bool alloc_or_decl) const { return
ty()->IsScalar(alloc_or_decl); }
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_method("tirx.TensorIsScalar", &TensorVar::IsScalar);
-}
-
} // namespace tirx
} // namespace tvm
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index 4a44b8e91..23ac20a6f 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -60,14 +60,6 @@ bool
TensorTypeObj::HasDefaultLayout(ffi::Optional<sym::Analyzer> analyzer) cons
return ffi::GetRef<TileLayout>(tile).IsContiguous(std::move(analyzer));
}
-bool TensorTypeObj::IsScalar(bool alloc_or_decl) const {
- // TODO(@bohan): logical scope is not considered
- return shape.size() == 1 && tvm::prim::IsOne(shape[0]) && data_alignment ==
64 &&
- (!layout.has_value() ||
- (layout.as<TileLayoutObj>() &&
-
layout.value().Canonicalize().as_or_throw<TileLayout>().IsTrivial()));
-}
-
ffi::Optional<int64_t> TensorTypeObj::ConstantAllocationSize() const {
int64_t result = 1;
for (const PrimExpr& extent : shape) {
diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc
b/src/tirx/transform/force_narrow_index_to_i32.cc
index 54b8ccc06..371917fc5 100644
--- a/src/tirx/transform/force_narrow_index_to_i32.cc
+++ b/src/tirx/transform/force_narrow_index_to_i32.cc
@@ -26,6 +26,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/layout.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
@@ -64,13 +65,20 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer {
}
}
+ /*! \brief Whether `ty` has the layout of TVMScript scalar storage, not
merely one element. */
+ static bool IsScalarTensor(const TensorType& ty) {
+ return ty->shape.size() == 1 && prim::IsOne(ty->shape[0]) &&
ty->data_alignment == 64 &&
+ (!ty->layout.has_value() ||
+ (ty->layout.as<TileLayoutObj>() &&
+
ty->layout.value().Canonicalize().as_or_throw<TileLayout>().IsTrivial()));
+ }
+
/*! \brief Reject allocated integer tensors wider than int32. */
void CheckAllocatedTensor(const TensorVar& buf) const {
// Scalar assignments in TVMScript use local scalar storage. Keep its
explicit
// dtype (e.g. an int64 opaque call result) and cast at narrowed index
uses.
- // IsScalar checks the scalar layout contract, not merely the allocation
size.
static const auto local = runtime::MemoryScope::Get("local");
- bool is_local_scalar = buf.scope() == local && buf.IsScalar();
+ bool is_local_scalar = buf.scope() == local && IsScalarTensor(buf.ty());
if (!is_local_scalar && buf.dtype().MatchesCode(DLDataTypeCode::kDLInt) &&
buf.dtype().bits() > 32) {
TVM_FFI_THROW(InternalError)
