This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 502a57d93 [FEAT][RUST] Deduplicate register_ops! blocks across
libraries (#42)
502a57d93 is described below
commit 502a57d938877bf7f96625f8b8f8e57696906622
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Oct 26 14:17:47 2026 -0400
[FEAT][RUST] Deduplicate register_ops! blocks across libraries (#42)
## Summary
The operators of `register_ops!` blocks were registered only when
something called
`tirx_register_ops`: `import tvm_tirx` did for `libtvm_tirx_rust` (#36),
but a library of
operators that another crate builds on `tvm-tirx` had to be registered
by hand. And each `cdylib`
links its own copy of every crate it is built from, so two libraries
that depend on the same crate
of operators both contain its blocks, and the second one to run them
failed:
```
AttributeError: Operator tirx.my_clamp is registered before
```
A block now runs when its library loads, as a C++
`TVM_FFI_STATIC_INIT_BLOCK` does, and once per
process: a block that another library has registered, from the same
version of its crate, is
skipped. Python only loads the library: its operators are registered
and, as for every operator of
a `register_ops!` block (#58), get their TVMScript builders, so the
parser and printer treat them
as those of a C++ library, whichever crate defines them:
```rust
// my_ops: a crate of operators.
tvm_tirx::register_ops! {
OpDef::new("tirx.my_clamp", "Clamp an index to the range `[0, bound)`.")
.signature((sig::arg::<IntExpr>("index", ""),
sig::arg::<IntExpr>("bound", "")))
.set_infer_type(first_arg_type);
}
```
```rust
// lib_a and lib_b: two cdylibs that both depend on my_ops.
use my_ops as _; // a crate that is never named is not linked
```
```python
import tvm_ffi
from tvm_tirx.script import tirx as T
tvm_ffi.load_module("target/release/liblib_a.so") # registers
tirx.my_clamp and T.my_clamp
tvm_ffi.load_module("target/release/liblib_b.so") # failed before; now
skips my_ops' block
```
Each block adds a constructor to its library (`.init_array`, or
`__mod_init_func` on macOS). The
first one runs every block of the library and keeps the library loaded,
as the operators keep its
hooks. An error at that point is printed rather than aborting the
process, and `tirx_register_ops`
raises it again, as `import tvm_tirx` does for `libtvm_tirx_rust`. A
block that has registered its
operators is marked by a global function
`ir.registered_op_block.<module:line@version#hash>` in the
TVM FFI registry, which the libraries of a process share; nothing
changes on the C++ side.
The builders come from `tirx.script.expose_script_op`, which `import
tvm_tirx` registers, so a
library loaded before `tvm_tirx` is imported registers its operators
without builders.
---
python/tvm_tirx/backend/rust/__init__.py | 3 +-
rust/examples/define_ops.rs | 7 +-
rust/src/ir/op_def.rs | 196 +++++++++++++++++++++++++++----
rust/src/lib.rs | 31 +++--
rust/tests/ir/op_def.rs | 86 ++++++++++++--
tests/python/ir/test_op_rust.py | 164 +++++++++++++++++++-------
6 files changed, 392 insertions(+), 95 deletions(-)
diff --git a/python/tvm_tirx/backend/rust/__init__.py
b/python/tvm_tirx/backend/rust/__init__.py
index 13aebfd49..74d8e5ea7 100644
--- a/python/tvm_tirx/backend/rust/__init__.py
+++ b/python/tvm_tirx/backend/rust/__init__.py
@@ -37,7 +37,8 @@ def register_backend():
except (OSError, FileNotFoundError, RuntimeError):
pass
else:
- # Unlike a runtime sidecar, the library registers its operators only
when asked to.
+ # Loading the library registers its operators, and keeps it, and the
hooks they have,
+ # loaded for the process; ``tirx_register_ops`` raises the error of a
block that failed.
tvm_ffi.load_module(_LOADED_LIBS["tvm_tirx_rust"]._name).tirx_register_ops()
return None
diff --git a/rust/examples/define_ops.rs b/rust/examples/define_ops.rs
index 931a8069b..4bdb49e32 100644
--- a/rust/examples/define_ops.rs
+++ b/rust/examples/define_ops.rs
@@ -21,10 +21,9 @@
//!
//! `cargo build --example define_ops` builds this as
`target/debug/examples/libdefine_ops.so`,
//! which exports `tirx_register_ops`, as every library built on `tvm_tirx`
does.
-//! `tests/python/ir/test_op_rust.py` loads it and calls `tirx_register_ops`,
as `import tvm_tirx`
-//! does for `libtvm_tirx_rust`, which registers these operators and gives
each its TVMScript
-//! builder: `T.example_clamp(i, n)`, `T.example_mark()`, `my_ext.hint(i,
attrs={..})` and
-//! `with T.example_loop(n) as j:`.
+//! `tests/python/ir/test_op_rust.py` loads it, which registers these
operators and gives each its
+//! TVMScript builder, as `import tvm_tirx` loads `libtvm_tirx_rust`:
`T.example_clamp(i, n)`,
+//! `T.example_mark()`, `my_ext.hint(i, attrs={..})` and `with
T.example_loop(n) as j:`.
use tvm_tirx::ir::op_attr::{CALL_EFFECT_KIND, FIXED_RETURN_TYPE};
use tvm_tirx::ir::sig::{self, IntExpr};
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 1b96b13f0..db244c1a3 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -28,11 +28,10 @@
//! `OpDef::signature` does from `sig::arg<T>`; each signature is its own
compiled function.
//!
//! The [`register_ops!`](crate::register_ops) blocks are gathered at link
time, wherever in the
-//! program or library they are, and [`register_defined_ops`] runs them all.
Unlike a C++ static
-//! init block, a block does not run when its library loads. Every library
built on this crate
-//! exports [`register_defined_ops`] as the TVM FFI function
`tirx_register_ops`, which
-//! `import tvm_tirx` calls on the library built from this crate
(`libtvm_tirx_rust`) right after
-//! it loads it, so Python sees its operators as it sees the C++ ones.
+//! 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 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, registers its operators once.
use std::any::TypeId;
use std::cell::RefCell;
@@ -40,11 +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::OnceLock;
+use std::sync::{Once, OnceLock};
use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFIObject, TVMFFITypeIndex};
use tvm_ffi::{
- Any, AnyView, Array, Error, ObjectArc, ObjectRefCore, Result, String,
INTERNAL_ERROR,
+ Any, AnyView, Array, Error, Function, ObjectArc, ObjectRefCore, Result,
String, INTERNAL_ERROR,
TYPE_ERROR, VALUE_ERROR,
};
@@ -322,9 +321,9 @@ 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 first lookup of an operator that is
not registered runs
-/// the [`register_ops!`](crate::register_ops) blocks first, so that a Rust
program can use the
-/// operators they define without registering them itself.
+/// [`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.
///
/// ```no_run
/// use tvm_tirx::ir::sig::{self, IntExpr};
@@ -442,6 +441,10 @@ pub struct OpRegistration {
/// The module and line of the block.
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>>,
@@ -455,16 +458,40 @@ pub static OP_REGISTRATIONS: [OpRegistration];
impl OpRegistration {
#[doc(hidden)]
- pub const fn __new(module: &'static str, line: u32, register: fn() ->
Result<()>) -> Self {
- Self { module, line, register, outcome: OnceLock::new() }
+ pub const fn __new(
+ module: &'static str,
+ line: u32,
+ version: &'static str,
+ text: &str,
+ register: fn() -> Result<()>,
+ ) -> Self {
+ let text_hash = text_hash(text);
+ Self { 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_op_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_op_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.
- fn register(&self) -> Result<()> {
+ #[doc(hidden)]
+ pub fn register(&self) -> Result<()> {
let mut failure = None;
let outcome = self.outcome.get_or_init(|| {
- (self.register)().map_err(|error| {
+ self.run().map_err(|error| {
let text = format!("{}: {}", error.kind(), error.message());
failure = Some(error);
text
@@ -480,20 +507,51 @@ impl OpRegistration {
}
}
+ /// 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 its
+ /// operators are 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) 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.
+/// 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 once.
///
-/// `import tvm_tirx` calls it on `libtvm_tirx_rust` through the TVM FFI
function
-/// `tirx_register_ops`, which every library built on this crate exports. A
Rust program calls it
-/// before it builds IR, as the registry does not synchronize registration
with readers; otherwise
-/// the first lookup through an [`OpHandle`] of an operator that is not
registered calls it.
+/// 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_ops`, 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_ops() -> Result<()> {
let mut blocks: Vec<&OpRegistration> = OP_REGISTRATIONS.iter().collect();
blocks.sort_by_key(|block| (block.module, block.line));
@@ -509,6 +567,69 @@ pub fn register_defined_ops() -> Result<()> {
Err(Error::new(failures[0].1.kind(), &message, ""))
}
+/// The constructor of every [`register_ops!`](crate::register_ops) 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_ops`] 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_ops) {
+ Ok(Ok(())) => {}
+ Ok(Err(error)) => eprintln!(
+ "{library}: registering operators failed: {}: {}",
+ error.kind(),
+ error.message()
+ ),
+ Err(_) => eprintln!("{library}: registering operators panicked"),
+ }
+ });
+}
+
+/// Keep the library that holds this code loaded for the rest of the process,
as the operators 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 {
@@ -551,15 +672,26 @@ fn doc_text(doc: &str) -> std::string::String {
/// .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque);
/// }
///
-/// tvm_tirx::ir::register_defined_ops()?;
+/// // 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>(())
/// ```
///
-/// Unlike a C++ static init block, a block does not run when its library
loads, but when the
-/// operators are registered: [`register_defined_ops`] runs every block of the
program or library,
-/// which `import tvm_tirx` does through `tirx_register_ops`, and an
[`OpHandle`] runs them on the
-/// first lookup of an operator that is not registered. A block runs once.
+/// 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_ops`]. 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_ops`] (`tirx_register_ops`) 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);* $(;)?) => {
@@ -570,8 +702,20 @@ macro_rules! register_ops {
}
#[$crate::__private::linkme::distributed_slice($crate::ir::OP_REGISTRATIONS)]
#[linkme(crate = $crate::__private::linkme)]
- static REGISTRATION: $crate::ir::OpRegistration =
- $crate::ir::OpRegistration::__new(module_path!(), line!(),
register);
+ static REGISTRATION: $crate::ir::OpRegistration =
$crate::ir::OpRegistration::__new(
+ module_path!(),
+ line!(),
+ match option_env!("CARGO_PKG_VERSION") {
+ Some(version) => version,
+ None => "",
+ },
+ stringify!($($def);*),
+ register,
+ );
+ #[used]
+ #[cfg_attr(target_os = "linux", link_section = ".init_array")]
+ #[cfg_attr(target_os = "macos", link_section =
"__DATA,__mod_init_func")]
+ static REGISTER_AT_LOAD: extern "C" fn() =
$crate::ir::__register_at_load;
};
};
}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index ed1f56047..261f235f8 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -129,23 +129,30 @@
//! ```
//!
//! An [`ir::OpHandle`] in a `static` keeps an operator as C++ keeps one in
`const Op& xxx_op()`.
-//! A block does not run when its library loads, as a C++ one does, but when
the operators are
-//! registered: [`ir::register_defined_ops`] runs every block, wherever it is,
and every library
-//! built on the crate exports it as the TVM FFI function `tirx_register_ops`;
a handle runs them on
-//! the first lookup of an operator that is not registered.
+//! 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_ops`],
which reports the
+//! errors of the blocks that failed, as the TVM FFI function
`tirx_register_ops`.
//!
//! # The library Python loads
//!
//! The crate is also built as a shared library, which the TIRx build installs
as
//! `libtvm_tirx_rust` next to `libtvm_tirx_compiler` when it finds a cargo as
new as the crate's
//! `rust-version` (`TVM_TIRX_USE_RUST`, `AUTO` by default). `import tvm_tirx`
loads it through
-//! the backend autoload (`tvm_tirx.backend.rust`) and calls
`tirx_register_ops`, so the operators
-//! the crate defines are registered by the time the import returns, as the
C++ ones are: `Op.get`
-//! finds them, a `Call` built in Python runs their hooks, and the TVMScript
printer prints their
-//! calls. The TVMScript builder of an operator (`T.<name>` for `tirx.<name>`,
`my_ext.<name>` for
-//! `my_ext.<name>`) is generated by `tvm_tirx.tirx.script.expose_script_op`,
unless Python has
-//! one written for it, as for a C++ operator. `examples/define_ops.rs` builds
a library of its own
-//! the same way, which `tvm_ffi.load_module(path).tirx_register_ops()`
registers.
+//! the backend autoload (`tvm_tirx.backend.rust`), which registers the
operators the crate defines
+//! by the time the import returns, as the C++ ones are, and calls
`tirx_register_ops` 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.
+//! The TVMScript builder of an operator (`T.<name>` for `tirx.<name>`,
`my_ext.<name>` for
+//! `my_ext.<name>`) is generated by `tvm_tirx.tirx.script.expose_script_op`,
unless Python has one
+//! written for it, as for a C++ operator.
+//!
+//! A library of operators of its own, a `cdylib` built on the crate as
`examples/define_ops.rs`
+//! is, registers its operators when Python loads it, by `tvm_ffi.load_module`
or `ctypes`, as a
+//! C++ library does, and stays loaded for the process. Each library links its
own copy of the
+//! crates it is built from, so several libraries can have the same block, as
those of a crate they
+//! share: the first library that runs it registers its operators and the
others skip it, so the
+//! libraries load together. A crate that a library depends on but never names
is not linked into
+//! it, nor are its blocks: name it, as `use my_ops as _;` does.
//!
//! # Defining a type in Rust
//!
@@ -290,7 +297,7 @@ pub mod tirx;
pub use tvm_ffi;
-// The registration entry of every library built on the crate, which `import
tvm_tirx` calls on
+// 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_ops`, the TVM FFI function
`tirx_register_ops` of
// the library as `tvm_ffi.load_module` loads it.
tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ops,
ir::register_defined_ops);
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 56857164a..928cfa5b5 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -32,14 +32,14 @@ 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_ops, AttrKey, Attrs, Call, CallEffectKind, CallObj,
DictAttrs, Evaluate, Expr,
- FuncType, IntImm, Op, OpDef, OpHandle, PrimType, RegionStmt,
RegionStmtObj, StringImm, Type,
- Var, VarObj, UNKNOWN_LOC,
+ FuncType, IntImm, Op, OpDef, OpHandle, OpRegistration, PrimType,
RegionStmt, RegionStmtObj,
+ StringImm, Type, Var, VarObj, OP_REGISTRATIONS, UNKNOWN_LOC,
};
use tvm_tirx::tirx::{CallFFIKernelAttr, TensorVar};
use tvm_tirx::tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
use tvm_tirx::tvm_ffi::{
- Any, Array, Error, ErrorKind, Map, ObjectRefCast, ObjectRefCore, Result,
String as FfiString,
- ATTRIBUTE_ERROR, INTERNAL_ERROR, TYPE_ERROR, VALUE_ERROR,
+ Any, Array, Error, ErrorKind, Function, Map, ObjectRefCast, ObjectRefCore,
Result,
+ String as FfiString, ATTRIBUTE_ERROR, INTERNAL_ERROR, TYPE_ERROR,
VALUE_ERROR,
};
use super::{int32, parse_dtype, tensor_var};
@@ -457,6 +457,21 @@ mod failing {
}
}
+/// The blocks of `module` in this test binary.
+fn blocks_of(module: &str) -> Vec<&'static OpRegistration> {
+ 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()))
+ };
+ OP_REGISTRATIONS.iter().filter(|block|
of_module(&block.identity())).collect()
+}
+
+fn is_block_registered(block: &OpRegistration) -> bool {
+ Function::get_global(&block.marker()).is_ok()
+}
+
static IDENTITY: OpHandle = OpHandle::new("rust_test.defined_identity");
static LAUNCH: OpHandle = OpHandle::new("rust_test.launch");
static PANICKING: OpHandle = OpHandle::new("rust_test.panicking_region");
@@ -528,11 +543,10 @@ fn region(
fn register_ops_blocks_define_operators() {
let _serial = serial();
- // The first lookup through a handle of an operator not registered runs
every block; a block
- // that fails does not keep the others from registering theirs.
- assert!(Op::get("rust_test.defined_identity").is_err(), "not registered
before");
- let identity = IDENTITY.get().unwrap();
- assert!(identity.same_as(&Op::get("rust_test.defined_identity").unwrap()));
+ // The blocks ran when the test binary loaded; a block that fails does not
keep the others from
+ // registering theirs.
+ let identity = Op::get("rust_test.defined_identity").unwrap();
+ assert!(identity.same_as(&IDENTITY.get().unwrap()));
assert_eq!(identity.doc.as_str(), "Return the value.\n\nThe documentation
is dedented.");
assert_eq!(identity.args_info.get(0).unwrap().name.as_str(), "value");
let effect = Op::attr_map::<i64>("TCallEffectKind").unwrap().unwrap();
@@ -540,7 +554,7 @@ fn register_ops_blocks_define_operators() {
let elsewhere = Op::get("rust_test.defined_elsewhere").unwrap();
assert_eq!(elsewhere.doc.as_str(), "An operator in another module.");
let fixed = Op::attr_map::<Type>("TFixedReturnType").unwrap().unwrap();
-
assert!(fixed.get(&elsewhere).unwrap().unwrap().same_as(&PrimType::void()));
+ assert!(structural_equal(&fixed.get(&elsewhere).unwrap().unwrap(),
&PrimType::void()));
// `register_defined_ops` 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.
@@ -557,6 +571,13 @@ fn register_ops_blocks_define_operators() {
assert!(Op::get("rust_test.after_broken").is_err(), "the block stopped");
assert_eq!(register_defined_ops().unwrap_err().message(), message);
+ // A block that succeeds is marked registered for the other libraries of
the process; one that
+ // fails is not, so it fails in each of them.
+
assert!(blocks_of("").into_iter().chain(blocks_of("::elsewhere")).all(is_block_registered));
+ let failing = blocks_of("::failing");
+ assert_eq!(failing.len(), 2);
+ assert!(!failing.into_iter().any(is_block_registered));
+
// The handle keeps the operator: the Rust `FInferType` types a Call of
it, and another thread
// sees the same operator.
assert!(IDENTITY.get().unwrap().same_as(&identity));
@@ -631,3 +652,48 @@ 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 = || OpRegistration::__new("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 = OpRegistration::__new("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| OpRegistration::__new("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 b5c43f6d1..57928ae5c 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -14,11 +14,13 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Operators defined in Rust are registered when their library is loaded, and
used as C++ ones are.
+"""Operators defined in Rust are registered when their library loads, as C++
ones are.
-``rust/examples/define_ops.rs`` is built by cargo and loaded by a fresh
interpreter, as
-``import tvm_tirx`` loads libtvm_tirx_rust, which then uses its operators as
it uses C++ ones:
-registering them also gives them their TVMScript builders.
+``rust/examples/define_ops.rs`` is built by cargo into a library of its own,
which a fresh
+interpreter loads, as ``import tvm_tirx`` loads libtvm_tirx_rust, and then
uses its operators as it
+uses C++ ones: registering them also gives them their TVMScript builders. A
copy of the library is
+a second library with the same ``register_ops!`` blocks, as libraries that
link the same crate
+have.
"""
import os
@@ -72,22 +74,61 @@ def define_ops_library():
return target_dir / "debug" / "examples" / name
-def run_python(code, library):
- """Run `code` in a fresh interpreter, with the path of `library` as its
argument.
+def run_python(*code, libraries):
+ """Run `code` in a fresh interpreter, with the paths of `libraries` as its
arguments, and
+ return its standard error.
The code runs from a file: TVMScript reads the source of the functions it
parses, which Python
keeps for code passed with `-c` only from 3.13 on.
"""
with tempfile.TemporaryDirectory() as tmp:
script = Path(tmp) / "script.py"
- script.write_text(textwrap.dedent(code))
+ script.write_text("\n".join(textwrap.dedent(part) for part in code))
result = subprocess.run(
- [sys.executable, str(script), str(library)],
+ [sys.executable, str(script), *map(str, libraries)],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
+ return result.stderr
+
+
+# With the TVMScript builders that registering the example operators gives
them, TVMScript writes,
+# prints and parses their calls as it does those of C++ operators: a call with
a value, a call
+# emitted as a statement, `Call.attrs`, and a region, whose body parameter the
Rust hook creates.
+ROUND_TRIP = """
+ import tvm_ffi
+ import tvm_tirx
+ from tvm_tirx import ir
+ from tvm_tirx.script import tirx as T
+
+ # Each builder is generated from the signature of its operator under the
name the printer
+ # gives its calls: `T.<name>` for `tirx.<name>`, and outside `tirx.`, in
the parser namespace
+ # its root names.
+ assert T.example_clamp.__doc__ == ir.Op.get("tirx.example_clamp").doc
+ assert T.example_loop.__doc__ == ir.Op.get("tirx.example_loop").doc
+ from tvm_tirx.script.parser import my_ext
+
+ assert my_ext.hint.__doc__ == ir.Op.get("my_ext.hint").doc
+
+ @T.function
+ def func(i: T.int32, n: T.int32):
+ T.evaluate(T.example_clamp(i + 1, n))
+ T.example_mark()
+ T.evaluate(my_ext.hint(i, attrs={"note": "hot"}))
+ with T.example_loop(n) as j:
+ T.evaluate(T.example_clamp(j, n))
+
+ text = func.script()
+ assert "T.example_clamp(i + 1, n)" in text, text
+ assert "T.example_mark()" in text, text
+ assert 'my_ext.hint(i, attrs={"note": "hot"})' in text, text
+ assert "with T.example_loop(n) as j:" in text, text
+ parsed = tvm_tirx.script.parse(text, extra_vars={"T": T})
+ assert tvm_ffi.structural_equal(parsed, func), parsed.script()
+ assert parsed.script() == text
+"""
def test_library_is_loaded_on_import():
@@ -110,19 +151,16 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
import sys
import tvm_ffi
- import tvm_tirx
from tvm_tirx import ir
- from tvm_tirx.script import tirx as T
- # Registered by the library, with the doc comment and signature the
Rust definition gives.
+ # Registered when the library loads, as libtvm_tirx_rust is when
`import tvm_tirx` loads
+ # it, with the doc comment and signature the Rust definition gives.
tvm_ffi.load_module(sys.argv[1]).tirx_register_ops()
clamp = ir.Op.get("tirx.example_clamp")
assert clamp.doc == "Clamp an index to the range `[0, bound)`."
assert [arg.name for arg in clamp.args_info] == ["index", "bound"]
mark = ir.Op.get("tirx.example_mark")
assert mark.doc.startswith("Mark a point in the program.\\n\\nThe call
has no value")
- hint = ir.Op.get("my_ext.hint")
- loop = ir.Op.get("tirx.example_loop")
# A Call built in Python runs the Rust type inference and validator.
i, n = ir.Var("i", "int32"), ir.Var("n", "int32")
@@ -135,35 +173,77 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
assert "`Call.args[1]` (`bound`) expected `ir.IntExpr`" in
str(error), error
else:
raise AssertionError("a float bound passed validation")
+ """,
+ ROUND_TRIP,
+ libraries=[define_ops_library],
+ )
+
+
+def
test_a_library_registers_its_operators_when_python_loads_it(define_ops_library,
tmp_path):
+ copy = tmp_path / define_ops_library.name
+ shutil.copy(define_ops_library, copy)
+ stderr = run_python(
+ """
+ import ctypes
+ import gc
+ import sys
+
+ import tvm_ffi
+ from tvm_tirx import ir
+
+ library, copy = sys.argv[1:]
+
+ # The library of another crate: importing tvm_tirx does not register
its operators.
+ try:
+ ir.Op.get("tirx.example_clamp")
+ except Exception:
+ pass
+ else:
+ raise AssertionError("registered before the library loaded")
+
+ # Loading it registers them, as loading a C++ library does. The
library stays loaded when
+ # Python drops its module, as the operators keep its hooks.
+ module = tvm_ffi.load_module(library)
+ del module
+ gc.collect()
+ clamp = ir.Op.get("tirx.example_clamp")
+ i, n = ir.Var("i", "int32"), ir.Var("n", "int32")
+ assert tvm_ffi.structural_equal(ir.Call(clamp, [i, n]).ty, i.ty)
+
+ # The copy has the blocks of the library, as two libraries that link
the same crate do:
+ # loading it registers nothing again.
+ ctypes.CDLL(copy)
+ assert ir.Op.get("tirx.example_clamp").same_as(clamp)
+ """,
+ ROUND_TRIP,
+ libraries=[define_ops_library, copy],
+ )
+ assert "registering operators failed" not in stderr, stderr
+
- # Registration gave each its TVMScript builder, generated from its
signature under the name
- # the printer gives its calls.
- assert T.example_clamp.__doc__ == clamp.doc
- assert T.example_loop.__doc__ == loop.doc
- # Outside `tirx.`, in the parser namespace its root names.
- from tvm_tirx.script.parser import my_ext
-
- assert my_ext.hint.__doc__ == hint.doc
-
- # With them, TVMScript writes, prints and parses the calls as it does
those of C++
- # operators: a call with a value, a call emitted as a statement,
`Call.attrs`, and a region,
- # whose body parameter the Rust hook creates.
- @T.function
- def func(i: T.int32, n: T.int32):
- T.evaluate(T.example_clamp(i + 1, n))
- T.example_mark()
- T.evaluate(my_ext.hint(i, attrs={"note": "hot"}))
- with T.example_loop(n) as j:
- T.evaluate(T.example_clamp(j, n))
-
- text = func.script()
- assert "T.example_clamp(i + 1, n)" in text, text
- assert "T.example_mark()" in text, text
- assert 'my_ext.hint(i, attrs={"note": "hot"})' in text, text
- assert "with T.example_loop(n) as j:" in text, text
- parsed = tvm_tirx.script.parse(text, extra_vars={"T": T})
- assert tvm_ffi.structural_equal(parsed, func), parsed.script()
- assert parsed.script() == text
+def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ops_library):
+ stderr = run_python(
+ """
+ import sys
+
+ import tvm_ffi
+ from tvm_tirx import ir
+
+ # Another library registered the name of an example operator: the
block that defines it
+ # fails when the library loads, and the other block registers its
operators.
+ tvm_ffi.get_global_func("ir.RegisterOp")("tirx.example_mark", "")
+ module = tvm_ffi.load_module(sys.argv[1])
+ ir.Op.get("tirx.example_clamp")
+
+ # `tirx_register_ops` raises the error again.
+ try:
+ module.tirx_register_ops()
+ except Exception as error:
+ assert "tirx.example_mark is registered before" in str(error),
error
+ else:
+ raise AssertionError("the block that failed was not reported")
""",
- define_ops_library,
+ libraries=[define_ops_library],
)
+ assert "registering operators failed" in stderr, stderr
+ assert "Operator tirx.example_mark is registered before" in stderr, stderr
