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 7ec946256 [FEAT][RUST] Generate TVMScript builders for registered
operators (#58)
7ec946256 is described below
commit 7ec9462564c2be39ad165f814963696b374c6ecf
Author: Linzhang Li <[email protected]>
AuthorDate: Sun Oct 25 14:40:19 2026 -0400
[FEAT][RUST] Generate TVMScript builders for registered operators (#58)
## Summary
An operator registered from Rust is now printed and parsed by TVMScript
without a builder written
for it in Python. `tvm_tirx.tirx.script.expose_script_op(name)`
generates the builder of a
registered operator from its registration, under the name the printer
gives its calls, and
`OpDef::op` calls it for every operator a `register_ops!` block defines,
once Python has imported
`tvm_tirx`. An operator registered from Python or C++ gets its builder
by calling it.
```rust
tvm_tirx::register_ops! {
OpDef::new("tirx.demo_clamp", "Clamp an index to the range `[0,
bound)`.")
.signature((sig::arg::<IntExpr>("index", ""),
sig::arg::<IntExpr>("bound", "")))
.set_infer_type(first_arg_type);
OpDef::new("my_ext.tiles", "Run the body for each tile in `[0,
count)`.")
.signature(sig::arg::<IntExpr>("count", ""))
.set_region_body_params(tile_index);
}
```
```python
tvm_ffi.load_module(path).tirx_register_ops() # the user's Rust library
@T.function
def kernel(i: T.int32, n: T.int32):
T.evaluate(T.demo_clamp(i, n))
with my_ext.tiles(n) as t:
T.evaluate(T.demo_clamp(t, n))
assert tvm_ffi.structural_equal(tvm_tirx.script.parse(kernel.script(),
extra_vars={"T": T}), kernel)
```
---
docs/reference/python/tirx.script.rst | 1 +
python/tvm_tirx/tirx/script/ir_builder/op.py | 66 ++++++++++++++++
rust/examples/define_ops.rs | 9 +--
rust/src/ir/op_def.rs | 9 ++-
rust/src/lib.rs | 7 +-
tests/python/ir/test_op_rust.py | 33 +++-----
tests/python/tirx/script/test_expose_script_op.py | 92 +++++++++++++++++++++++
7 files changed, 185 insertions(+), 32 deletions(-)
diff --git a/docs/reference/python/tirx.script.rst
b/docs/reference/python/tirx.script.rst
index 6aa50b31d..e74a3e4ab 100644
--- a/docs/reference/python/tirx.script.rst
+++ b/docs/reference/python/tirx.script.rst
@@ -368,6 +368,7 @@ Functions
min
max
register_script_namespace
+ expose_script_op
abs
acos
acosh
diff --git a/python/tvm_tirx/tirx/script/ir_builder/op.py
b/python/tvm_tirx/tirx/script/ir_builder/op.py
index cf815098c..4285d8232 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/op.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/op.py
@@ -18,8 +18,12 @@
from __future__ import annotations
+import sys
+from types import SimpleNamespace
from typing import Any
+import tvm_ffi
+
from tvm_tirx import ir
from tvm_tirx import ir as _ir
from tvm_tirx import tirx as _tir
@@ -258,6 +262,67 @@ def register_script_namespace(name: str, namespace:
object) -> object:
return namespace
+def expose_script_op(name: str) -> Any:
+ """Give a registered Op the TVMScript builder its printed calls name, and
return it.
+
+ The printer spells a call to ``tirx.<path>``, ``prim.<path>`` or
``ir.<path>`` as
+ ``T.<path>``, and a call to any other Op by its name, whose first part is
then a
+ parser namespace. The builder is generated from the Op's registration, as
backend
+ builders are, and a missing namespace of its path is created. An existing
callable
+ keeps the name, so a hand-written builder takes precedence. Operators of
Rust
+ ``register_ops!`` blocks get their builders when they are registered.
+
+ Parameters
+ ----------
+ name : str
+ Name of a registered Op.
+
+ Returns
+ -------
+ builder : Callable
+ The builder the printed calls name.
+ """
+ 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
+ if parts[0] in ("tirx", "prim", "ir"):
+ container, parts = sys.modules["tvm_tirx.tirx.script.ir_builder"],
parts[1:]
+
+ def lookup(part):
+ if container is None:
+ return parser._NAMESPACES.get(part)
+ return getattr(container, part, None)
+
+ def bind(part, value):
+ if container is None:
+ parser.register_namespace(part, value)
+ else:
+ setattr(container, part, value)
+
+ *namespaces, leaf = parts
+ for part in namespaces:
+ if lookup(part) is None:
+ bind(part, SimpleNamespace())
+ container = lookup(part)
+ 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
+
+
+# Rust operators call this when they are registered.
+tvm_ffi.register_global_func("tirx.script.expose_script_op", expose_script_op)
+
+
abs = _auto_emit_void_call(_tir_op.abs) # pylint: disable=redefined-builtin
@@ -831,6 +896,7 @@ __all__ = [
"exp",
"exp2",
"exp10",
+ "expose_script_op",
"fabs",
"floor",
"floordiv",
diff --git a/rust/examples/define_ops.rs b/rust/examples/define_ops.rs
index 7bea93fc3..931a8069b 100644
--- a/rust/examples/define_ops.rs
+++ b/rust/examples/define_ops.rs
@@ -22,9 +22,8 @@
//! `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. As a C++
operator
-//! does, each has its TVMScript builder written in Python, which the test
writes:
-//! `T.example_clamp(i, n)`, `T.example_mark()`, `T.example_hint(i,
attrs={..})` and
+//! 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:`.
use tvm_tirx::ir::op_attr::{CALL_EFFECT_KIND, FIXED_RETURN_TYPE};
@@ -51,8 +50,8 @@ tvm_tirx::register_ops! {
.set_infer_type(first_arg_type)
.set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
- // `Call.attrs` holds a `DictAttrs`.
- OpDef::new("tirx.example_hint", "Return the value, with hints for the code
generator.")
+ // `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);
diff --git a/rust/src/ir/op_def.rs b/rust/src/ir/op_def.rs
index 5eb608047..1b96b13f0 100644
--- a/rust/src/ir/op_def.rs
+++ b/rust/src/ir/op_def.rs
@@ -93,9 +93,14 @@ impl OpDef {
Self(registered.and_then(|_| Op::get(name)))
}
- /// The operator, or the error of the step that failed.
+ /// The operator, or the error of the step that failed. Once Python has
imported `tvm_tirx`,
+ /// this also gives the operator its TVMScript builder
(`tirx.script.expose_script_op`).
pub fn op(self) -> Result<Op> {
- self.0
+ let op = self.0?;
+ if let Ok(expose) =
tvm_ffi::Function::get_global("tirx.script.expose_script_op") {
+ expose.call_tuple((op.name.clone(),))?;
+ }
+ Ok(op)
}
/// C++ `OpDef::signature`: replace the argument descriptions and the
attrs type, and the
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 1706935ef..ed1f56047 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -142,9 +142,10 @@
//! 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. As for a C++ operator, the TVMScript builder of an operator
(`T.<name>` for
-//! `tirx.<name>`) is a Python function written for it.
`examples/define_ops.rs` builds a library
-//! of its own the same way, which
`tvm_ffi.load_module(path).tirx_register_ops()` registers.
+//! 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.
//!
//! # Defining a type in Rust
//!
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index a40da764a..b5c43f6d1 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -18,7 +18,7 @@
``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:
-their TVMScript builders are Python functions written for them.
+registering them also gives them their TVMScript builders.
"""
import os
@@ -113,8 +113,6 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
import tvm_tirx
from tvm_tirx import ir
from tvm_tirx.script import tirx as T
- from tvm_tirx.script.ir_builder.base import auto_emit_void_call
- from tvm_tirx.tirx.script import ir_builder
# Registered by the library, with the doc comment and signature the
Rust definition gives.
tvm_ffi.load_module(sys.argv[1]).tirx_register_ops()
@@ -123,7 +121,7 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
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("tirx.example_hint")
+ 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.
@@ -138,23 +136,14 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
else:
raise AssertionError("a float bound passed validation")
- # The TVMScript builders, written in Python as those of the C++
operators are in
- # `tirx/script/ir_builder/op.py`, under the names the printer gives
the calls.
- def example_clamp(index, bound):
- return ir.Call(clamp, [index, bound])
+ # 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
- @auto_emit_void_call
- def example_mark():
- return ir.Call(mark, [])
-
- def example_hint(value, attrs=None):
- return ir.Call(hint, [value], attrs=attrs)
-
- def example_loop(extent):
- return T.region(loop, [extent])
-
- for builder in (example_clamp, example_mark, example_hint,
example_loop):
- setattr(ir_builder, builder.__name__, builder)
+ 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,
@@ -163,14 +152,14 @@ def
test_rust_operators_are_registered_by_their_library(define_ops_library):
def func(i: T.int32, n: T.int32):
T.evaluate(T.example_clamp(i + 1, n))
T.example_mark()
- T.evaluate(T.example_hint(i, attrs={"note": "hot"}))
+ 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 'T.example_hint(i, attrs={"note": "hot"})' 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()
diff --git a/tests/python/tirx/script/test_expose_script_op.py
b/tests/python/tirx/script/test_expose_script_op.py
new file mode 100644
index 000000000..ece863efe
--- /dev/null
+++ b/tests/python/tirx/script/test_expose_script_op.py
@@ -0,0 +1,92 @@
+# 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.
+"""TVMScript builders generated for registered operators."""
+
+import pytest
+import tvm_ffi
+
+import tvm_tirx
+from tvm_tirx.ir import Call, Op, PrimType, 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
+
+
+def define(name, args, result):
+ """Register an operator once per process, as the registry keeps operators
for good."""
+ if name not in Op.list_op_names():
+ tvm_ffi.get_global_func("ir.RegisterOp")(name, f"Test operator
{name}.")
+ Op.get(name).set_signature(args)
+ register_op_attr(name, "TFixedReturnType", PrimType(result))
+ return Op.get(name)
+
+
[email protected]
+def exposed():
+ """The names a test exposes on T, which are removed after it."""
+ names = []
+ yield names
+ for name in names:
+ vars(builder).pop(name, None)
+ parser._NAMESPACES.pop(name, None)
+
+
+def test_builders_round_trip_calls_under_every_root(exposed):
+ exposed += ["test_script_scale", "test_script_mark", "test_script_ns",
"test_script_twice"]
+ exposed += ["test_script_root", "test_script_bare"]
+ scale = define("tirx.test_script_scale", ["x", "factor"], "int32")
+ mark = define("tirx.test_script_mark", ["tag"], "void")
+ touch = define("tirx.test_script_ns.touch", ["x"], "void")
+ twice = define("ir.test_script_twice", ["x"], "int32")
+ other = define("test_script_root.sub.op", ["x"], "void")
+ bare = define("test_script_bare", ["x"], "void")
+ for op in (scale, mark, touch, twice, other, bare):
+ T.expose_script_op(op.name)
+ test_script_root, test_script_bare = parser.test_script_root,
parser.test_script_bare
+ assert T.expose_script_op(scale.name) is T.test_script_scale
+ assert T.test_script_scale.__doc__ == scale.doc
+ assert_structural_equal(T.test_script_scale(x=1, factor=2), Call(scale,
[1, 2]))
+
+ @T.function
+ def func(i: T.int32):
+ T.evaluate(T.test_script_scale(i, 2))
+ T.test_script_mark(i)
+ T.test_script_ns.touch(i)
+ T.evaluate(T.test_script_twice(i))
+ test_script_root.sub.op(i)
+ test_script_bare(i)
+
+ text = func.script()
+ assert "T.test_script_ns.touch(i)" in text, text
+ assert "T.test_script_twice(i)" in text, text
+ assert "test_script_root.sub.op(i)" in text, text
+ assert "\n test_script_bare(i)" in text, text
+ assert_structural_equal(tvm_tirx.script.parse(text, extra_vars={"T": T}),
func)
+
+
+def test_existing_names_keep_their_values(exposed):
+ exposed += ["test_script_written"]
+ written = define("tirx.test_script_written", ["x"], "int32")
+
+ def test_script_written(x):
+ return Call(written, [x])
+
+ builder.test_script_written = test_script_written
+ assert T.expose_script_op(written.name) is test_script_written
+ 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")
