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 8afa23043 [REFACTOR][RUST] Define operators and types in
tvm_ffi::static_init_block! (#105)
8afa23043 is described below
commit 8afa23043008b576f526bd2187a7b5e031585ad4
Author: Linzhang Li <[email protected]>
AuthorDate: Wed Oct 28 08:35:04 2026 -0400
[REFACTOR][RUST] Define operators and types in tvm_ffi::static_init_block!
(#105)
## Summary
This PR defines Rust operators and types the way C++ does. The
definitions are ordinary items, and
a `tvm_ffi::static_init_block!` registers them when the library loads,
the counterpart of
`TVM_FFI_STATIC_INIT_BLOCK`. The crate's own `register_ops!` and
`register_types!` macros, and the
registry behind them, are removed.
```rust
/// The type of a pair of values, printed as `my_ext.Pair(first, second)`.
#[rs_class("my_ext.PairType", final)]
pub struct PairTypeObj {
base: TypeObj,
pub first: Type,
pub second: Type,
}
def_object_ref!(pub PairType, Type, PairTypeObj);
impl Placeholder for PairTypeObj {
fn placeholder() -> Self {
Self::new(TypeObj::placeholder(), Type::placeholder(),
Type::placeholder())
}
}
tvm_ffi::static_init_block! {
OpDef::new("my_ext.clamp", "Clamp an index to the range `[0, bound)`.")
.signature((sig::arg::<IntExpr>("index", ""),
sig::arg::<IntExpr>("bound", "")))
.set_infer_type(index_type)
.op()?;
register_type::<PairType>()?;
}
```
Signed-off-by: Yuchuan <[email protected]>
---
3rdparty/tvm-ffi | 2 +-
AGENTS.md | 4 +-
python/tvm_tirx/backend/rust/__init__.py | 5 +-
rust/Cargo.toml | 1 -
rust/examples/define_ir.rs | 51 +++--
rust/src/ir/op_def.rs | 331 ++-----------------------------
rust/src/ir/type_def.rs | 165 ++++++---------
rust/src/lib.rs | 61 +++---
rust/tests/ir/op_def.rs | 154 ++++++--------
tests/python/ir/test_op_rust.py | 4 +-
10 files changed, 212 insertions(+), 566 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index dbdcdacee..1a2e530fa 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit dbdcdaceeef5b0b03251e8dea87725a1c1e2cde6
+Subproject commit 1a2e530fa68f5a251e3dea8b24d2a7ba19614a14
diff --git a/AGENTS.md b/AGENTS.md
index 2b28a8075..4c2f3188a 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -139,8 +139,8 @@ uv run --no-sync cargo test --manifest-path rust/Cargo.toml
The editable install also builds the crate as `libtvm_tirx_rust` when it finds
a
cargo as new as the crate's `rust-version` (`TVM_TIRX_USE_RUST=AUTO`; rustup's
-cargo is preferred). `import tvm_tirx` loads it and registers the operators of
its
-`register_ops!` blocks, so re-run the editable install after changing them.
+cargo is preferred). `import tvm_tirx` loads it and registers the operators
and types
+of its static init blocks, so re-run the editable install after changing them.
## Linting
diff --git a/python/tvm_tirx/backend/rust/__init__.py
b/python/tvm_tirx/backend/rust/__init__.py
index 38e1e2d6c..2c0640a60 100644
--- a/python/tvm_tirx/backend/rust/__init__.py
+++ b/python/tvm_tirx/backend/rust/__init__.py
@@ -37,8 +37,9 @@ def register_backend():
except (OSError, FileNotFoundError, RuntimeError):
pass
else:
- # Loading the library registers its operators and types, and keeps it,
and their hooks,
- # loaded for the process; ``tirx_register_ir`` raises the error of a
block that failed.
+ # Loading the library runs its static init blocks, which register its
operators and
+ # types, and keeps it, and their hooks, loaded for the process;
``tirx_register_ir``
+ # raises the error of a block that failed.
tvm_ffi.load_module(_LOADED_LIBS["tvm_tirx_rust"]._name).tirx_register_ir()
return None
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 80f96d402..24cdff321 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -30,7 +30,6 @@ name = "tvm_tirx"
crate-type = ["rlib", "cdylib"]
[dependencies]
-linkme = "0.3"
tvm-ffi = { path = "../3rdparty/tvm-ffi/rust/tvm-ffi" }
# A library of operators and a type defined in Rust, which
`tests/python/ir/test_op_rust.py` loads.
diff --git a/rust/examples/define_ir.rs b/rust/examples/define_ir.rs
index 0d1677fd2..76e267258 100644
--- a/rust/examples/define_ir.rs
+++ b/rust/examples/define_ir.rs
@@ -31,9 +31,10 @@
use tvm_tirx::ir::op_attr::{CALL_EFFECT_KIND, FIXED_RETURN_TYPE};
use tvm_tirx::ir::sig::{self, IntExpr};
use tvm_tirx::ir::{
- CallEffectKind, CallObj, DictAttrs, Expr, OpDef, PrimType, Type, TypeObj,
Var, UNKNOWN_LOC,
+ register_type, CallEffectKind, CallObj, DictAttrs, Expr, OpDef,
Placeholder, PrimType, Type,
+ TypeObj, Var, UNKNOWN_LOC,
};
-use tvm_tirx::tvm_ffi::{self, rs_class, Array, ObjectRefCast, Result};
+use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class, Array, ObjectRefCast,
Result};
fn first_arg_type(call: &CallObj) -> Result<Type> {
Ok(call.args.get(0)?.ty.clone())
@@ -45,45 +46,54 @@ fn loop_index(call: &CallObj) -> Result<Array<Var>> {
Ok(Array::new(vec![Var::new("i", ty, UNKNOWN_LOC)]))
}
-tvm_tirx::register_ops! {
+// A type, which Python constructs from its fields, `T.ExamplePair(T.int32,
T.float32)`, as it
+// constructs a C++ type. Its object makes the placeholder `__ffi_new__` fills.
+/// The type of a pair of values.
+#[rs_class("tirx.ExamplePairType", final)]
+pub struct ExamplePairTypeObj {
+ base: TypeObj,
+ pub first: Type,
+ pub second: Type,
+}
+def_object_ref!(pub ExamplePairType, Type, ExamplePairTypeObj);
+
+impl Placeholder for ExamplePairTypeObj {
+ fn placeholder() -> Self {
+ Self::new(TypeObj::placeholder(), Type::placeholder(),
Type::placeholder())
+ }
+}
+
+tvm_ffi::static_init_block! {
OpDef::new("tirx.example_clamp", "Clamp an index to the range `[0,
bound)`.")
.signature((
sig::arg::<IntExpr>("index", "The index."),
sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
))
.set_infer_type(first_arg_type)
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure)
+ .op()?;
// `Call.attrs` holds a `DictAttrs`; outside `tirx.`, the calls print as
`my_ext.hint(..)`.
OpDef::new("my_ext.hint", "Return the value, with hints for the code
generator.")
.signature((sig::arg::<Expr>("value", "The value."),
sig::call_attrs::<DictAttrs>()))
.set_infer_type(first_arg_type)
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation);
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation)
+ .op()?;
// A region operation, which binds an index in its body.
OpDef::new("tirx.example_loop", "Run the body for each index in `[0,
extent)`.")
.signature(sig::arg::<IntExpr>("extent", "The number of iterations."))
- .set_region_body_params(loop_index);
-}
+ .set_region_body_params(loop_index)
+ .op()?;
-// A type, which Python constructs from its fields, `T.ExamplePair(T.int32,
T.float32)`, as it
-// constructs a C++ type.
-tvm_tirx::register_types! {
- /// The type of a pair of values.
- #[rs_class("tirx.ExamplePairType", final)]
- pub struct ExamplePairTypeObj {
- base: TypeObj,
- pub first: Type,
- pub second: Type,
- }
- def_object_ref!(pub ExamplePairType, Type, ExamplePairTypeObj);
+ register_type::<ExamplePairType>()?;
}
// A block can be anywhere in the library.
mod effects {
use super::*;
- tvm_tirx::register_ops! {
+ tvm_ffi::static_init_block! {
OpDef::new("tirx.example_mark", "
Mark a point in the program.
@@ -91,6 +101,7 @@ mod effects {
")
.signature(())
.set_attr(FIXED_RETURN_TYPE, PrimType::void())
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque);
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque)
+ .op()?;
}
}
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 7bc2aa3b5..5122bb02e 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -18,7 +18,7 @@
*/
//! Registering operators from Rust: C++ `OpDef`, with its signature
descriptors ([`sig`]) and
-//! attribute keys ([`op_attr`]); the [`register_ops!`](crate::register_ops)
blocks that define
+//! attribute keys ([`op_attr`]), whose chains a
[`tvm_ffi::static_init_block!`] runs to define
//! operators, as C++ defines them in a `TVM_FFI_STATIC_INIT_BLOCK`; and the
handles that keep
//! operators in `static`s.
//!
@@ -27,11 +27,11 @@
//! packing arguments. A signature generates its validator from its argument
types, as C++
//! `OpDef::signature` does from `sig::arg<T>`; each signature is its own
compiled function.
//!
-//! The [`register_ops!`](crate::register_ops) and
[`register_types!`](crate::register_types)
-//! blocks are gathered at link time, wherever in the program or library they
are, and run when it
-//! loads, as a C++ static init block does: once Python loads a library built
on this crate, by any
-//! means, it sees its operators and types as it sees those of a C++ library.
A block that several
-//! libraries of a process have, as each links its own copy of a crate they
share, runs once.
+//! A [`tvm_ffi::static_init_block!`] is gathered at link time, wherever in
the program or library
+//! it is, and runs when it loads, as a C++ static init block does, so that
once Python loads a
+//! library built on this crate, by any means, it sees the operators and types
the block defines
+//! as it sees those of a C++ library. A block that several libraries of a
process have, as each
+//! links its own copy of a crate they share, runs once.
use std::any::TypeId;
use std::cell::RefCell;
@@ -39,12 +39,11 @@ use std::ffi::c_void;
use std::marker::PhantomData;
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
-use std::sync::{Once, OnceLock};
+use std::sync::OnceLock;
use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFIObject, TVMFFITypeIndex};
use tvm_ffi::{
- Any, AnyView, Array, Error, Function, ObjectArc, ObjectRefCore, Result,
String, INTERNAL_ERROR,
- TYPE_ERROR, VALUE_ERROR,
+ Any, AnyView, Array, Error, ObjectArc, ObjectRefCore, Result, String,
TYPE_ERROR, VALUE_ERROR,
};
use super::attrs::Attrs;
@@ -60,8 +59,8 @@ use sig::{Checks, Specs};
/// Each step updates the registry at once, as in C++. After a failed step the
rest are skipped,
/// and [`OpDef::op`] returns the error. The registry does not synchronize
registration with
/// readers, so register an operator before building IR with it, as a
-/// [`register_ops!`](crate::register_ops) block does. An operator and its
hooks live for the rest
-/// of the process.
+/// [`tvm_ffi::static_init_block!`] does. An operator and its hooks live for
the rest of the
+/// process.
///
/// ```no_run
/// use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
@@ -80,6 +79,7 @@ use sig::{Checks, Specs};
/// .op()?;
/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
/// ```
+#[must_use = "end the chain with `.op()?`, which returns its error and exposes
the operator"]
pub struct OpDef(Result<Op>);
impl OpDef {
@@ -322,23 +322,24 @@ impl From<CallEffectKind> for Any {
/// A handle to an operator, kept in a `static` as C++ keeps one in `const Op&
xxx_op()`. The
/// operator is looked up by name once per thread; [`OpHandle::get`] returns
it and
-/// [`OpHandle::is`] compares with it. The
[`register_ops!`](crate::register_ops) blocks run when
-/// their program or library loads; where the loader does not run them, the
first lookup of an
-/// operator that is not registered runs them.
+/// [`OpHandle::is`] compares with it. The [`tvm_ffi::static_init_block!`]s
that define operators
+/// run when their program or library loads; where the loader does not run
them, the first lookup
+/// of an operator that is not registered runs them, with
[`tvm_ffi::run_static_init_blocks`].
///
/// ```no_run
/// use tvm_tirx::ir::sig::{self, IntExpr};
/// use tvm_tirx::ir::{Call, CallObj, Expr, OpDef, OpHandle, PrimType, Type,
Var, UNKNOWN_LOC};
-/// use tvm_tirx::tvm_ffi::Result;
+/// use tvm_tirx::tvm_ffi::{self, Result};
///
/// fn first_arg_type(call: &CallObj) -> Result<Type> {
/// Ok(call.args.get(0)?.ty.clone())
/// }
///
-/// tvm_tirx::register_ops! {
+/// tvm_ffi::static_init_block! {
/// OpDef::new("my_ext.clamp", "Clamp an index to a bound.")
/// .signature((sig::arg::<IntExpr>("index", ""),
sig::arg::<IntExpr>("bound", "")))
-/// .set_infer_type(first_arg_type);
+/// .set_infer_type(first_arg_type)
+/// .op()?;
/// }
///
/// static CLAMP: OpHandle = OpHandle::new("my_ext.clamp");
@@ -369,8 +370,7 @@ thread_local! {
}
impl OpHandle {
- /// The operator `name`, registered in C++, in Python or by a
- /// [`register_ops!`](crate::register_ops) block.
+ /// The operator `name`, registered in C++, in Python or by a Rust
[`OpDef`].
pub const fn new(name: &'static str) -> Self {
Self { name, slot: OnceLock::new(), object:
AtomicPtr::new(std::ptr::null_mut()) }
}
@@ -413,7 +413,7 @@ impl OpHandle {
Some(op) => op,
None => {
let op = Op::get(self.name).or_else(|missing| {
- let registered = register_defined_ir();
+ let registered = tvm_ffi::run_static_init_blocks();
Op::get(self.name).map_err(|_| match registered {
Ok(()) => missing,
Err(failed) => {
@@ -436,209 +436,6 @@ impl OpHandle {
}
}
-/// A [`register_ops!`](crate::register_ops) or
[`register_types!`](crate::register_types) block,
-/// gathered at link time.
-#[doc(hidden)]
-pub struct Registration {
- /// The macro of the block, and its module and line.
- kind: &'static str,
- module: &'static str,
- line: u32,
- /// The version of the block's crate and a hash of the block's text, which
with the module and
- /// line identify the block in every library that links the crate.
- version: &'static str,
- text_hash: u64,
- register: fn() -> Result<()>,
- /// The outcome of the block, which runs once: the error's text if it
failed.
- outcome: OnceLock<std::result::Result<(), std::string::String>>,
-}
-
-/// The blocks of [`register_ops!`](crate::register_ops) and
-/// [`register_types!`](crate::register_types), gathered at link time from the
whole program or
-/// library.
-#[doc(hidden)]
-#[linkme::distributed_slice]
-pub static REGISTRATIONS: [Registration];
-
-impl Registration {
- #[doc(hidden)]
- pub const fn __new(
- kind: &'static str,
- module: &'static str,
- line: u32,
- version: &'static str,
- text: &str,
- register: fn() -> Result<()>,
- ) -> Self {
- let text_hash = text_hash(text);
- Self { kind, module, line, version, text_hash, register, outcome:
OnceLock::new() }
- }
-
- /// The identity of the block in the process, `module:line@version#hash`:
the same in every
- /// library that links the block's crate, and another one if the crate
version or the block's
- /// text differs.
- #[doc(hidden)]
- pub fn identity(&self) -> std::string::String {
- format!("{}:{}@{}#{:016x}", self.module, self.line, self.version,
self.text_hash)
- }
-
- /// The name of the global function that marks the block registered in the
process,
- /// `ir.registered_block.<identity>`. The libraries of a process share the
global function
- /// table of TVM FFI, and their statics are their own.
- #[doc(hidden)]
- pub fn marker(&self) -> std::string::String {
- format!("ir.registered_block.{}", self.identity())
- }
-
- /// Run the block if it has not run. The caller that runs a failing block
gets its error;
- /// later ones get an error that names the block.
- #[doc(hidden)]
- pub fn register(&self) -> Result<()> {
- let mut failure = None;
- let outcome = self.outcome.get_or_init(|| {
- self.run().map_err(|error| {
- let text = format!("{}: {}", error.kind(), error.message());
- failure = Some(error);
- text
- })
- });
- match (failure, outcome) {
- (Some(error), _) => Err(error),
- (None, Ok(())) => Ok(()),
- (None, Err(text)) => {
- let message = format!("The {} block at {} failed: {text}",
self.kind, self.place());
- Err(Error::new(INTERNAL_ERROR, &message, ""))
- }
- }
- }
-
- /// Run the block unless another library of the process has run it. Every
library links its
- /// own copy of the crates it is built from, so the same block can be in
several, and what it
- /// defines is registered once, with the hooks of the first library that
runs it. The block
- /// is marked once it succeeds, so that one that fails fails in every
library.
- fn run(&self) -> Result<()> {
- let marker = self.marker();
- if Function::get_global(&marker).is_err() {
- (self.register)()?;
- Function::register_global(&marker, Function::from_packed(|_|
Ok(Any::default())))?;
- }
- Ok(())
- }
-
- /// Where the block is, for an error.
- fn place(&self) -> std::string::String {
- format!("{}:{}", self.module, self.line)
- }
-}
-
-/// The FNV-1a hash of `text` without its whitespace, which does not depend on
how the compiler
-/// spaces the tokens of a block it stringifies.
-const fn text_hash(text: &str) -> u64 {
- let bytes = text.as_bytes();
- let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
- let mut i = 0;
- while i < bytes.len() {
- if !bytes[i].is_ascii_whitespace() {
- hash = (hash ^ bytes[i] as u64).wrapping_mul(0x0100_0000_01b3);
- }
- i += 1;
- }
- hash
-}
-
-/// Run every [`register_ops!`](crate::register_ops) and
[`register_types!`](crate::register_types)
-/// block of this program or library that has not run, in the order of their
modules and lines;
-/// running again does nothing. Every block is tried, and the error lists
those that failed. A
-/// block that another library of the process has registered, from the same
version of its crate
-/// and the same text, is skipped: libraries that link the same crate register
its operators and
-/// types once.
-///
-/// The blocks run when their program or library loads, so this reports the
errors of those that
-/// failed. Every library built on this crate exports it as the TVM FFI
function
-/// `tirx_register_ir`, which `import tvm_tirx` calls on `libtvm_tirx_rust` to
raise them. Where
-/// the loader does not run the blocks, the first lookup through an
[`OpHandle`] of an operator
-/// that is not registered calls it.
-pub fn register_defined_ir() -> Result<()> {
- let mut blocks: Vec<&Registration> = REGISTRATIONS.iter().collect();
- blocks.sort_by_key(|block| (block.module, block.line));
- let mut failures: Vec<(&Registration, Error)> =
- blocks.into_iter().filter_map(|block| block.register().err().map(|e|
(block, e))).collect();
- if failures.len() <= 1 {
- return failures.pop().map_or(Ok(()), |(_, error)| Err(error));
- }
- let mut kinds: Vec<&str> = failures.iter().map(|(block, _)|
block.kind).collect();
- kinds.sort();
- kinds.dedup();
- let mut message = format!("{} {} blocks failed:", failures.len(),
kinds.join(" and "));
- for (block, error) in &failures {
- message.push_str(&format!("\n {}: {}: {}", block.place(),
error.kind(), error.message()));
- }
- Err(Error::new(failures[0].1.kind(), &message, ""))
-}
-
-/// The constructor of every [`register_ops!`](crate::register_ops) and
-/// [`register_types!`](crate::register_types) block, which the loader runs
when the program or
-/// library loads: the first one keeps the library loaded and runs every block
of it. An error
-/// cannot reach a caller here, so it is printed; [`register_defined_ir`]
reports it again.
-#[doc(hidden)]
-pub extern "C" fn __register_at_load() {
- static LOADED: Once = Once::new();
- LOADED.call_once(|| {
- let library = keep_loaded().unwrap_or_else(||
"register_ops!".to_string());
- match catch_unwind(register_defined_ir) {
- Ok(Ok(())) => {}
- Ok(Err(error)) => eprintln!(
- "{library}: registering operators and types failed: {}: {}",
- error.kind(),
- error.message()
- ),
- Err(_) => eprintln!("{library}: registering operators and types
panicked"),
- }
- });
-}
-
-/// Keep the library that holds this code loaded for the rest of the process,
as the operators and
-/// types it registers keep its hooks: unloading it, as `tvm_ffi.load_module`
does when Python drops the
-/// module, would leave them dangling. Returns the path of the library.
-#[cfg(any(target_os = "linux", target_os = "macos"))]
-fn keep_loaded() -> Option<std::string::String> {
- use std::ffi::{c_char, c_int, CStr};
-
- #[repr(C)]
- struct DlInfo {
- fname: *const c_char,
- fbase: *mut c_void,
- sname: *const c_char,
- saddr: *mut c_void,
- }
- extern "C" {
- fn dladdr(addr: *const c_void, info: *mut DlInfo) -> c_int;
- fn dlopen(filename: *const c_char, flags: c_int) -> *mut c_void;
- }
- // RTLD_LAZY | RTLD_NOLOAD | RTLD_NODELETE.
- #[cfg(target_os = "linux")]
- const FLAGS: c_int = 0x1 | 0x4 | 0x1000;
- #[cfg(target_os = "macos")]
- const FLAGS: c_int = 0x1 | 0x10 | 0x80;
-
- let null = std::ptr::null_mut();
- let mut info = DlInfo { fname: null as _, fbase: null, sname: null as _,
saddr: null };
- // SAFETY: `dladdr` fills `info` for an address of this library, whose
path stays valid while
- // the library is loaded; `dlopen` of a loaded library only marks it, and
its handle is kept.
- unsafe {
- if dladdr(keep_loaded as *const c_void, &mut info) == 0 ||
info.fname.is_null() {
- return None;
- }
- dlopen(info.fname, FLAGS);
- Some(CStr::from_ptr(info.fname).to_string_lossy().into_owned())
- }
-}
-
-#[cfg(not(any(target_os = "linux", target_os = "macos")))]
-fn keep_loaded() -> Option<std::string::String> {
- None
-}
-
/// 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 {
@@ -653,94 +450,6 @@ fn doc_text(doc: &str) -> std::string::String {
lines.join("\n").trim().to_string()
}
-/// Define operators as C++ does in a `TVM_FFI_STATIC_INIT_BLOCK`: each item
is an [`OpDef`]
-/// chain, which the block runs in order, stopping at the first that fails.
-///
-/// ```no_run
-/// use tvm_tirx::ir::op_attr::{CALL_EFFECT_KIND, FIXED_RETURN_TYPE};
-/// use tvm_tirx::ir::sig::{self, IntExpr};
-/// use tvm_tirx::ir::{CallEffectKind, CallObj, Op, OpDef, PrimType, Type};
-/// use tvm_tirx::tvm_ffi::Result;
-///
-/// fn first_arg_type(call: &CallObj) -> Result<Type> {
-/// Ok(call.args.get(0)?.ty.clone())
-/// }
-///
-/// tvm_tirx::register_ops! {
-/// OpDef::new("my_ext.clamp", "Clamp an index to a bound.")
-/// .signature((sig::arg::<IntExpr>("index", ""),
sig::arg::<IntExpr>("bound", "")))
-/// .set_infer_type(first_arg_type)
-/// .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
-///
-/// OpDef::new("my_ext.mark", "
-/// Mark a point in the program.
-///
-/// The documentation is dedented.
-/// ")
-/// .set_attr(FIXED_RETURN_TYPE, PrimType::void())
-/// .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque);
-/// }
-///
-/// // Registered when the program loaded.
-/// assert_eq!(Op::get("my_ext.clamp")?.doc.as_str(), "Clamp an index to a
bound.");
-/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-/// ```
-///
-/// As a C++ static init block, a block runs when its program or library
loads: it adds a
-/// constructor, which the loader runs, and the first constructor runs every
block of the program
-/// or library with [`register_defined_ir`]. Python sees the operators of a
library once it loads
-/// it, by any means, and the library stays loaded for the process, as the
operators keep its
-/// hooks. An error cannot reach a caller when the library loads: it is
printed, and
-/// [`register_defined_ir`] (`tirx_register_ir`) and the handles report it
again. Where the
-/// loader does not run the constructor, an [`OpHandle`] runs the blocks on
the first lookup of an
-/// operator that is not registered.
-///
-/// A block runs once per process. Every library built on the crate links its
own copy of the
-/// crates it uses, so a block of a crate that several libraries link is in
each of them; the first
-/// library that runs it registers its operators, and the others skip it. A
block of another
-/// version of the crate, or whose text differs, is another block, and
registering its operators
-/// again is an error. A change to a hook's body alone does not make another
block: rebuild every
-/// library that links a crate after changing it.
-#[macro_export]
-macro_rules! register_ops {
- ($($def:expr);* $(;)?) => {
- $crate::__register_block!("register_ops!", stringify!($($def);*), || {
- $($crate::ir::OpDef::op($def)?;)*
- Ok(())
- });
- };
-}
-
-/// The registration of a [`register_ops!`](crate::register_ops) or
-/// [`register_types!`](crate::register_types) block: gathered at link time,
and run by a
-/// constructor when its program or library loads.
-#[doc(hidden)]
-#[macro_export]
-macro_rules! __register_block {
- ($kind:literal, $text:expr, $register:expr) => {
- const _: () = {
-
#[$crate::__private::linkme::distributed_slice($crate::ir::REGISTRATIONS)]
- #[linkme(crate = $crate::__private::linkme)]
- static REGISTRATION: $crate::ir::Registration =
$crate::ir::Registration::__new(
- $kind,
- module_path!(),
- line!(),
- match option_env!("CARGO_PKG_VERSION") {
- Some(version) => version,
- None => "",
- },
- $text,
- $register,
- );
- #[used]
- #[cfg_attr(target_os = "linux", link_section = ".init_array")]
- #[cfg_attr(target_os = "macos", link_section =
"__DATA,__mod_init_func")]
- #[cfg_attr(windows, link_section = ".CRT$XCU")]
- static REGISTER_AT_LOAD: extern "C" fn() =
$crate::ir::__register_at_load;
- };
- };
-}
-
/// 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;
diff --git a/rust/src/ir/type_def.rs b/rust/src/ir/type_def.rs
index ef9852d84..d68dd4231 100644
--- a/rust/src/ir/type_def.rs
+++ b/rust/src/ir/type_def.rs
@@ -17,16 +17,16 @@
* under the License.
*/
-//! Registering IR types from Rust: the
[`register_types!`](crate::register_types) blocks that
-//! define `rs_class` types, as C++ defines types with `refl::ObjectDef` in a
-//! `TVM_FFI_STATIC_INIT_BLOCK`, and the [`Placeholder`] values of their
fields.
+//! Registering IR types from Rust: [`register_type`], which a
[`tvm_ffi::static_init_block!`]
+//! calls for an `rs_class` type, as C++ registers a type with
`refl::ObjectDef` in a
+//! `TVM_FFI_STATIC_INIT_BLOCK`, and the [`Placeholder`] values of its fields.
//!
//! `rs_class` registers the fields of a type, but not a way to make one. C++
`refl::ObjectDef`
//! also registers a creator of an empty object, which tvm-ffi fills through
the reflection of the
//! fields to construct the type from Python (the `__ffi_init__` it generates)
and to read it from
-//! a JSON graph. A Rust object cannot be empty, as its references are not
null, so a block
-//! registers `__ffi_new__`, which makes an object whose fields hold
[`Placeholder`]s, and tvm-ffi
-//! then sets each field as it sets those of a C++ object.
+//! a JSON graph. A Rust object cannot be empty, as its references are not
null, so
+//! [`register_type`] registers `__ffi_new__`, which makes an object whose
fields hold
+//! [`Placeholder`]s, and tvm-ffi then sets each field as it sets those of a
C++ object.
use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeRegisterAttr};
use tvm_ffi::{Any, AnyCompatible, Array, Function, ObjectCore, ObjectRefCore,
Result, String};
@@ -36,11 +36,12 @@ use super::location::Location;
use super::r#type::{MissingType, PrimType, Type, TypeObj};
use super::support;
-/// A valid value of a field type, which an object of a
[`register_types!`](crate::register_types)
-/// block holds until reflection sets the field: `__ffi_new__` makes an object
with one in every
-/// field, and a field `__ffi_init__` does not take (`#[reflect(init_off)]`)
keeps it. A block
-/// implements it for the objects and references it defines, and this crate
for the common field
-/// types; implement it for a field type of your own crate.
+/// A valid value of a field type, which an object of a type [`register_type`]
registers holds
+/// until reflection sets the field: `__ffi_new__` makes an object with one in
every field, and a
+/// field `__ffi_init__` does not take (`#[reflect(init_off)]`) keeps it. The
object of a type you
+/// register implements it, from the placeholders of its parent's object and
its fields; this
+/// crate implements it for the common field types, and a field type of your
own crate implements
+/// it too.
pub trait Placeholder {
/// The value.
fn placeholder() -> Self;
@@ -106,18 +107,59 @@ impl Placeholder for Expr {
}
}
-/// Register the type of the reference `R` of a
[`register_types!`](crate::register_types) block:
-/// `__ffi_new__`, then the `__ffi_init__` tvm-ffi generates from the
reflected fields, then, once
-/// the Python receiver is available, request its constructor in registered
destinations
-/// (`ir.expose_type`).
-/// An absent receiver is a no-op; an installed receiver's error is returned.
-#[doc(hidden)]
-pub fn __register_type<R>() -> Result<()>
+/// Register the IR type of the reference `R`, an `rs_class` type, as C++
registers one with
+/// `refl::ObjectDef`: `__ffi_new__`, from the [`Placeholder`] of the object,
then the
+/// `__ffi_init__` tvm-ffi generates from the reflected fields, then, when the
receiver is
+/// installed, a request for its constructor in registered destinations
(`ir.expose_type`); an
+/// absent receiver is a no-op, and an installed receiver's error is returned.
With an explicitly
+/// declared `my_ext` type API destination, Python then constructs the type
from its fields,
+/// `my_ext.Pair(T.int32, second=T.float32)`, as it constructs a C++ type, and
a JSON graph holds
+/// it. Import `tvm_tirx` before loading the library so exposure requests
reach Python. Import the
+/// extension module in Python scope, or supply it to text parsing through
imports, `extra_vars` or
+/// explicit parser namespace registration; type exposure does not add parser
defaults.
+///
+/// Call it in a [`tvm_ffi::static_init_block!`], with the
[`OpDef`](crate::ir::OpDef) chains of
+/// the operators, so that the type is registered when its program or library
loads, as a C++ one
+/// is; the type's struct, reference and `impl`s are written outside the
block, as any other:
+///
+/// ```no_run
+/// use tvm_tirx::ir::{register_type, Location, Placeholder, Type, TypeObj};
+/// use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class, ObjectRefCore};
+///
+/// /// The type of a pair of values, printed as `my_ext.Pair(first, second)`.
+/// #[rs_class("my_ext.PairType", final)]
+/// pub struct PairTypeObj {
+/// base: TypeObj,
+/// pub first: Type,
+/// #[reflect(kw_only)]
+/// pub second: Type,
+/// }
+/// def_object_ref!(pub PairType, Type, PairTypeObj);
+///
+/// impl Placeholder for PairTypeObj {
+/// fn placeholder() -> Self {
+/// Self::new(TypeObj::placeholder(), Type::placeholder(),
Type::placeholder())
+/// }
+/// }
+///
+/// impl PairType {
+/// pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
+/// PairTypeObj::new(TypeObj::new(Location::unknown()), first,
second).into()
+/// }
+/// }
+///
+/// tvm_ffi::static_init_block! {
+/// register_type::<PairType>()?;
+/// }
+/// ```
+pub fn register_type<R>() -> Result<()>
where
- R: ObjectRefCore + AnyCompatible + Placeholder + 'static,
+ R: ObjectRefCore + AnyCompatible + From<R::ContainerType> + 'static,
+ R::ContainerType: Placeholder,
{
let type_index = R::ContainerType::type_index();
- let new = Function::from_typed(|| -> Result<R> { Ok(R::placeholder()) });
+ let new =
+ Function::from_typed(|| -> Result<R> {
Ok(R::from(R::ContainerType::placeholder())) });
register_type_attr(type_index, "__ffi_new__", Any::from(new))?;
tvm_ffi::cached_global_func!("ffi._RegisterFFIInit").call_tuple((type_index,))?;
if let Ok(expose) = Function::get_global("ir.expose_type") {
@@ -134,86 +176,3 @@ fn register_type_attr(type_index: i32, name: &str, mut
value: Any) -> Result<()>
TVMFFITypeRegisterAttr(type_index, &name, Any::as_data_ptr(&mut value))
})
}
-
-/// Define IR types as C++ does with `refl::ObjectDef` in a
`TVM_FFI_STATIC_INIT_BLOCK`: each item
-/// is an `rs_class` object struct and the `def_object_ref!` of its reference,
written as outside a
-/// block.
-///
-/// ```no_run
-/// use tvm_tirx::ir::{Type, TypeObj};
-/// use tvm_tirx::tvm_ffi::{self, rs_class};
-///
-/// tvm_tirx::register_types! {
-/// /// The type of a pair of values, printed as `my_ext.Pair(first,
second)`.
-/// #[rs_class("my_ext.PairType", final)]
-/// pub struct PairTypeObj {
-/// base: TypeObj,
-/// pub first: Type,
-/// #[reflect(kw_only)]
-/// pub second: Type,
-/// }
-/// def_object_ref!(pub PairType, Type, PairTypeObj);
-/// }
-///
-/// ```
-///
-/// The block registers its types when its program or library loads, as a
-/// [`register_ops!`](crate::register_ops) block registers its operators, and
once per process.
-/// With an explicitly declared `my_ext` type API destination, Python
constructs a type from its
-/// fields as `my_ext.Pair(T.int32, second=T.float32)`, like a C++ type, and a
JSON graph holds it.
-/// Import `tvm_tirx` before loading the library so exposure requests reach
Python. Import the
-/// extension module in Python scope, or supply it to text parsing through
imports, `extra_vars`
-/// or explicit parser namespace registration; type exposure does not add
parser defaults.
-/// The type of each field and the parent's object implement [`Placeholder`].
-#[macro_export]
-macro_rules! register_types {
- ($(
- $(#[$attr:meta])*
- $vis:vis struct $obj:ident {
- $base:ident: $base_ty:ty
- $(, $(#[$field_attr:meta])* $field_vis:vis $field:ident:
$field_ty:ty)* $(,)?
- }
- def_object_ref!($ref_vis:vis $ref:ident, $parent:ty, $ref_obj:ident
$(,)?);
- )*) => {
- $(
- $(#[$attr])*
- $vis struct $obj {
- $base: $base_ty,
- $($(#[$field_attr])* $field_vis $field: $field_ty,)*
- }
-
- $crate::tvm_ffi::def_object_ref!($ref_vis $ref, $parent, $ref_obj);
-
- impl $crate::ir::Placeholder for $obj {
- fn placeholder() -> Self {
- $obj::new(
- <$base_ty as $crate::ir::Placeholder>::placeholder(),
- $(<$field_ty as
$crate::ir::Placeholder>::placeholder(),)*
- )
- }
- }
-
- impl $crate::ir::Placeholder for $ref {
- fn placeholder() -> Self {
- <$obj as $crate::ir::Placeholder>::placeholder().into()
- }
- }
- )*
-
- $crate::__register_block!(
- "register_types!",
- stringify!($(
- $(#[$attr])*
- $vis struct $obj {
- $base: $base_ty
- $(, $(#[$field_attr])* $field_vis $field: $field_ty)*
- }
- def_object_ref!($ref_vis $ref, $parent, $ref_obj);
- )*),
- || {
- $($crate::ir::__register_type::<$ref>()?;)*
- Ok(())
- }
- );
- };
-}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index ed1c902f3..d82481dc5 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -99,30 +99,44 @@
//! `floordiv(&x + 1, 4)?` is the C++ `floordiv(x + 1, 4)`. An `Ex` passes
into any function or
//! constructor that takes an [`ir::IntoExpr`], as any expression does.
//!
-//! Operators are defined from Rust in [`register_ops!`] blocks, as C++
defines them in a
-//! `TVM_FFI_STATIC_INIT_BLOCK`: each item is an [`ir::OpDef`] chain, the C++
`OpDef`, with the
-//! signature ([`ir::sig`], whose descriptors generate the validator), the
attributes, by the keys
-//! of [`ir::op_attr`] and [`tirx::op_attr`], and the type inference,
validator and region hooks as
-//! Rust functions, which C++ calls natively:
+//! Operators and types are defined from Rust in a
[`tvm_ffi::static_init_block!`], as C++ defines
+//! them in a `TVM_FFI_STATIC_INIT_BLOCK`. An operator is an [`ir::OpDef`]
chain, the C++ `OpDef`,
+//! with the signature ([`ir::sig`], whose descriptors generate the
validator), the attributes, by
+//! the keys of [`ir::op_attr`] and [`tirx::op_attr`], and the type inference,
validator and
+//! region hooks as Rust functions, which C++ calls natively; a type, written
outside the block as
+//! any other ([Defining a type in Rust](#defining-a-type-in-rust)), is
registered by
+//! [`ir::register_type`]:
//!
//! ```no_run
//! use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
//! use tvm_tirx::ir::sig::{self, IntExpr};
-//! use tvm_tirx::ir::{CallEffectKind, CallObj, OpDef, OpHandle, Type};
-//! use tvm_tirx::tvm_ffi::Result;
+//! use tvm_tirx::ir::{register_type, CallEffectKind, CallObj, OpDef,
OpHandle, Type};
+//! use tvm_tirx::tvm_ffi::{self, Result};
+//! # use tvm_tirx::ir::{Placeholder, TypeObj};
+//! # use tvm_tirx::tvm_ffi::{def_object_ref, rs_class, ObjectRefCore};
+//! # #[rs_class("my_ext.PairType", final)]
+//! # pub struct PairTypeObj { base: TypeObj, pub first: Type, pub second:
Type }
+//! # def_object_ref!(pub PairType, Type, PairTypeObj);
+//! # impl Placeholder for PairTypeObj {
+//! # fn placeholder() -> Self {
+//! # Self::new(TypeObj::placeholder(), Type::placeholder(),
Type::placeholder())
+//! # }
+//! # }
//!
//! fn index_type(call: &CallObj) -> Result<Type> {
//! Ok(call.args.get(0)?.ty.clone())
//! }
//!
-//! tvm_tirx::register_ops! {
+//! tvm_ffi::static_init_block! {
//! OpDef::new("my_ext.clamp", "Clamp an index to the range `[0, bound)`.")
//! .signature((
//! sig::arg::<IntExpr>("index", "The index."),
//! sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
//! ))
//! .set_infer_type(index_type)
-//! .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
+//! .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure)
+//! .op()?;
+//! register_type::<PairType>()?;
//! }
//!
//! static CLAMP: OpHandle = OpHandle::new("my_ext.clamp");
@@ -130,8 +144,8 @@
//!
//! An [`ir::OpHandle`] in a `static` keeps an operator as C++ keeps one in
`const Op& xxx_op()`.
//! A block runs when its program or library loads, as a C++ one does,
wherever in it the block
-//! is. Every library built on the crate exports [`ir::register_defined_ir`],
which reports the
-//! errors of the blocks that failed, as the TVM FFI function
`tirx_register_ir`.
+//! is. Every library built on the crate exports
[`tvm_ffi::run_static_init_blocks`], which
+//! reports the errors of the blocks that failed, as the TVM FFI function
`tirx_register_ir`.
//!
//! # The library Python loads
//!
@@ -142,8 +156,8 @@
//! crate defines by the time the import returns, as the C++ ones are, and
calls `tirx_register_ir`
//! to raise the error of a block that failed. Python uses the operators as it
uses the C++ ones:
//! `Op.get` finds them, a `Call` built in Python runs their hooks, and the
TVMScript printer prints
-//! their calls. Completing an operator or a type of a
[`register_types!`](crate::register_types)
-//! block requests API exposure through the `ir.expose_op` or `ir.expose_type`
FFI receiver.
+//! their calls. Completing an operator with [`ir::OpDef::op`] or a type with
[`ir::register_type`]
+//! requests API exposure through the `ir.expose_op` or `ir.expose_type` FFI
receiver.
//! Python installs these receivers before loading native libraries and
publishes the selected
//! definitions to registered destinations, including destinations initialized
later. Op API
//! destinations construct raw calls; script destinations add builder
behavior. Type destinations
@@ -202,9 +216,10 @@
//! The generated code names `tvm_ffi`, which `use tvm_tirx::tvm_ffi` brings
into scope in a crate
//! that depends only on `tvm-tirx`. The parent's object must have a public
`new` and implement
//! `ShallowCopy`; the bindings generate both for a complete type that is not
final, as `ir.Type`.
-//! Written with its `def_object_ref!` in a
[`register_types!`](crate::register_types) block, the
-//! type is also registered with the operators, so that Python constructs it
from its fields and
-//! TVMScript parses it, as a C++ type.
+//! Registered by [`ir::register_type`] in a [`tvm_ffi::static_init_block!`],
with the operators,
+//! the type is also constructed by Python from its fields and parsed by
TVMScript, as a C++ type;
+//! its object must implement [`ir::Placeholder`], written from the
placeholders of the parent's
+//! object and the fields.
//!
//! # Generated and hand-written code
//!
@@ -305,13 +320,7 @@ pub mod tirx;
pub use tvm_ffi;
-// The errors of the 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.
-tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir,
ir::register_defined_ir);
-
-/// What the exported macros expand to.
-#[doc(hidden)]
-pub mod __private {
- pub use linkme;
-}
+// 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.
+tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir,
tvm_ffi::run_static_init_blocks);
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 0601595bb..f0a66987f 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -17,10 +17,10 @@
* under the License.
*/
-//! Tests of `ir::op_def`: operators registered from Rust, by an `OpDef` or in
`register_ops!`
-//! blocks, which `register_defined_ir` or the first lookup through an
`OpHandle` runs, are used by
+//! Tests of `ir::op_def`: operators registered from Rust, by an `OpDef` or in
static init blocks,
+//! which run when the test binary loads or at the first lookup through an
`OpHandle`, are used by
//! the native registry, `Call` constructor and validation, and their type
inference, validators
-//! and region contracts are native hooks.
+//! and region contracts are native hooks; and types registered with them.
//!
//! The registry does not synchronize registration, so the tests take
`serial()`; the tests
//! register operators of different names.
@@ -31,19 +31,20 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use tvm_tirx::ir::op_attr::{CALL_EFFECT_KIND, FIXED_RETURN_TYPE};
use tvm_tirx::ir::sig::{self, IntExpr, PrimExpr, PrimVar};
use tvm_tirx::ir::{
- register_defined_ir, AttrKey, Attrs, Call, CallEffectKind, CallObj,
DictAttrs, Evaluate, Expr,
- FuncType, IntImm, Op, OpDef, OpHandle, PrimType, RegionStmt,
RegionStmtObj, Registration,
- StringImm, Type, Var, VarObj, REGISTRATIONS, UNKNOWN_LOC,
+ register_type, AttrKey, Attrs, Call, CallEffectKind, CallObj, DictAttrs,
Evaluate, Expr,
+ FuncType, IntImm, Op, OpDef, OpHandle, Placeholder, PrimType, RegionStmt,
RegionStmtObj,
+ StringImm, Type, TypeObj, Var, VarObj, UNKNOWN_LOC,
};
use tvm_tirx::tirx::{CallFFIKernelAttr, TensorVar};
+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::{
- Any, AnyView, Array, Error, ErrorKind, Function, Map, ObjectCore,
ObjectRefCast, ObjectRefCore,
- Result, String as FfiString, ATTRIBUTE_ERROR, INTERNAL_ERROR, TYPE_ERROR,
VALUE_ERROR,
+ self, def_object_ref, rs_class, run_static_init_blocks, Any, AnyView,
Array, Error, ErrorKind,
+ Function, Map, ObjectCore, ObjectRefCast, ObjectRefCore, Result, String as
FfiString,
+ ATTRIBUTE_ERROR, INTERNAL_ERROR, TYPE_ERROR, VALUE_ERROR,
};
use super::{int32, parse_dtype, tensor_var};
-use types::{BlockType, BlockTypeObj};
use crate::common::{global, serial, structural_equal};
fn var(name: &str, dtype: &str) -> Var {
@@ -413,7 +414,23 @@ fn accept(_call: &CallObj) -> Result<()> {
Ok(())
}
-tvm_tirx::register_ops! {
+/// A type and a count, registered with the operators when the test binary
loads.
+#[rs_class("rust_test.BlockType", final)]
+pub struct BlockTypeObj {
+ base: TypeObj,
+ pub elem: Type,
+ #[reflect(kw_only)]
+ pub count: i64,
+}
+def_object_ref!(pub BlockType, Type, BlockTypeObj);
+
+impl Placeholder for BlockTypeObj {
+ fn placeholder() -> Self {
+ Self::new(TypeObj::placeholder(), Type::placeholder(),
i64::placeholder())
+ }
+}
+
+tvm_ffi::static_init_block! {
OpDef::new("rust_test.defined_identity", "
Return the value.
@@ -421,73 +438,61 @@ tvm_tirx::register_ops! {
")
.signature(sig::arg::<Expr>("value", "The value."))
.set_infer_type(identity_type)
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation);
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation)
+ .op()?;
// `tirx.launch_thread`'s region contract (`src/tirx/op/region.cc`),
written in Rust: one
// thread index of the extent's type.
OpDef::new("rust_test.launch", "Bind a thread index below an extent.")
.signature((sig::arg::<StringImm>("tag", ""),
sig::arg::<IntExpr>("extent", "")))
.set_region_body_params(launch_body_params)
- .set_region_validate(no_results_or_attributes);
+ .set_region_validate(no_results_or_attributes)
+ .op()?;
OpDef::new("rust_test.panicking_region", "")
- .set_region_body_params(|_: &CallObj| -> Result<Array<Var>> {
panic!("hook bug") });
+ .set_region_body_params(|_: &CallObj| -> Result<Array<Var>> {
panic!("hook bug") })
+ .op()?;
+
+ register_type::<BlockType>()?;
}
mod elsewhere {
use super::*;
- tvm_tirx::register_ops! {
+ tvm_ffi::static_init_block! {
OpDef::new("rust_test.defined_elsewhere", "An operator in another
module.")
- .set_attr(FIXED_RETURN_TYPE, PrimType::void());
- }
-}
-
-// A type of a `register_types!` block, registered with the operators when the
test binary loads.
-mod types {
- use tvm_tirx::ir::{Type, TypeObj};
- use tvm_tirx::tvm_ffi::{self, rs_class};
-
- tvm_tirx::register_types! {
- /// A type and a count.
- #[rs_class("rust_test.BlockType", final)]
- pub struct BlockTypeObj {
- base: TypeObj,
- pub elem: Type,
- #[reflect(kw_only)]
- pub count: i64,
- }
- def_object_ref!(pub BlockType, Type, BlockTypeObj);
+ .set_attr(FIXED_RETURN_TYPE, PrimType::void())
+ .op()?;
}
}
-// Two blocks that fail, which `register_defined_ir` reports together while it
runs the others: a
+// Two blocks that fail, which `run_static_init_blocks` reports together while
it runs the others: a
// name C++ registers, and a second custom validator, after which the block
stops.
mod failing {
use super::*;
- tvm_tirx::register_ops! {
- OpDef::new("tirx.launch_thread", "");
+ tvm_ffi::static_init_block! {
+ OpDef::new("tirx.launch_thread", "").op()?;
}
- tvm_tirx::register_ops! {
- OpDef::new("rust_test.defined_broken",
"").set_validator(accept).set_validator(accept);
- OpDef::new("rust_test.after_broken", "");
+ tvm_ffi::static_init_block! {
+ OpDef::new("rust_test.defined_broken",
"").set_validator(accept).set_validator(accept).op()?;
+ OpDef::new("rust_test.after_broken", "").op()?;
}
}
-/// The blocks of `module` in this test binary.
-fn blocks_of(module: &str) -> Vec<&'static Registration> {
+/// The static init blocks of `module` in this test binary.
+fn blocks_of(module: &str) -> Vec<&'static StaticInitBlock> {
let prefix = format!("ir::op_def{module}:");
let of_module = |identity: &str| {
identity
.strip_prefix(&prefix)
.is_some_and(|rest| rest.starts_with(|c: char| c.is_ascii_digit()))
};
- REGISTRATIONS.iter().filter(|block| of_module(&block.identity())).collect()
+ STATIC_INIT_BLOCKS.iter().filter(|block|
of_module(&block.identity())).collect()
}
-fn is_block_registered(block: &Registration) -> bool {
+fn is_block_registered(block: &StaticInitBlock) -> bool {
Function::get_global(&block.marker()).is_ok()
}
@@ -559,7 +564,7 @@ fn region(
}
#[test]
-fn register_ops_blocks_define_operators() {
+fn static_init_blocks_define_operators_and_types() {
let _serial = serial();
// The blocks ran when the test binary loaded; a block that fails does not
keep the others from
@@ -575,11 +580,11 @@ fn register_ops_blocks_define_operators() {
let fixed = Op::attr_map::<Type>("TFixedReturnType").unwrap().unwrap();
assert!(structural_equal(&fixed.get(&elsewhere).unwrap().unwrap(),
&PrimType::void()));
- // `register_defined_ir` reports the blocks that failed together, each
with its first error;
+ // `run_static_init_blocks` reports the blocks that failed together, each
with its first error;
// a block stops at its first failure, and a failed block does not run
again.
- let error = register_defined_ir().unwrap_err();
+ let error = run_static_init_blocks().unwrap_err();
let message = error.message();
- assert!(message.starts_with("2 register_ops! blocks failed:"),
"{message}");
+ assert!(message.starts_with("2 static init blocks failed:"), "{message}");
assert!(message.contains("op_def::failing:"), "{message}");
assert!(
message.contains("ValueError: Validator of rust_test.defined_broken is
already registered"),
@@ -588,10 +593,10 @@ fn register_ops_blocks_define_operators() {
assert!(message.contains("AttributeError: "), "{message}");
assert!(message.contains("is registered before"), "{message}");
assert!(Op::get("rust_test.after_broken").is_err(), "the block stopped");
- assert_eq!(register_defined_ir().unwrap_err().message(), message);
+ assert_eq!(run_static_init_blocks().unwrap_err().message(), message);
- // The type of a `register_types!` block is registered with them: tvm-ffi
constructs it from its
- // fields, positionally or by keyword, and reads it back from a JSON graph.
+ // The type `register_type` registers with them: tvm-ffi constructs it
from its fields,
+ // positionally or by keyword, and reads it back from a JSON graph.
let init = Function::from_type_attr(BlockTypeObj::type_index(),
"__ffi_init__").unwrap();
let kwargs = global("ffi.GetKwargsObject").call_tuple(()).unwrap();
let (elem, count, four) = (Type::from(PrimType::int(32)),
FfiString::from("count"), 4i64);
@@ -637,7 +642,7 @@ fn register_ops_blocks_define_operators() {
assert_error(
UNREGISTERED.get(),
ATTRIBUTE_ERROR,
- "Operator rust_test.unregistered is not registered\n2 register_ops!
blocks failed:",
+ "Operator rust_test.unregistered is not registered\n2 static init
blocks failed:",
);
assert!(!UNREGISTERED.is(&call.op));
@@ -694,50 +699,3 @@ fn register_ops_blocks_define_operators() {
);
assert!(built.is_err(), "a panicking hook fails the constructor");
}
-
-#[test]
-fn a_block_another_library_registered_is_skipped() {
- let _serial = serial();
- fn register() -> Result<()> {
- OpDef::new("rust_test.registered_once", "").op()?;
- Ok(())
- }
- // One block in two libraries, as one of a crate both link.
- let text = r#"OpDef::new("rust_test.registered_once", "")"#;
- let block = || Registration::__new("register_ops!", "ext::shared", 3,
"0.1.0", text, register);
- let (first, second) = (block(), block());
- assert_eq!(first.marker(), second.marker());
-
- // The first library registers the operator and marks the block; the
second skips it, where
- // registering the name again would fail.
- first.register().unwrap();
- assert!(is_block_registered(&first));
- let op = Op::get("rust_test.registered_once").unwrap();
- second.register().unwrap();
- assert!(Op::get("rust_test.registered_once").unwrap().same_as(&op));
-
- // Another version of the crate is another block, whose name is already
registered.
- let other = Registration::__new("register_ops!", "ext::shared", 3,
"0.2.0", text, register);
- assert!(other.register().unwrap_err().message().contains("is registered
before"));
-}
-
-#[test]
-fn register_ops_block_identity_is_its_crate_version_and_text() {
- fn register() -> Result<()> {
- Ok(())
- }
- let identity = |version, text| {
- Registration::__new("register_ops!", "ext::ops", 7, version, text,
register)
- };
- let text = r#"OpDef::new("ext.clamp",
"Clamp.").set_infer_type(first_arg_type)"#;
- let block = identity("0.1.0", text).identity();
- assert!(block.starts_with("ext::ops:[email protected]#"), "{block}");
-
- // The same block in another library, whose compiler may space the tokens
otherwise.
- let spaced = "OpDef :: new(\"ext.clamp\", \"Clamp.\")\n .
set_infer_type(first_arg_type)";
- assert_eq!(identity("0.1.0", spaced).identity(), block);
- // A block of another version of the crate, or whose text differs, is
another block.
- assert_ne!(identity("0.1.1", text).identity(), block);
- let renamed = text.replace("ext.clamp", "ext.clip");
- assert_ne!(identity("0.1.0", &renamed).identity(), block);
-}
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index a0b61cdfd..a72dd18b1 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -268,7 +268,7 @@ def
test_a_library_registers_its_operators_when_python_loads_it(define_ir_librar
ROUND_TRIP,
libraries=[define_ir_library, copy],
)
- assert "registering operators and types failed" not in stderr, stderr
+ assert "the static init block at" not in stderr, stderr
def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ir_library):
@@ -295,5 +295,5 @@ def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ir_library
""",
libraries=[define_ir_library],
)
- assert "registering operators and types failed" in stderr, stderr
+ assert "the static init block at" in stderr, stderr
assert "Operator tirx.example_mark is registered before" in stderr, stderr
