This is an automated email from the ASF dual-hosted git repository.
Seven-Streams 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 37b583953 [REFACTOR][RUST] Look up operators with Op::get and remove
OpHandle (#160)
37b583953 is described below
commit 37b583953d693c232139d299d6b92ccc96c98475
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Nov 2 10:53:06 2026 -0500
[REFACTOR][RUST] Look up operators with Op::get and remove OpHandle (#160)
## Summary
Operators defined in Rust are registered in
`tvm_ffi::static_init_block!`s, which run when their
library loads, so `Op::get` finds them wherever C++ `Op::Get` would find
theirs. A pass looks its
operators up with `Op::get` and tests a call with
`call.op.same_as(&op)`, as in C++:
```rust
let cmul = Op::get("my_ext.cmul")?;
if call.op.same_as(&cmul) { ... }
```
`OpHandle`, the `static` handle that looked an operator up on first use,
is removed. The extension
example and its tutorial use `Op::get`.
---
docs/advanced/defining_ir_extensions.rst | 40 +++++-----
examples/extension/src/lib.rs | 69 +++++++++++------
rust/src/ir/op_def.rs | 127 +------------------------------
rust/src/lib.rs | 12 +--
rust/tests/ir/op_def.rs | 42 ++++------
5 files changed, 90 insertions(+), 200 deletions(-)
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index 1ba916634..50458ea45 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -153,25 +153,26 @@ that pair. This lets the rules apply to the step-by-step
function from
R[i] = product_re
Im[i] = product_im
-Only scalar arithmetic remains. The pass refers to the four operators through
-handles:
+Only scalar arithmetic remains. The pass looks up the four operators by name,
+as C++ does with ``Op::Get``:
.. literalinclude:: ../../examples/extension/src/lib.rs
:language: rust
- :start-after: // begin handles
- :end-before: // end handles
+ :start-after: // begin operators
+ :end-before: // end operators
-A handle names a registered operator and looks it up on first use.
-``CMUL.is(&call.op)`` tests whether a call's ``op`` is ``my_ext.cmul``, and
-``COMPLEX.get()?`` returns ``my_ext.complex`` to build a new call. ``Lowering``
-implements these rewrites with them:
+The operators are registered when the library loads, and ``Op::get`` returns
+the registry's single object for a name. ``call.op.same_as(&ops.cmul)``
+therefore tests whether a call's ``op`` is ``my_ext.cmul`` by comparing
+pointers. A lookup goes through the registry, so the pass fetches the operators
+once per run and keeps them. ``Lowering`` implements the rewrites with them:
.. literalinclude:: ../../examples/extension/src/lib.rs
:language: rust
:start-after: // begin lowering
:end-before: // end lowering
- :prepend: #[derive(Default)]
- struct Lowering {
+ :prepend: struct Lowering {
+ ops: ComplexOps,
bindings: Vec<(Var, Var, Var)>,
}
impl Lowering {
@@ -182,13 +183,13 @@ implements these rewrites with them:
fn map_sequence(&mut self, seq: SeqStmt) -> SeqStmt { ... }
:append: }
-``bindings`` records each complex variable with its two scalar replacements,
-and ``parts`` reads the real and imaginary parts of an operand, either a
-``complex(re, im)`` call or a variable already split. ``#[dispatch(map)]``
-sends each node to the ``map_`` method of its type, and leaves other nodes
-unchanged: ``map_call`` applies the three rules, ``map_bind`` splits a complex
-binding, and ``map_sequence`` flattens the sequence that a split binding leaves
-inside another.
+``ops`` holds the four operators, ``bindings`` records each complex variable
+with its two scalar replacements, and ``parts`` reads the real and imaginary
+parts of an operand, either a ``complex(re, im)`` call or a variable already
+split. ``#[dispatch(map)]`` sends each node to the ``map_`` method of its type,
+and leaves other nodes unchanged: ``map_call`` applies the three rules,
+``map_bind`` splits a complex binding, and ``map_sequence`` flattens the
+sequence that a split binding leaves inside another.
The pass is built from an ordinary named function:
@@ -202,7 +203,8 @@ rewritten function; ``Some`` replaces the original. The
first ``structural_map``
walks the function with ``Lowering``. ``PostOrder`` visits children before
their
parent, so the operands of ``cmul`` are lowered before ``cmul`` itself. The
second ``structural_map`` only checks the result. It takes a tuple of three
-functions, and calls each on the nodes of its parameter type:
+functions, and calls each on the nodes of its parameter type; the closure
+passes the operators to ``check_call_lowered``:
.. code-block:: rust
@@ -213,7 +215,7 @@ functions, and calls each on the nodes of its parameter
type:
fn check_expr(expr: Expr) -> Result<Expr> { ... }
/// Reject a call to one of the four complex operators that the rewrite did
not consume.
- fn check_call_lowered(call: Call) -> Result<Call> { ... }
+ fn check_call_lowered(ops: &ComplexOps, call: Call) -> Result<Call> { ... }
Each returns its node unchanged and raises a ``ValueError`` on an unsupported
use, such as a complex function parameter, so the pass reports such a use
diff --git a/examples/extension/src/lib.rs b/examples/extension/src/lib.rs
index 0f6fdeda1..5fa78ebb7 100644
--- a/examples/extension/src/lib.rs
+++ b/examples/extension/src/lib.rs
@@ -24,8 +24,8 @@ use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
use tvm_tirx::ir::sig::{self, PrimExpr};
use tvm_tirx::ir::transform::{Pass, PassContext, PassInfo};
use tvm_tirx::ir::{
- dtype, register_type, Bind, Call, CallEffectKind, CallObj, Expr, IRModule,
Op, OpDef, OpHandle,
- PrimType, SeqStmt, Stmt, Type, TypeObj, Var,
+ dtype, register_type, Bind, Call, CallEffectKind, CallObj, Expr, IRModule,
Op, OpDef, PrimType,
+ SeqStmt, Stmt, Type, TypeObj, Var,
};
use tvm_tirx::tirx::transform::FunctionPass;
use tvm_tirx::tirx::Function;
@@ -115,12 +115,11 @@ fn infer_part(call: &CallObj) -> Result<Type> {
// Explicit result types bypass inference, so validation checks the inputs and
result again.
fn validate_call(call: &CallObj) -> Result<()> {
- let expected = if COMPLEX.is(&call.op) {
- infer_complex(call)?
- } else if CMUL.is(&call.op) {
- infer_cmul(call)?
- } else {
- infer_part(call)?
+ let op: Op = call.op.clone().try_cast()?;
+ let expected = match op.name.as_str() {
+ "my_ext.complex" => infer_complex(call)?,
+ "my_ext.cmul" => infer_cmul(call)?,
+ _ => infer_part(call)?,
};
let matches: bool = tvm_ffi::cached_global_func!("ffi.StructuralEqual")
.call_tuple((&expected, &call.ty, false, false))?
@@ -162,15 +161,33 @@ tvm_ffi::static_init_block! {
.set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
}
-// begin handles
-static COMPLEX: OpHandle = OpHandle::new("my_ext.complex");
-static CMUL: OpHandle = OpHandle::new("my_ext.cmul");
-static REAL: OpHandle = OpHandle::new("my_ext.real");
-static IMAG: OpHandle = OpHandle::new("my_ext.imag");
-// end handles
+// begin operators
+struct ComplexOps {
+ complex: Op,
+ cmul: Op,
+ real: Op,
+ imag: Op,
+}
+
+impl ComplexOps {
+ fn get() -> Result<Self> {
+ Ok(Self {
+ complex: Op::get("my_ext.complex")?,
+ cmul: Op::get("my_ext.cmul")?,
+ real: Op::get("my_ext.real")?,
+ imag: Op::get("my_ext.imag")?,
+ })
+ }
+
+ /// Whether `op`, such as a call's `op`, is one of the four operators.
+ fn contains(&self, op: &Expr) -> bool {
+ [&self.complex, &self.cmul, &self.real,
&self.imag].into_iter().any(|own| op.same_as(own))
+ }
+}
+// end operators
-#[derive(Default)]
struct Lowering {
+ ops: ComplexOps,
// Each complex local variable is replaced by two scalar local variables.
bindings: Vec<(Var, Var, Var)>,
}
@@ -178,7 +195,7 @@ struct Lowering {
impl Lowering {
fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> {
if let Some(call) = expr.as_node::<CallObj>() {
- if COMPLEX.is(&call.op) {
+ if call.op.same_as(&self.ops.complex) {
return Ok((call.args.get(0)?, call.args.get(1)?));
}
}
@@ -199,10 +216,11 @@ impl Lowering {
// begin lowering
fn map_call(&mut self, call: Call) -> Result<Expr> {
- if COMPLEX.is(&call.op) || CMUL.is(&call.op) || REAL.is(&call.op) ||
IMAG.is(&call.op) {
+ let ops = &self.ops;
+ if ops.contains(&call.op) {
validate_call(&call)?;
}
- if CMUL.is(&call.op) {
+ if call.op.same_as(&ops.cmul) {
let (a, b) = self.parts(&call.args.get(0)?)?;
let (c, d) = self.parts(&call.args.get(1)?)?;
// (a + bi)(c + di) = (ac - bd) + (ad + bc)i
@@ -210,7 +228,7 @@ impl Lowering {
let im = &a * &d + &b * &c;
return Ok(Call::new(
None,
- COMPLEX.get()?,
+ ops.complex.clone(),
[re, im],
None,
Vec::<Type>::new(),
@@ -218,9 +236,9 @@ impl Lowering {
)?
.into());
}
- if REAL.is(&call.op) || IMAG.is(&call.op) {
+ if call.op.same_as(&ops.real) || call.op.same_as(&ops.imag) {
let (re, im) = self.parts(&call.args.get(0)?)?;
- return Ok(if REAL.is(&call.op) { re } else { im });
+ return Ok(if call.op.same_as(&ops.real) { re } else { im });
}
Ok(call.into())
}
@@ -262,8 +280,8 @@ fn check_expr(expr: Expr) -> Result<Expr> {
}
/// Reject a call to one of the four complex operators that the rewrite did
not consume.
-fn check_call_lowered(call: Call) -> Result<Call> {
- if COMPLEX.is(&call.op) || CMUL.is(&call.op) || REAL.is(&call.op) ||
IMAG.is(&call.op) {
+fn check_call_lowered(ops: &ComplexOps, call: Call) -> Result<Call> {
+ if ops.contains(&call.op) {
return Err(Error::new(
tvm_ffi::VALUE_ERROR,
"LowerComplex left an unsupported complex call",
@@ -279,11 +297,12 @@ fn lower_function(
_module: IRModule,
_ctx: PassContext,
) -> Result<Option<Function>> {
- let mut lowering = Lowering::default();
+ let mut lowering = Lowering { ops: ComplexOps::get()?, bindings:
Vec::new() };
let rewritten = tvm_ffi::structural_map(func, &mut lowering,
WalkOrder::PostOrder)?;
+ let ops = &lowering.ops;
let checked = tvm_ffi::structural_map(
rewritten,
- (check_type, check_call_lowered, check_expr),
+ (check_type, |call| check_call_lowered(ops, call), check_expr),
WalkOrder::PostOrder,
)?;
let function: Function = checked.try_into()?;
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 07d846578..c91b09345 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -19,8 +19,7 @@
//! Registering operators from Rust: C++ `OpDef`, with its signature
descriptors ([`sig`]) and
//! 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.
+//! operators, as C++ defines them in a `TVM_FFI_STATIC_INIT_BLOCK`.
//!
//! Every hook an operator gets here is native, as the C++ ones are: a Rust
function compiled into
//! the C ABI adapter that C++ `NativeFunctionView` calls through a function
pointer, without
@@ -34,17 +33,12 @@
//! links its own copy of a crate they share, runs once.
use std::any::TypeId;
-use std::cell::RefCell;
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 tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFIObject, TVMFFITypeIndex};
-use tvm_ffi::{
- Any, AnyView, Array, Error, ObjectArc, ObjectRefCore, Result, String,
TYPE_ERROR, VALUE_ERROR,
-};
+use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex};
+use tvm_ffi::{Any, AnyView, Array, Error, ObjectRefCore, Result, String,
TYPE_ERROR, VALUE_ERROR};
use super::attrs::Attrs;
use super::expr::{CallObj, Expr, Var};
@@ -345,121 +339,6 @@ 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 [`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::{self, Result};
-///
-/// fn first_arg_type(call: &CallObj) -> Result<Type> {
-/// Ok(call.args.get(0)?.ty.clone())
-/// }
-///
-/// 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);
-/// }
-///
-/// static CLAMP: OpHandle = OpHandle::new("my_ext.clamp");
-/// static LAUNCH_THREAD: OpHandle = OpHandle::new("tirx.launch_thread");
-///
-/// let index = Var::new("i", PrimType::int(32), UNKNOWN_LOC);
-/// let bound = Var::new("n", PrimType::int(32), UNKNOWN_LOC);
-/// let args = [Expr::from(&index), Expr::from(&bound)];
-/// let clamped = Call::new(None, CLAMP.get()?, args, None,
Vec::<Type>::new(), UNKNOWN_LOC)?;
-/// assert!(CLAMP.is(&clamped.op) && !LAUNCH_THREAD.is(&clamped.op));
-/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
-/// ```
-pub struct OpHandle {
- name: &'static str,
- /// This handle's slot in each thread's cache of fetched operators.
- slot: OnceLock<usize>,
- /// The operator once fetched, for identity tests: operators are registry
singletons that
- /// live for the process, and `is` is the hot path of every pass that
looks at `Call::op`.
- object: AtomicPtr<TVMFFIObject>,
-}
-
-/// The number of slots handed out to handles.
-static SLOTS: AtomicUsize = AtomicUsize::new(0);
-
-thread_local! {
- /// The operators this thread fetched, by handle slot.
- static FETCHED: RefCell<Vec<Option<Op>>> = const {
RefCell::new(Vec::new()) };
-}
-
-impl OpHandle {
- /// 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()) }
- }
-
- pub fn name(&self) -> &'static str {
- self.name
- }
-
- /// The operator. One that is not registered runs the blocks first; if it
is still missing, the
- /// error says so, followed by the failures of the blocks, if any.
- pub fn get(&self) -> Result<Op> {
- self.with(Op::clone)
- }
-
- /// Whether `value`, such as a Call's `op`, is this operator (C++
`value.same_as(op)`). No
- /// value is an operator that cannot be fetched, such as one this build
does not register.
- pub fn is(&self, value: &impl ObjectRefCore) -> bool {
- let cached = self.object.load(Ordering::Acquire);
- if !cached.is_null() {
- // SAFETY: the pointer is only compared, never dereferenced.
- let pointer = unsafe {
ObjectArc::as_raw(ObjectRefCore::data(value)) };
- return pointer.cast::<TVMFFIObject>().cast_mut() == cached;
- }
- self.with(|op| {
- // Keep one reference so the pointer stays valid for the process
(a second thread
- // racing here leaks one more, which is harmless for a registry
singleton).
- // SAFETY: the raw pointer is never turned back into an
`ObjectArc`.
- let pointer = unsafe {
ObjectArc::into_raw(Op::into_data(op.clone())) };
- self.object.store(pointer.cast::<TVMFFIObject>().cast_mut(),
Ordering::Release);
- value.same_as(op)
- })
- .unwrap_or(false)
- }
-
- /// `use_op` applied to the operator, fetched into this thread's slot the
first time.
- fn with<R>(&self, use_op: impl FnOnce(&Op) -> R) -> Result<R> {
- let slot = *self.slot.get_or_init(|| SLOTS.fetch_add(1,
Ordering::Relaxed));
- let cached = FETCHED.with(|fetched|
fetched.borrow().get(slot).cloned().flatten());
- let op = match cached {
- Some(op) => op,
- None => {
- let op = Op::get(self.name).or_else(|missing| {
- let registered = tvm_ffi::run_static_init_blocks();
- Op::get(self.name).map_err(|_| match registered {
- Ok(()) => missing,
- Err(failed) => {
- let message = format!("{}\n{}", missing.message(),
failed.message());
- Error::new(missing.kind(), &message, "")
- }
- })
- })?;
- FETCHED.with(|fetched| {
- let mut fetched = fetched.borrow_mut();
- if fetched.len() <= slot {
- fetched.resize(slot + 1, None);
- }
- fetched[slot] = Some(op.clone());
- });
- op
- }
- };
- Ok(use_op(&op))
- }
-}
-
/// The text of `doc` without the indentation its lines share and without the
blank lines around
/// it. Text that is not indented is unchanged.
fn doc_text(doc: &str) -> std::string::String {
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index f92a0945e..83fc8b6fb 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -110,7 +110,7 @@
//! ```no_run
//! use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
//! use tvm_tirx::ir::sig::{self, IntExpr};
-//! use tvm_tirx::ir::{register_type, CallEffectKind, CallObj, OpDef,
OpHandle, Type};
+//! use tvm_tirx::ir::{register_type, CallEffectKind, CallObj, Op, OpDef,
Type};
//! use tvm_tirx::tvm_ffi::{self, Result};
//! # use tvm_tirx::ir::TypeObj;
//! # use tvm_tirx::tvm_ffi::{def_object_ref, rs_class, ObjectRefCore};
@@ -133,13 +133,15 @@
//! register_type::<PairType>()?;
//! }
//!
-//! static CLAMP: OpHandle = OpHandle::new("my_ext.clamp");
+//! let clamp = Op::get("my_ext.clamp")?;
+//! # Ok::<(), tvm_ffi::Error>(())
//! ```
//!
-//! 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
[`tvm_ffi::run_static_init_blocks`], which
-//! reports the errors of the blocks that failed, as the TVM FFI function
`tirx_register_ir`.
+//! is. [`ir::Op::get`] then returns the operator, as C++ `Op::Get` does: the
registry keeps one
+//! object per name, so `call.op.same_as(&clamp)` tests whether a call is to
it. 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
//!
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 80d981346..10d521e50 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -18,8 +18,7 @@
*/
//! 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
+//! which run when the test binary loads, are used by the native registry,
`Call` constructor and validation, and their type inference, validators
//! and region contracts are native hooks; and types registered with them.
//!
//! The registry does not synchronize registration, so the tests take
`serial()`; the tests
@@ -32,8 +31,8 @@ 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_type, AttrKey, Attrs, Call, CallEffectKind, CallObj, DictAttrs,
Evaluate, Expr,
- FuncType, IntImm, Op, OpDef, OpHandle, PrimType, RegionStmt,
RegionStmtObj, StringImm, Type,
- TypeObj, Var, VarObj, UNKNOWN_LOC,
+ FuncType, IntImm, Op, OpDef, 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};
@@ -487,12 +486,6 @@ fn is_block_registered(block: &StaticInitBlock) -> 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");
-static LAUNCH_THREAD: OpHandle = OpHandle::new("tirx.launch_thread");
-static UNREGISTERED: OpHandle = OpHandle::new("rust_test.unregistered");
-
fn launch_body_params(call: &CallObj) -> Result<Array<Var>> {
let tag = call.args.get(0)?.try_cast::<StringImm>()?;
let extent_ty: PrimType = call.args.get(1)?.ty.clone().try_cast()?;
@@ -593,7 +586,6 @@ fn static_init_blocks_define_operators_and_types() {
// 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();
@@ -652,29 +644,26 @@ fn static_init_blocks_define_operators_and_types() {
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));
+ // The Rust `FInferType` types a Call of the operator. The registry keeps
one object per name,
+ // for every lookup and thread, so `same_as` tests which operator a Call
is to.
let x = var("x", "int32");
let call = call(&identity, vec![Expr::from(&x)]).unwrap();
assert!(call.ty.same_as(&x.ty), "the Rust FInferType types the Call");
- assert!(IDENTITY.is(&call.op));
- assert!(!LAUNCH_THREAD.is(&call.op));
- let name = std::thread::spawn(||
IDENTITY.get().unwrap().name.as_str().to_string());
+ assert!(call.op.same_as(&Op::get("rust_test.defined_identity").unwrap()));
+ let native = Op::get("tirx.launch_thread").unwrap();
+ assert!(!call.op.same_as(&native));
+ let name = std::thread::spawn(|| {
+
Op::get("rust_test.defined_identity").unwrap().name.as_str().to_string()
+ });
assert_eq!(name.join().unwrap(), "rust_test.defined_identity");
-
- // A handle to an operator C++ registers looks it up; one to an operator
nothing registers
- // fails, with the failures of the blocks.
-
assert!(LAUNCH_THREAD.get().unwrap().same_as(&Op::get("tirx.launch_thread").unwrap()));
assert_error(
- UNREGISTERED.get(),
+ Op::get("rust_test.unregistered"),
ATTRIBUTE_ERROR,
- "Operator rust_test.unregistered is not registered\n2 static init
blocks failed:",
+ "Operator rust_test.unregistered is not registered",
);
- assert!(!UNREGISTERED.is(&call.op));
// The region hooks are native function pointers.
- let launch = LAUNCH.get().unwrap();
+ let launch = Op::get("rust_test.launch").unwrap();
for key in ["FRegionGetBodyParams", "FRegionValidate"] {
let hooks = Op::attr_map::<*mut c_void>(key).unwrap().unwrap();
assert!(hooks.get(&launch).unwrap().is_some(), "{key}");
@@ -682,7 +671,6 @@ fn static_init_blocks_define_operators_and_types() {
// The `RegionStmt` constructor checks a region of the Rust operator as it
checks one of
// `tirx.launch_thread`: the same outcome, and the same error but for the
operator's name.
- let native = LAUNCH_THREAD.get().unwrap();
let extent: Expr = int32(8);
let wide: Expr = IntImm::from_dtype(parse_dtype("int64"), 8,
UNKNOWN_LOC).unwrap().into();
let flag: Expr = IntImm::bool(true).into();
@@ -712,7 +700,7 @@ fn static_init_blocks_define_operators_and_types() {
}
// A panic in a hook becomes an error instead of unwinding into C++.
- let panicking = PANICKING.get().unwrap();
+ 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(
