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 4d894c4e7 [FEAT][RUST] Complete OpDef chains when they are dropped
(#116)
4d894c4e7 is described below
commit 4d894c4e7e2985653adbc6bd8ca2eb006069a7c0
Author: Linzhang Li <[email protected]>
AuthorDate: Wed Oct 28 15:13:14 2026 -0400
[FEAT][RUST] Complete OpDef chains when they are dropped (#116)
## Summary
Every `OpDef` chain had to end in `.op()?`, which returns the error of a
failed step and requests
the operator's APIs through `ir.expose_op`. In a
`tvm_ffi::static_init_block!`, where the operator
is only defined, the `.op()?` carried nothing but those two duties, and
forgetting it left the
operator unexposed and its error unreported (only `#[must_use]` warned).
C++ needs no such call:
an `OpDef` completes as its statement ends.
This PR completes a chain when it is dropped, at the end of its
statement:
```rust
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)
.set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
register_type::<PairType>()?;
}
```
`.op()` stays for code that needs the `Op` or wants the `Result`: it
completes the chain at once,
and the drop then does nothing. Every existing chain that ends in
`.op()?` compiles and behaves as
before.
---
rust/examples/define_ir.rs | 12 +++----
rust/src/ir/op_def.rs | 78 ++++++++++++++++++++++++++++++----------------
rust/src/lib.rs | 29 +++++++++--------
rust/tests/ir/op_def.rs | 57 ++++++++++++++++++++++++++-------
4 files changed, 114 insertions(+), 62 deletions(-)
diff --git a/rust/examples/define_ir.rs b/rust/examples/define_ir.rs
index 76e267258..f582c7179 100644
--- a/rust/examples/define_ir.rs
+++ b/rust/examples/define_ir.rs
@@ -70,21 +70,18 @@ tvm_ffi::static_init_block! {
sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
))
.set_infer_type(first_arg_type)
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure)
- .op()?;
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
// `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)
- .op()?;
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation);
// 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)
- .op()?;
+ .set_region_body_params(loop_index);
register_type::<ExamplePairType>()?;
}
@@ -101,7 +98,6 @@ mod effects {
")
.signature(())
.set_attr(FIXED_RETURN_TYPE, PrimType::void())
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque)
- .op()?;
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::Opaque);
}
}
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 5122bb02e..07d846578 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -56,31 +56,33 @@ use sig::{Checks, Specs};
/// C++ `OpDef`: register a new operator from Rust.
///
-/// 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
-/// [`tvm_ffi::static_init_block!`] does. An operator and its hooks live for
the rest of the
+/// Each step updates the registry at once, as in C++. After a failed step the
rest are skipped.
+/// The chain completes when it is dropped, at the end of its statement, as a
C++ `OpDef` does:
+/// the operator's APIs are requested in registered destinations, and the
error of a failed step
+/// panics, which a [`tvm_ffi::static_init_block!`] reports as the error of
its block.
+/// [`OpDef::op`] completes it instead, and returns the operator or the error.
The registry does
+/// not synchronize registration with readers, so register an operator before
building IR with it,
+/// as a [`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;
/// use tvm_tirx::ir::sig::{self, IntExpr};
/// use tvm_tirx::ir::{CallEffectKind, CallObj, OpDef, Type};
-/// 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())
/// }
///
-/// let op = 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)
-/// .op()?;
-/// # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
+/// 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_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
+/// }
/// ```
-#[must_use = "end the chain with `.op()?`, which returns its error and exposes
the operator"]
-pub struct OpDef(Result<Op>);
+pub struct OpDef(Option<Result<Op>>);
impl OpDef {
/// Register a new operator (`ir.RegisterOp`), with `doc` as its
documentation, less the
@@ -89,18 +91,13 @@ impl OpDef {
pub fn new(name: &str, doc: &str) -> Self {
let registered = tvm_ffi::cached_global_func!("ir.RegisterOp")
.call_tuple((String::from(name), String::from(doc_text(doc))));
- Self(registered.and_then(|_| Op::get(name)))
+ Self(Some(registered.and_then(|_| Op::get(name))))
}
- /// The completed operator, or the error of the step that failed. Request
its APIs in
- /// registered destinations through `ir.expose_op` after its signature and
attributes are ready.
- /// An absent receiver is a no-op; an installed receiver's error is
returned.
- pub fn op(self) -> Result<Op> {
- let op = self.0?;
- if let Ok(expose) = tvm_ffi::Function::get_global("ir.expose_op") {
- expose.call_tuple((op.name.clone(),))?;
- }
- Ok(op)
+ /// Complete the chain now: the operator, or the error of the step that
failed, which
+ /// dropping the chain would panic with.
+ pub fn op(mut self) -> Result<Op> {
+ complete(self.take())
}
/// C++ `OpDef::signature`: replace the argument descriptions and the
attrs type, and the
@@ -238,11 +235,39 @@ impl OpDef {
})
}
- fn then(self, step: impl FnOnce(&Op) -> Result<()>) -> Self {
- Self(self.0.and_then(|op| step(&op).map(|_| op)))
+ fn then(mut self, step: impl FnOnce(&Op) -> Result<()>) -> Self {
+ Self(Some(self.take().and_then(|op| step(&op).map(|_| op))))
+ }
+
+ /// The state of the chain, which leaves nothing for `drop` to complete.
+ fn take(&mut self) -> Result<Op> {
+ self.0.take().expect("an `OpDef` is completed once")
}
}
+impl Drop for OpDef {
+ fn drop(&mut self) {
+ // Dropped while a panic unwinds, the chain did not finish, and a
second panic would abort.
+ if std::thread::panicking() {
+ return;
+ }
+ if let Some(Err(error)) = self.0.take().map(complete) {
+ panic!("{}: {}", error.kind(), error.message());
+ }
+ }
+}
+
+/// Request the APIs of a defined operator in registered destinations through
`ir.expose_op`,
+/// after its signature and attributes are ready. An absent receiver is a
no-op; an installed
+/// receiver's error is returned.
+fn complete(state: Result<Op>) -> Result<Op> {
+ let op = state?;
+ if let Ok(expose) = tvm_ffi::Function::get_global("ir.expose_op") {
+ expose.call_tuple((op.name.clone(),))?;
+ }
+ Ok(op)
+}
+
/// C++ `ffi::reflection::AttrKey<T>`: the name of an operator attribute and
the type of its
/// values, which [`OpDef::set_attr`] takes. [`op_attr`] and
/// [`tirx::op_attr`](crate::tirx::op_attr) have the keys C++ declares, and
`AttrKey::new` names
@@ -338,8 +363,7 @@ impl From<CallEffectKind> for Any {
/// 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)
-/// .op()?;
+/// .set_infer_type(first_arg_type);
/// }
///
/// static CLAMP: OpHandle = OpHandle::new("my_ext.clamp");
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index a7d1aeaac..30ef5250d 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -134,8 +134,7 @@
//! sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
//! ))
//! .set_infer_type(index_type)
-//! .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure)
-//! .op()?;
+//! .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
//! register_type::<PairType>()?;
//! }
//!
@@ -149,19 +148,19 @@
//!
//! # 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`), which registers the
operators and types the
-//! 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 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
-//! publish constructors, such as `T.<Name>` for `tirx.<Name>Type`.
Handwritten Python bindings take
+//! 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`), which registers the operators and
types the 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, when its [`ir::OpDef`] chain ends, 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 publish
+//! constructors, such as `T.<Name>` for `tirx.<Name>Type`. Handwritten Python
bindings take
//! precedence in each destination. Without receivers, native registration
proceeds without API
//! exposure.
//!
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index f0a66987f..c50390505 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -438,20 +438,17 @@ tvm_ffi::static_init_block! {
")
.signature(sig::arg::<Expr>("value", "The value."))
.set_infer_type(identity_type)
- .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation)
- .op()?;
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::ExprAnnotation);
// `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)
- .op()?;
+ .set_region_validate(no_results_or_attributes);
OpDef::new("rust_test.panicking_region", "")
- .set_region_body_params(|_: &CallObj| -> Result<Array<Var>> {
panic!("hook bug") })
- .op()?;
+ .set_region_body_params(|_: &CallObj| -> Result<Array<Var>> {
panic!("hook bug") });
register_type::<BlockType>()?;
}
@@ -461,13 +458,13 @@ mod elsewhere {
tvm_ffi::static_init_block! {
OpDef::new("rust_test.defined_elsewhere", "An operator in another
module.")
- .set_attr(FIXED_RETURN_TYPE, PrimType::void())
- .op()?;
+ .set_attr(FIXED_RETURN_TYPE, PrimType::void());
}
}
// 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.
+// name C++ registers, which `.op()?` returns, and a second custom validator,
which the chain panics
+// with when it is dropped, after which the block stops.
mod failing {
use super::*;
@@ -476,8 +473,8 @@ mod failing {
}
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()?;
+ OpDef::new("rust_test.defined_broken",
"").set_validator(accept).set_validator(accept);
+ OpDef::new("rust_test.after_broken", "");
}
}
@@ -563,6 +560,38 @@ fn region(
})
}
+/// The text of a panic `catch_unwind` caught.
+fn panic_text(payload: Box<dyn std::any::Any + Send>) -> String {
+ match payload.downcast::<String>() {
+ Ok(text) => *text,
+ Err(payload) =>
payload.downcast_ref::<&str>().unwrap_or(&"").to_string(),
+ }
+}
+
+#[test]
+fn a_chain_completes_when_it_is_dropped() {
+ let _serial = serial();
+ // Dropped at the end of its statement, the chain exposes the operator.
+ OpDef::new("rust_test.dropped", "").set_attr(CALL_EFFECT_KIND,
CallEffectKind::Pure);
+ let effect = Op::attr_map::<i64>("TCallEffectKind").unwrap().unwrap();
+ assert_eq!(effect.get(&Op::get("rust_test.dropped").unwrap()).unwrap(),
Some(1));
+
+ // A failed chain panics with its error when it is dropped.
+ let caught = std::panic::catch_unwind(|| {
+ OpDef::new("tirx.launch_thread", "");
+ });
+ let text = panic_text(caught.unwrap_err());
+ assert!(text.starts_with("AttributeError: "), "{text}");
+ assert!(text.ends_with("Operator tirx.launch_thread is registered
before"), "{text}");
+
+ // Dropped while another panic unwinds, it does not panic again, which
would abort.
+ let caught = std::panic::catch_unwind(|| {
+ let _def = OpDef::new("tirx.launch_thread", "");
+ panic!("the first panic");
+ });
+ assert_eq!(panic_text(caught.unwrap_err()), "the first panic");
+}
+
#[test]
fn static_init_blocks_define_operators_and_types() {
let _serial = serial();
@@ -586,8 +615,12 @@ fn static_init_blocks_define_operators_and_types() {
let message = error.message();
assert!(message.starts_with("2 static init blocks failed:"), "{message}");
assert!(message.contains("op_def::failing:"), "{message}");
+ // A chain that is dropped reports its error as the panic of the block.
assert!(
- message.contains("ValueError: Validator of rust_test.defined_broken is
already registered"),
+ message.contains(
+ "InternalError: panicked: ValueError: Validator of
rust_test.defined_broken is already \
+ registered"
+ ),
"{message}"
);
assert!(message.contains("AttributeError: "), "{message}");
