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 75b9411b5 [FEAT][RUST] Register types of register_types! blocks for
Python and TVMScript (#87)
75b9411b5 is described below
commit 75b9411b5a176b67342211a0368ac0a53ff5109e
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Oct 26 18:47:02 2026 -0400
[FEAT][RUST] Register types of register_types! blocks for Python and
TVMScript (#87)
An IR type defined in Rust could be built and printed, but Python could
not construct it, a JSON
graph could not read it back, and TVMScript could not parse it:
`rs_class` registers the fields of
a type, but not a way to make one. A C++ type gets that from
`refl::ObjectDef`, which also
registers a creator of an empty object that tvm-ffi fills through the
reflection of the fields.
A type is now written as before, an `rs_class` object and its
`def_object_ref!`, inside a new
`register_types!` block. As a `register_ops!` block does (#42), the
block runs when its library
loads, once per process, and Python then treats the type as it treats a
C++ one:
```rust
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);
}
```
```python
import tvm_ffi
import tvm_tirx
from tvm_tirx import ir
from tvm_tirx.script import parser
from tvm_tirx.script import tirx as T
tvm_ffi.load_module("target/release/libmy_ext.so") # registers
my_ext.PairType
my_ext = parser.my_ext
pair = my_ext.Pair(T.int32, second=T.float32)
@T.function
def func(p: my_ext.Pair(T.int32, T.float32)):
T.evaluate(p)
assert tvm_ffi.structural_equal(tvm_tirx.script.parse(func.script(),
extra_vars={"T": T}), func)
assert tvm_ffi.structural_equal(ir.load_json(ir.save_json(func)), func)
```
---
docs/reference/python/tirx.script.rst | 1 +
python/tvm_tirx/backend/rust/__init__.py | 6 +-
python/tvm_tirx/tirx/script/ir_builder/op.py | 79 +++++++-
rust/Cargo.toml | 4 +-
rust/examples/{define_ops.rs => define_ir.rs} | 28 ++-
rust/src/ir.rs | 2 +
rust/src/ir/op_def.rs | 115 +++++++-----
rust/src/ir/type_def.rs | 215 ++++++++++++++++++++++
rust/src/lib.rs | 43 +++--
rust/tests/ir/op_def.rs | 76 ++++++--
tests/python/ir/test_op_rust.py | 86 ++++++---
tests/python/tirx/script/test_expose_script_op.py | 57 +++++-
12 files changed, 585 insertions(+), 127 deletions(-)
diff --git a/docs/reference/python/tirx.script.rst
b/docs/reference/python/tirx.script.rst
index e74a3e4ab..51a755b35 100644
--- a/docs/reference/python/tirx.script.rst
+++ b/docs/reference/python/tirx.script.rst
@@ -369,6 +369,7 @@ Functions
max
register_script_namespace
expose_script_op
+ expose_script_type
abs
acos
acosh
diff --git a/python/tvm_tirx/backend/rust/__init__.py
b/python/tvm_tirx/backend/rust/__init__.py
index 74d8e5ea7..38e1e2d6c 100644
--- a/python/tvm_tirx/backend/rust/__init__.py
+++ b/python/tvm_tirx/backend/rust/__init__.py
@@ -37,9 +37,9 @@ def register_backend():
except (OSError, FileNotFoundError, RuntimeError):
pass
else:
- # 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()
+ # 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.
+
tvm_ffi.load_module(_LOADED_LIBS["tvm_tirx_rust"]._name).tirx_register_ir()
return None
diff --git a/python/tvm_tirx/tirx/script/ir_builder/op.py
b/python/tvm_tirx/tirx/script/ir_builder/op.py
index 4285d8232..f1233f0be 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/op.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/op.py
@@ -282,9 +282,28 @@ def expose_script_op(name: str) -> Any:
builder : Callable
The builder the printed calls name.
"""
+ op = _ir.Op.get(name)
+ leaf, lookup, bind = _script_slot(name)
+ existing = lookup(leaf)
+ if existing is not None:
+ if not callable(existing):
+ raise ValueError(f"Op {name!r} conflicts with a non-callable
Python attribute")
+ return existing
+ call = _ir.op._make_op_api(op, __name__)
+ if op.get_attr("FRegionGetBodyParams") is None:
+ call = _auto_emit_void_call(call)
+ bind(leaf, call)
+ return call
+
+
+def _script_slot(name: str):
+ """The last part of ``name``, and the lookup and binding of a name where
its first parts are.
+
+ ``tirx.<path>``, ``prim.<path>`` and ``ir.<path>`` are printed as
``T.<path>``, and the first
+ part of any other name is a parser namespace. A missing namespace of the
path is created.
+ """
from tvm_tirx.script import parser # pylint:
disable=import-outside-toplevel
- op = _ir.Op.get(name)
parts = name.split(".")
# None stands for the parser's namespaces, which hold the first part of
other names.
container = None
@@ -307,20 +326,61 @@ def expose_script_op(name: str) -> Any:
if lookup(part) is None:
bind(part, SimpleNamespace())
container = lookup(part)
+ return leaf, lookup, bind
+
+
+def expose_script_type(type_key: str) -> Any:
+ """Give a registered Type the class its printed constructor names, and
return it.
+
+ The printer spells a Type ``<namespace>.<Name>Type`` without a printer
hook as a call of
+ ``<Name>`` in its namespace, with its fields as arguments, which is placed
as
+ :func:`expose_script_op` places a builder. The class is the one registered
for the type,
+ or, when the type has none or only the fallback class of an object that
reached Python
+ first, a class registered now, which constructs the type through its
``__ffi_init__``. As for
+ an Op, an existing callable keeps the name, so a hand-written builder
takes precedence, except
+ a class of the type, which is replaced: a hand-written class takes the
name when this is called
+ after it is registered. Types of Rust ``register_types!`` blocks get their
classes when they
+ are registered.
+
+ Parameters
+ ----------
+ type_key : str
+ Key of a registered Type.
+
+ Returns
+ -------
+ cls : Callable
+ The class, or the builder, the printed constructors name.
+ """
+ type_key = str(type_key)
+ info = tvm_ffi.core._lookup_or_register_type_info_from_type_key(type_key)
+ namespace, _, name = type_key.rpartition(".")
+ cls = info.type_cls
+ # The fallback class tvm_ffi makes for an object of a type without a class
has no __init__.
+ if cls is None or "__init__" not in vars(cls):
+ parent = info.parent_type_info
+ base = parent.type_cls
+ if base is None or "__init__" not in vars(base):
+ base = expose_script_type(parent.type_key)
+ attrs = {"__slots__": (), "__module__": namespace}
+ cls = tvm_ffi.register_object(type_key)(type(name, (base,), attrs))
+ # The printer drops the "Type" suffix (DefaultTypeTranslate).
+ if len(name) > 4 and name.endswith("Type"):
+ name = name[:-4]
+ leaf, lookup, bind = _script_slot(f"{namespace}.{name}")
existing = lookup(leaf)
- if existing is not None:
+ existing_info = getattr(existing, "__tvm_ffi_type_info__", None)
+ if existing is not None and getattr(existing_info, "type_key", None) !=
type_key:
if not callable(existing):
- raise ValueError(f"Op {name!r} conflicts with a non-callable
Python attribute")
+ raise ValueError(f"Type {type_key!r} conflicts with a non-callable
Python attribute")
return existing
- call = _ir.op._make_op_api(op, __name__)
- if op.get_attr("FRegionGetBodyParams") is None:
- call = _auto_emit_void_call(call)
- bind(leaf, call)
- return call
+ bind(leaf, cls)
+ return cls
-# Rust operators call this when they are registered.
+# Rust operators and types call these when they are registered.
tvm_ffi.register_global_func("tirx.script.expose_script_op", expose_script_op)
+tvm_ffi.register_global_func("tirx.script.expose_script_type",
expose_script_type)
abs = _auto_emit_void_call(_tir_op.abs) # pylint: disable=redefined-builtin
@@ -897,6 +957,7 @@ __all__ = [
"exp2",
"exp10",
"expose_script_op",
+ "expose_script_type",
"fabs",
"floor",
"floordiv",
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 3d63e5dbb..80f96d402 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -33,9 +33,9 @@ crate-type = ["rlib", "cdylib"]
linkme = "0.3"
tvm-ffi = { path = "../3rdparty/tvm-ffi/rust/tvm-ffi" }
-# A library of operators defined in Rust, which
`tests/python/ir/test_op_rust.py` loads.
+# A library of operators and a type defined in Rust, which
`tests/python/ir/test_op_rust.py` loads.
[[example]]
-name = "define_ops"
+name = "define_ir"
crate-type = ["cdylib"]
# Standalone crate: not a member of any enclosing workspace.
diff --git a/rust/examples/define_ops.rs b/rust/examples/define_ir.rs
similarity index 76%
rename from rust/examples/define_ops.rs
rename to rust/examples/define_ir.rs
index 4bdb49e32..d05332913 100644
--- a/rust/examples/define_ops.rs
+++ b/rust/examples/define_ir.rs
@@ -17,18 +17,21 @@
* under the License.
*/
-//! Operators defined in Rust, in a library of their own.
+//! Operators and a type defined in Rust, in a library of their own.
//!
-//! `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.
+//! `cargo build --example define_ir` builds this as
`target/debug/examples/libdefine_ir.so`,
+//! which exports `tirx_register_ir`, as every library built on `tvm_tirx`
does.
//! `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:`.
+//! `T.example_mark()`, `my_ext.hint(i, attrs={..})` and `with
T.example_loop(n) as j:`; and
+//! registers the type, with its class `T.ExamplePair`.
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, Var, UNKNOWN_LOC};
-use tvm_tirx::tvm_ffi::{Array, ObjectRefCast, Result};
+use tvm_tirx::ir::{
+ CallEffectKind, CallObj, DictAttrs, Expr, OpDef, PrimType, Type, TypeObj,
Var, UNKNOWN_LOC,
+};
+use tvm_tirx::tvm_ffi::{self, rs_class, Array, ObjectRefCast, Result};
fn first_arg_type(call: &CallObj) -> Result<Type> {
Ok(call.args.get(0)?.ty.clone())
@@ -61,6 +64,19 @@ tvm_tirx::register_ops! {
.set_region_body_params(loop_index);
}
+// 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);
+}
+
// A block can be anywhere in the library.
mod effects {
use super::*;
diff --git a/rust/src/ir.rs b/rust/src/ir.rs
index 27872d54b..ca5c92e2d 100644
--- a/rust/src/ir.rs
+++ b/rust/src/ir.rs
@@ -45,6 +45,7 @@ pub mod stmt;
pub(crate) mod support;
pub mod transform;
pub mod r#type;
+pub mod type_def;
pub mod unique_name_supply;
pub use attrs::*;
@@ -58,6 +59,7 @@ pub use op::*;
pub use op_def::*;
pub use stmt::*;
pub use r#type::*;
+pub use type_def::*;
pub use unique_name_supply::*;
// tvm-ffi-stubgen(begin): import-section
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index bd5adec79..05b055f6d 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -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) 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 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.
+//! 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.
use std::any::TypeId;
use std::cell::RefCell;
@@ -412,7 +412,7 @@ impl OpHandle {
Some(op) => op,
None => {
let op = Op::get(self.name).or_else(|missing| {
- let registered = register_defined_ops();
+ let registered = register_defined_ir();
Op::get(self.name).map_err(|_| match registered {
Ok(()) => missing,
Err(failed) => {
@@ -435,10 +435,12 @@ impl OpHandle {
}
}
-/// A [`register_ops!`](crate::register_ops) block, gathered at link time.
+/// A [`register_ops!`](crate::register_ops) or
[`register_types!`](crate::register_types) block,
+/// gathered at link time.
#[doc(hidden)]
-pub struct OpRegistration {
- /// The module and line of the block.
+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
@@ -450,15 +452,17 @@ pub struct OpRegistration {
outcome: OnceLock<std::result::Result<(), std::string::String>>,
}
-/// The blocks of [`register_ops!`](crate::register_ops), gathered at link
time from the whole
-/// program or library.
+/// 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 OP_REGISTRATIONS: [OpRegistration];
+pub static REGISTRATIONS: [Registration];
-impl OpRegistration {
+impl Registration {
#[doc(hidden)]
pub const fn __new(
+ kind: &'static str,
module: &'static str,
line: u32,
version: &'static str,
@@ -466,7 +470,7 @@ impl OpRegistration {
register: fn() -> Result<()>,
) -> Self {
let text_hash = text_hash(text);
- Self { module, line, version, text_hash, register, outcome:
OnceLock::new() }
+ 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
@@ -478,11 +482,11 @@ impl OpRegistration {
}
/// 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
+ /// `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_op_block.{}", self.identity())
+ format!("ir.registered_block.{}", self.identity())
}
/// Run the block if it has not run. The caller that runs a failing block
gets its error;
@@ -501,15 +505,15 @@ impl OpRegistration {
(Some(error), _) => Err(error),
(None, Ok(())) => Ok(()),
(None, Err(text)) => {
- let message = format!("The register_ops! block at {} failed:
{text}", self.place());
+ 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 its
- /// operators are registered once, with the hooks of the first library
that runs it. The block
+ /// 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();
@@ -541,55 +545,59 @@ const fn text_hash(text: &str) -> u64 {
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. 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.
+/// 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_ops`, which `import tvm_tirx` calls on `libtvm_tirx_rust`
to raise them. Where
+/// `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_ops() -> Result<()> {
- let mut blocks: Vec<&OpRegistration> = OP_REGISTRATIONS.iter().collect();
+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<(&OpRegistration, Error)> =
+ 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 message = format!("{} register_ops! blocks failed:",
failures.len());
+ 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) 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.
+/// 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_ops) {
+ match catch_unwind(register_defined_ir) {
Ok(Ok(())) => {}
Ok(Err(error)) => eprintln!(
- "{library}: registering operators failed: {}: {}",
+ "{library}: registering operators and types failed: {}: {}",
error.kind(),
error.message()
),
- Err(_) => eprintln!("{library}: registering operators panicked"),
+ 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 it
-/// registers keep its hooks: unloading it, as `tvm_ffi.load_module` does when
Python drops the
+/// 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> {
@@ -679,10 +687,10 @@ fn doc_text(doc: &str) -> std::string::String {
///
/// 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
+/// 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_ops`] (`tirx_register_ops`) and the handles report it
again. Where the
+/// [`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.
///
@@ -695,22 +703,33 @@ fn doc_text(doc: &str) -> std::string::String {
#[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 _: () = {
- fn register() -> $crate::tvm_ffi::Result<()> {
- $($crate::ir::OpDef::op($def)?;)*
- Ok(())
- }
-
#[$crate::__private::linkme::distributed_slice($crate::ir::OP_REGISTRATIONS)]
+
#[$crate::__private::linkme::distributed_slice($crate::ir::REGISTRATIONS)]
#[linkme(crate = $crate::__private::linkme)]
- static REGISTRATION: $crate::ir::OpRegistration =
$crate::ir::OpRegistration::__new(
+ static REGISTRATION: $crate::ir::Registration =
$crate::ir::Registration::__new(
+ $kind,
module_path!(),
line!(),
match option_env!("CARGO_PKG_VERSION") {
Some(version) => version,
None => "",
},
- stringify!($($def);*),
- register,
+ $text,
+ $register,
);
#[used]
#[cfg_attr(target_os = "linux", link_section = ".init_array")]
diff --git a/rust/src/ir/type_def.rs b/rust/src/ir/type_def.rs
new file mode 100644
index 000000000..5f8165d36
--- /dev/null
+++ b/rust/src/ir/type_def.rs
@@ -0,0 +1,215 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * 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.
+//!
+//! `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.
+
+use tvm_ffi::tvm_ffi_sys::{TVMFFIByteArray, TVMFFITypeRegisterAttr};
+use tvm_ffi::{Any, AnyCompatible, Array, Function, ObjectCore, ObjectRefCore,
Result, String};
+
+use super::expr::{Expr, IntImm};
+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.
+pub trait Placeholder {
+ /// The value.
+ fn placeholder() -> Self;
+}
+
+macro_rules! default_placeholder {
+ ($($ty:ty),*) => {
+ $(impl Placeholder for $ty {
+ fn placeholder() -> Self {
+ Self::default()
+ }
+ })*
+ };
+}
+
+default_placeholder!(bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64);
+
+impl Placeholder for String {
+ fn placeholder() -> Self {
+ String::from("")
+ }
+}
+
+impl<T: AnyCompatible + Clone + 'static> Placeholder for Array<T> {
+ fn placeholder() -> Self {
+ Array::new(vec![])
+ }
+}
+
+impl<T> Placeholder for Option<T> {
+ fn placeholder() -> Self {
+ None
+ }
+}
+
+impl Placeholder for Location {
+ fn placeholder() -> Self {
+ Location::unknown()
+ }
+}
+
+impl Placeholder for TypeObj {
+ fn placeholder() -> Self {
+ TypeObj::new(Location::unknown())
+ }
+}
+
+impl Placeholder for Type {
+ fn placeholder() -> Self {
+ MissingType::new().into()
+ }
+}
+
+impl Placeholder for PrimType {
+ fn placeholder() -> Self {
+ PrimType::void()
+ }
+}
+
+impl Placeholder for Expr {
+ fn placeholder() -> Self {
+ IntImm::new(Location::unknown(), PrimType::int(32), 0).into()
+ }
+}
+
+/// 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
+/// Python has imported `tvm_tirx`, its TVMScript class
(`tirx.script.expose_script_type`).
+#[doc(hidden)]
+pub fn __register_type<R>() -> Result<()>
+where
+ R: ObjectRefCore + AnyCompatible + Placeholder + 'static,
+{
+ let type_index = R::ContainerType::type_index();
+ let new = Function::from_typed(|| -> Result<R> { Ok(R::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("tirx.script.expose_script_type")
{
+
expose.call_tuple((String::from(support::type_key(type_index).as_str()),))?;
+ }
+ Ok(())
+}
+
+/// Register `value` as the type attribute `name` of the type of `type_index`.
+fn register_type_attr(type_index: i32, name: &str, mut value: Any) ->
Result<()> {
+ // SAFETY: `name` and `value` outlive the call, which copies them.
+ let name = unsafe { TVMFFIByteArray::from_str(name) };
+ tvm_ffi::check_safe_call!(unsafe {
+ 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.
+/// Python then constructs a type from its fields, `my_ext.Pair(T.int32,
second=T.float32)`, as it
+/// constructs a C++ type, a JSON graph holds it, and, when Python has
imported `tvm_tirx` before it
+/// loads the library, the type gets the class its TVMScript constructor
names. 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 261f235f8..01c5db7ef 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -130,29 +130,31 @@
//!
//! 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_ops`],
which reports the
-//! errors of the blocks that failed, as the TVM FFI function
`tirx_register_ops`.
+//! 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`.
//!
//! # 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 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.
+//! 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. 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, and a type of a
+//! [`register_types!`](crate::register_types) block gets the class its
TVMScript constructor names
+//! (`T.<Name>` for `tirx.<Name>Type`) from `expose_script_type`.
+//!
+//! A library of operators and types of its own, a `cdylib` built on the crate
as
+//! `examples/define_ir.rs` is, registers them 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 what it defines
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
//!
@@ -197,6 +199,9 @@
//! 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.
//!
//! # Generated and hand-written code
//!
@@ -298,9 +303,9 @@ 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_ops`, the TVM FFI function
`tirx_register_ops` of
+// `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_ops,
ir::register_defined_ops);
+tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir,
ir::register_defined_ir);
/// What the exported macros expand to.
#[doc(hidden)]
diff --git a/rust/tests/ir/op_def.rs b/rust/tests/ir/op_def.rs
index 928cfa5b5..0601595bb 100644
--- a/rust/tests/ir/op_def.rs
+++ b/rust/tests/ir/op_def.rs
@@ -18,7 +18,7 @@
*/
//! Tests of `ir::op_def`: operators registered from Rust, by an `OpDef` or in
`register_ops!`
-//! blocks, which `register_defined_ops` or the first lookup through an
`OpHandle` runs, are used by
+//! blocks, which `register_defined_ir` or the first lookup through an
`OpHandle` runs, are used by
//! the native registry, `Call` constructor and validation, and their type
inference, validators
//! and region contracts are native hooks.
//!
@@ -31,18 +31,19 @@ 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_ops, AttrKey, Attrs, Call, CallEffectKind, CallObj,
DictAttrs, Evaluate, Expr,
- FuncType, IntImm, Op, OpDef, OpHandle, OpRegistration, PrimType,
RegionStmt, RegionStmtObj,
- StringImm, Type, Var, VarObj, OP_REGISTRATIONS, UNKNOWN_LOC,
+ 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,
};
use tvm_tirx::tirx::{CallFFIKernelAttr, TensorVar};
use tvm_tirx::tvm_ffi::tvm_ffi_sys::TVMFFITypeIndex;
use tvm_tirx::tvm_ffi::{
- Any, Array, Error, ErrorKind, Function, Map, ObjectRefCast, ObjectRefCore,
Result,
- String as FfiString, ATTRIBUTE_ERROR, INTERNAL_ERROR, TYPE_ERROR,
VALUE_ERROR,
+ 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 {
@@ -442,7 +443,25 @@ mod elsewhere {
}
}
-// Two blocks that fail, which `register_defined_ops` reports together while
it runs the others: a
+// 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);
+ }
+}
+
+// Two blocks that fail, which `register_defined_ir` 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::*;
@@ -458,17 +477,17 @@ mod failing {
}
/// The blocks of `module` in this test binary.
-fn blocks_of(module: &str) -> Vec<&'static OpRegistration> {
+fn blocks_of(module: &str) -> Vec<&'static Registration> {
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()
+ REGISTRATIONS.iter().filter(|block| of_module(&block.identity())).collect()
}
-fn is_block_registered(block: &OpRegistration) -> bool {
+fn is_block_registered(block: &Registration) -> bool {
Function::get_global(&block.marker()).is_ok()
}
@@ -556,9 +575,9 @@ 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_ops` reports the blocks that failed together, each
with its first error;
+ // `register_defined_ir` 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_ops().unwrap_err();
+ let error = register_defined_ir().unwrap_err();
let message = error.message();
assert!(message.starts_with("2 register_ops! blocks failed:"),
"{message}");
assert!(message.contains("op_def::failing:"), "{message}");
@@ -569,7 +588,30 @@ 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_ops().unwrap_err().message(), message);
+ assert_eq!(register_defined_ir().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.
+ 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);
+ let built = init
+ .call_packed(&[
+ AnyView::from(&elem),
+ AnyView::from(&kwargs),
+ AnyView::from(&count),
+ AnyView::from(&four),
+ ])
+ .unwrap();
+ let block = BlockType::try_from(built).unwrap();
+ assert!(block.elem.same_as(&elem));
+ assert_eq!(block.count, 4);
+ let positional = init.call_packed(&[AnyView::from(&elem),
AnyView::from(&four)]);
+ assert!(positional.is_err_and(|error| error.kind() == TYPE_ERROR),
"`count` is keyword-only");
+ let json = global("ffi.ToJSONGraph").call_packed(&[AnyView::from(&block),
AnyView::new()]);
+ let read =
global("ffi.FromJSONGraph").call_packed(&[AnyView::from(&json.unwrap())]);
+ let read = BlockType::try_from(read.unwrap()).unwrap();
+ assert!(!read.same_as(&block) && structural_equal(&read, &block));
// 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.
@@ -662,7 +704,7 @@ fn a_block_another_library_registered_is_skipped() {
}
// 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 block = || Registration::__new("register_ops!", "ext::shared", 3,
"0.1.0", text, register);
let (first, second) = (block(), block());
assert_eq!(first.marker(), second.marker());
@@ -675,7 +717,7 @@ fn a_block_another_library_registered_is_skipped() {
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);
+ let other = Registration::__new("register_ops!", "ext::shared", 3,
"0.2.0", text, register);
assert!(other.register().unwrap_err().message().contains("is registered
before"));
}
@@ -684,7 +726,9 @@ 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 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}");
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index 57928ae5c..3e2e6e863 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -16,11 +16,11 @@
# under the License.
"""Operators defined in Rust are registered when their library loads, as C++
ones are.
-``rust/examples/define_ops.rs`` is built by cargo into a library of its own,
which a fresh
+``rust/examples/define_ir.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.
+uses C++ ones: registering them also gives them their TVMScript builders. Its
type is registered
+with them. A copy of the library is a second library with the same blocks, as
libraries that link
+the same crate have.
"""
import os
@@ -55,8 +55,8 @@ def find_cargo():
@pytest.fixture(scope="module")
-def define_ops_library():
- """The example library of Rust operators, built by cargo."""
+def define_ir_library():
+ """The example library of Rust operators and types, built by cargo."""
cargo = find_cargo()
if cargo is None or not shutil.which("tvm-ffi-config") or not
shutil.which("tvm-tirx-config"):
pytest.skip(
@@ -65,11 +65,11 @@ def define_ops_library():
)
manifest = RUST_DIR / "Cargo.toml"
subprocess.run(
- [cargo, "build", "--manifest-path", str(manifest), "--example",
"define_ops"], check=True
+ [cargo, "build", "--manifest-path", str(manifest), "--example",
"define_ir"], check=True
)
target_dir = Path(os.environ.get("CARGO_TARGET_DIR", RUST_DIR / "target"))
- name = {"darwin": "libdefine_ops.dylib", "win32": "define_ops.dll"}.get(
- sys.platform, "libdefine_ops.so"
+ name = {"darwin": "libdefine_ir.dylib", "win32": "define_ir.dll"}.get(
+ sys.platform, "libdefine_ir.so"
)
return target_dir / "debug" / "examples" / name
@@ -145,7 +145,7 @@ def test_library_is_loaded_on_import():
assert "tvm_tirx_rust" in _LOADED_LIBS
-def test_rust_operators_are_registered_by_their_library(define_ops_library):
+def test_rust_operators_are_registered_by_their_library(define_ir_library):
run_python(
"""
import sys
@@ -155,7 +155,7 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
# 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()
+ tvm_ffi.load_module(sys.argv[1])
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"]
@@ -175,13 +175,55 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
raise AssertionError("a float bound passed validation")
""",
ROUND_TRIP,
- libraries=[define_ops_library],
+ libraries=[define_ir_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)
+def test_rust_types_are_registered_by_their_library(define_ir_library):
+ run_python(
+ """
+ import sys
+
+ import tvm_ffi
+ import tvm_tirx
+ from tvm_tirx import ir
+ from tvm_tirx.script import tirx as T
+
+ # Registered with the operators when the library loads, the type gets
the class its
+ # TVMScript constructor names, which constructs it from its fields, as
a C++ type is
+ # constructed.
+ tvm_ffi.load_module(sys.argv[1])
+ pair = T.ExamplePair(T.int32, second=T.float32)
+ assert pair.script() == "T.ExamplePair(T.int32, T.float32)"
+ assert tvm_ffi.structural_equal(pair.first, T.int32)
+
+ # TVMScript prints and parses it, and a JSON graph holds it.
+ @T.function
+ def func(p: T.ExamplePair(T.int32, T.float32)):
+ T.evaluate(p)
+
+ text = func.script()
+ assert "p: T.ExamplePair(T.int32, T.float32)" in text, text
+ parsed = tvm_tirx.script.parse(text, extra_vars={"T": T})
+ assert tvm_ffi.structural_equal(parsed, func), parsed.script()
+ assert tvm_ffi.structural_equal(ir.load_json(ir.save_json(func)), func)
+
+ # A class written for it in Python takes the name.
+ @tvm_ffi.register_object("tirx.ExamplePairType")
+ class ExamplePair(ir.Type):
+ pass
+
+ assert T.expose_script_type("tirx.ExamplePairType") is ExamplePair
+ assert T.ExamplePair is ExamplePair
+ assert type(T.ExamplePair(T.int32, T.float32)) is ExamplePair
+ """,
+ libraries=[define_ir_library],
+ )
+
+
+def
test_a_library_registers_its_operators_when_python_loads_it(define_ir_library,
tmp_path):
+ copy = tmp_path / define_ir_library.name
+ shutil.copy(define_ir_library, copy)
stderr = run_python(
"""
import ctypes
@@ -216,12 +258,12 @@ def
test_a_library_registers_its_operators_when_python_loads_it(define_ops_libra
assert ir.Op.get("tirx.example_clamp").same_as(clamp)
""",
ROUND_TRIP,
- libraries=[define_ops_library, copy],
+ libraries=[define_ir_library, copy],
)
- assert "registering operators failed" not in stderr, stderr
+ assert "registering operators and types failed" not in stderr, stderr
-def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ops_library):
+def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ir_library):
stderr = run_python(
"""
import sys
@@ -235,15 +277,15 @@ def
test_a_block_that_fails_when_its_library_loads_is_reported(define_ops_librar
module = tvm_ffi.load_module(sys.argv[1])
ir.Op.get("tirx.example_clamp")
- # `tirx_register_ops` raises the error again.
+ # `tirx_register_ir` raises the error again.
try:
- module.tirx_register_ops()
+ module.tirx_register_ir()
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")
""",
- libraries=[define_ops_library],
+ libraries=[define_ir_library],
)
- assert "registering operators failed" in stderr, stderr
+ assert "registering operators and types failed" in stderr, stderr
assert "Operator tirx.example_mark is registered before" in stderr, stderr
diff --git a/tests/python/tirx/script/test_expose_script_op.py
b/tests/python/tirx/script/test_expose_script_op.py
index ece863efe..e5f311c31 100644
--- a/tests/python/tirx/script/test_expose_script_op.py
+++ b/tests/python/tirx/script/test_expose_script_op.py
@@ -14,13 +14,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TVMScript builders generated for registered operators."""
+"""TVMScript builders generated for registered operators, and classes for
registered types."""
import pytest
import tvm_ffi
+from tvm_ffi.dataclasses import py_class
import tvm_tirx
-from tvm_tirx.ir import Call, Op, PrimType, assert_structural_equal,
register_op_attr
+from tvm_tirx.ir import Call, Op, PrimType, Type, assert_structural_equal,
register_op_attr
from tvm_tirx.script import parser
from tvm_tirx.script import tirx as T
from tvm_tirx.tirx.script import ir_builder as builder
@@ -35,6 +36,32 @@ def define(name, args, result):
return Op.get(name)
+@py_class("tirx.TestScriptPairType")
+class ScriptPair(Type):
+ """A pair of types, printed as ``T.TestScriptPair(first, second)``."""
+
+ first: Type
+ second: Type
+
+
+@py_class("test_script_types.BoxType")
+class ScriptBox(Type):
+ """A type and a count, printed as ``test_script_types.Box(elem,
count)``."""
+
+ elem: Type
+ count: int
+
+
+@py_class("tirx.evaluateType")
+class ScriptEvaluate(Type):
+ """A type whose printed name, ``T.evaluate``, is taken by a builder."""
+
+
+@py_class("tirx.TestScriptConstantType")
+class ScriptConstant(Type):
+ """A type whose printed name, ``T.TestScriptConstant``, a test takes by a
non-callable."""
+
+
@pytest.fixture
def exposed():
"""The names a test exposes on T, which are removed after it."""
@@ -90,3 +117,29 @@ def test_existing_names_keep_their_values(exposed):
define("tirx.supports_mutable_declarations", ["x"], "int32")
with pytest.raises(ValueError, match="conflicts with a non-callable"):
T.expose_script_op("tirx.supports_mutable_declarations")
+
+
+def test_type_classes_round_trip_types_under_every_root(exposed):
+ exposed += ["TestScriptPair", "test_script_types"]
+ assert T.expose_script_type("tirx.TestScriptPairType") is ScriptPair
+ assert T.expose_script_type("test_script_types.BoxType") is ScriptBox
+ assert T.expose_script_type("tirx.TestScriptPairType") is T.TestScriptPair
+ test_script_types = parser.test_script_types
+
+ @T.function
+ def func(p: T.TestScriptPair(T.int32, T.float32), b:
test_script_types.Box(T.int8, 4)):
+ T.evaluate(p)
+ T.evaluate(b)
+
+ text = func.script()
+ assert "p: T.TestScriptPair(T.int32, T.float32)" in text, text
+ assert "b: test_script_types.Box(T.int8, 4)" in text, text
+ assert_structural_equal(tvm_tirx.script.parse(text, extra_vars={"T": T}),
func)
+
+
+def test_existing_names_keep_their_values_over_types(exposed):
+ exposed += ["TestScriptConstant"]
+ assert T.expose_script_type("tirx.evaluateType") is builder.evaluate
+ builder.TestScriptConstant = 1
+ with pytest.raises(ValueError, match="conflicts with a non-callable"):
+ T.expose_script_type("tirx.TestScriptConstantType")
