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 a28b1ea6b [DOCS] Lead the extension tutorial with its goal and add a
Python version (#180)
a28b1ea6b is described below
commit a28b1ea6b22cb2c44322920e76b74d270ffded3e
Author: Linzhang Li <[email protected]>
AuthorDate: Tue Nov 3 19:11:10 2026 -0500
[DOCS] Lead the extension tutorial with its goal and add a Python version
(#180)
## Summary
This PR reworks the "Extending TIRx" tutorial around its goal and adds a
Python version of the
complex-number extension, shown beside the Rust one. It also links
source files on GitHub instead
of offering downloads, and points the quick start to further material.
### Getting Started with Rust Extensions
(`docs/advanced/extending_ir_in_rust.rst`)
- **The case study comes first.** It states the goal: to show
"everything is a plugin" with
complex numbers that are written in TVMScript like built-in features and
compile with TIRx.
The page then shows the kernel, the scalar loop that compilation should
produce, and what the
extension needs to get there: a type, operators, and a lowering pass,
plus a note that an
extension can also add optimization passes.
- **"Why Rust" moved** to a short note in the next chapter, where the
Python and Rust versions
meet.
- **Project setup** now says that an extension can also be written in
pure Python.
### Defining Types, Operators, and Passes
(`docs/advanced/defining_ir_extensions.rst`)
- **Rust and Python tabs.** Each step (type, operator, type inference,
rewrites, pass function,
result checks, pass) shows the Rust code and the Python code in group
tabs, as the quick start
does for CPU and CUDA. The text compares the two versions.
- **Python version.** The intro gives the command that runs the example
with it. A note explains
why Rust is preferred for extensions.
- **Notes for the Rust macros.** `def_object_ref!` and
`static_init_block!` are explained in
notes. Another note shows that one `static_init_block!` can register
both types and operators,
and a see-also points `#[dispatch]` and `structural_map` to "Working
with IR in Rust".
- The page follows #178: the excerpt and the text use
`#[dispatch(map, (AnyView<'_>,), Result<Any>)]`.
### Using and Compiling the Extension
(`docs/advanced/using_ir_extensions.rst`)
- The three lines that declare the `my_ext` module and load the library
are explained one by
one.
### Code
- **`examples/extension/my_ext_py/my_ext.py`.** The extension in Python,
with the same names as the
Rust package, so `run_example.py` runs with either:
- the type is a `py_class`, published with `ir.expose_type`;
- the operators are registered with `ir.Op.register`, with Python type
inference;
- the pass is a `structural_map` over `(type, callback)` entries,
wrapped by
`tirx.transform.function_pass`.
A Python operator only has the argument-count validator that its
signature generates, so the
Python pass checks the types of calls with explicit result types before
lowering them.
- **`examples/extension/README.md`.** Mentions the Python version and
how to run it.
### Docs site
- **`docs/conf.py`.** The page header's download menu is replaced by a
button that links the
page's source file on GitHub, next to the fullscreen button. Pages
generated during the build
(`reference/python/generated/`, `reference/cpp/generated/`) have no
source in the repository
and get no button.
- **Links instead of downloads.** The tutorial's `:download:` links to
`src/lib.rs`,
`__init__.py` and `run_example.py` now point to the files on GitHub.
- **`docs/get_started/quick_start.rst`.** Further Reading starts with
[Modern GPU Programming for
MLSys](https://mlc.ai/modern-gpu-programming-for-mlsys/), recommended
for learning to write programs with TIRx, and lists the extension
tutorial.
- **`docs/ir_guides/working_with_ir_in_rust.rst`.** Adds the labels
`rust-dispatch` and
`rust-structural-utilities` for the tutorial's see-also.
### Tests
- **`tests/python/ir/test_op_rust.py`.**
- The checks of explicitly typed invalid calls move to
`VALIDATOR_CHECKS`, which only the Rust
package runs.
- The new `test_extension_prototype` runs `run_example.py` and the
shared checks with the
Python version, plus `PROTOTYPE_CHECKS`: its argument-count validator,
and the pass rejecting
explicitly typed invalid calls.
Verified locally on Linux:
- `tests/python/ir/test_op_rust.py`: all 7 tests pass on the base before
#178. On the rebased
branch, `test_extension_prototype` passes. The Rust tests were not rerun
after the rebase,
because the local build predates the TVM FFI update in #178; this PR
does not change the Rust
code.
- `PYTHONPATH=examples/extension/my_ext_py python
examples/extension/run_example.py` prints
`Complex multiplication matches NumPy.`
- A merged type-and-operator `static_init_block!`, as the new note
shows, builds and runs the
example.
- The Sphinx HTML build with `-W`, using CI's Sphinx 8.1.3 and docutils
0.21.2.
- pre-commit on the changed files.
---
docs/advanced/defining_ir_extensions.rst | 394 +++++++++++++++++++++--------
docs/advanced/extending_ir_in_rust.rst | 141 ++++++-----
docs/advanced/using_ir_extensions.rst | 35 ++-
docs/conf.py | 33 +++
docs/get_started/quick_start.rst | 6 +
docs/ir_guides/working_with_ir_in_rust.rst | 4 +
examples/extension/README.md | 7 +
examples/extension/my_ext_py/my_ext.py | 226 +++++++++++++++++
tests/python/ir/test_op_rust.py | 44 +++-
9 files changed, 700 insertions(+), 190 deletions(-)
diff --git a/docs/advanced/defining_ir_extensions.rst
b/docs/advanced/defining_ir_extensions.rst
index 7ec42e5f6..e58aa53fd 100644
--- a/docs/advanced/defining_ir_extensions.rst
+++ b/docs/advanced/defining_ir_extensions.rst
@@ -18,11 +18,33 @@
Defining Types, Operators, and Passes
=====================================
-In this chapter, we walk through the Rust library of the package set up in
-:doc:`extending_ir_in_rust`: a complex type, four operators, and a pass that
-lowers them to scalar arithmetic. The excerpts come from
-:download:`src/lib.rs <../../examples/extension/src/lib.rs>`, which also
contains
-the validation and traversal helpers they use.
+In this chapter, we define the extension set up in :doc:`extending_ir_in_rust`:
+a complex type, four operators, and a pass that lowers them to scalar
+arithmetic. Each step is shown in Rust, from the library
+`src/lib.rs
<https://github.com/apache/tvm-tirx/blob/main/examples/extension/src/lib.rs>`_,
+and in Python, from the module
+`my_ext_py/my_ext.py
<https://github.com/apache/tvm-tirx/blob/main/examples/extension/my_ext_py/my_ext.py>`_.
+Both files also contain the validation and traversal helpers that the excerpts
+use.
+
+The two versions define the same names, so the example program runs with
+either. The Python version is a single module that needs no build, which makes
+it a quick way to try out an extension before writing it in Rust. To run the
+example with it, put its directory on ``PYTHONPATH``:
+
+.. code-block:: bash
+
+ PYTHONPATH=examples/extension/my_ext_py python
examples/extension/run_example.py
+
+.. note::
+
+ **Why is Rust preferred for extensions?**
+
+ Python is the quickest way to try out an extension, but it is slow on large
+ programs.
+
+ Rust is fast and memory safe, which makes it a good fit for agents exploring
+ new types, operators, and passes.
Defining a Type
---------------
@@ -31,36 +53,80 @@ To work with complex numbers directly in the IR, we first
introduce a type for
them. ``ComplexType`` has one field, ``dtype``, the scalar type shared by the
real and imaginary parts, as in ``TensorType``:
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin type
- :end-before: // end type
-
-As in C++, the type is defined in two parts: ``ComplexTypeObj`` is the object
-that holds the fields, and ``ComplexType`` is a reference-counted reference to
-it, which IR nodes store and functions pass around. ``base: TypeObj`` makes the
-object an IR type, and ``rs_class`` registers its fields for reflection,
through
-which Python and IR JSON loading construct the type as they construct a C++
one.
-
-``def_object_ref!(pub ComplexType, Type, ComplexTypeObj)`` defines the
-reference. Its arguments are:
-
-- ``pub ComplexType``: the name and visibility of the reference.
-- ``Type``: the reference to the parent object, ``TypeObj``. A ``ComplexType``
- converts into a ``Type`` with ``into()`` wherever a ``Type`` is expected,
such
- as the result of type inference.
-- ``ComplexTypeObj``: the object it refers to. A ``ComplexTypeObj`` converts
- into a ``ComplexType`` with ``into()``, and a ``ComplexType`` dereferences to
- it, so ``ty.dtype`` reads the field.
-
-The ``static_init_block!`` runs ``register_type`` when the library loads, as a
-``TVM_FFI_STATIC_INIT_BLOCK`` does for C++, and Python then sees the type as
-``my_ext.Complex``.
-
-A ``ComplexType`` records the scalar type of the two parts, and a
-``complex(re, im)`` call supplies their values as its operands. ``dtype`` can
be
-any scalar float type, such as ``float32``, ``float64`` or ``bfloat16``; this
-tutorial uses ``float32``.
+.. tabs::
+
+ .. group-tab:: Rust
+
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin type
+ :end-before: // end type
+
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :start-after: # begin type
+ :end-before: # end type
+
+.. note::
+
+ **What does def_object_ref! do?**
+
+ ``def_object_ref!(ObjRef, Parent, Obj)`` defines ``ObjRef``, the reference
+ type of the object type ``Obj``, as the C++ macro
+ ``TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE`` does. An ``ObjRef`` holds
+ a reference-counted pointer to an ``Obj``, and the macro gives it these
+ conversions:
+
+ - an ``Obj`` converts into an ``ObjRef`` with ``into()``;
+ - an ``ObjRef`` dereferences to its ``Obj``, so its fields read directly;
+ - an ``ObjRef`` converts into ``Parent``, the reference to the parent
+ object, with ``into()``.
+
+ A visibility such as ``pub`` before ``ObjRef`` applies to the reference.
+ Here, ``def_object_ref!(pub ComplexType, Type, ComplexTypeObj)`` defines the
+ public reference ``ComplexType``: ``ty.dtype`` reads the field of a
+ ``ComplexType``, and a ``ComplexType`` converts into a ``Type`` wherever one
+ is expected, such as the result of type inference.
+
+In Rust, as in C++, the type is defined in two parts: ``ComplexTypeObj`` is the
+object that holds the fields, and ``ComplexType`` is a reference-counted
+reference to it, which IR nodes store and functions pass around.
+``base: TypeObj`` makes the object an IR type, and ``rs_class`` registers its
+fields for reflection, through which Python and IR JSON loading construct the
+type as they construct a C++ one. ``def_object_ref!`` defines the reference.
+
+The ``static_init_block!`` registers the type when the library loads.
+
+.. note::
+
+ **What does static_init_block! do?**
+
+ ``static_init_block! { ... }`` runs its body when the program or library
that
+ holds it loads, as a C++ ``TVM_FFI_STATIC_INIT_BLOCK`` does. The body
+ registers what the library defines, such as types, operators, and global
+ functions, and Python sees these registrations once it loads the library.
+ The body is the body of a function that returns ``Result<()>``, so ``?``
+ stops it at its first error.
+
+ Here, ``register_type::<ComplexType>()`` registers ``ComplexType`` as an IR
+ type, so that Python constructs it from its fields and IR JSON holds it, as
+ for a C++ type.
+
+In Python, one class is enough. ``py_class`` from ``tvm_ffi.dataclasses``
+registers it as an IR object: ``my_ext.ComplexType`` is its type key,
+``ir.Type`` its parent, and the annotated ``dtype`` its field. The class gets a
+constructor, ``ComplexType(dtype)``, and its field takes part in structural
+equality, printing, and IR JSON, as the fields of a C++ type do. The global
+function ``ir.expose_type`` then publishes the type, as ``register_type`` does
+in Rust.
+
+Either way, Python sees the type as ``my_ext.Complex``. A ``ComplexType``
+records the scalar type of the two parts, and a ``complex(re, im)`` call
+supplies their values as its operands. ``dtype`` can be any scalar float type,
+such as ``float32``, ``float64`` or ``bfloat16``; this tutorial uses
+``float32``.
Defining Operators
------------------
@@ -91,38 +157,85 @@ they are registered under the ``my_ext`` namespace:
- Extract the imaginary part
- ``T.float32``
-The ``complex`` operator is registered as an ``OpDef`` chain in a
-``static_init_block!``. Each step updates the operator registry, and the chain
-completes when its statement ends:
+Each operator has a signature, a type inference function, and the effect kind
+``Pure``, which marks it as a pure function that does not interact with any
+external state. The ``complex`` operator is registered as follows:
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin constructor
- :end-before: // end constructor
+.. tabs::
-``signature`` declares the arguments of the operator and of its Python and
-TVMScript builder. ``set_infer_type`` attaches a named Rust function that
-checks the arguments and computes the result type:
+ .. group-tab:: Rust
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin inference
- :end-before: // end inference
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin constructor
+ :end-before: // end constructor
-``set_validator`` repeats these checks for a call whose result type is given
-explicitly rather than inferred. ``Pure`` marks the operator as a pure
-function, which does not interact with any external state.
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :start-after: # begin constructor
+ :end-before: # end constructor
+
+In Rust, the operator is an ``OpDef`` chain in a ``static_init_block!``: each
+step updates the operator registry, and the chain completes when its statement
+ends. In Python, ``ir.Op.register`` adds the operator to the registry and
+returns it, and the helper ``_define`` sets its signature and attributes;
+``expose`` publishes its builder, ``my_ext.complex``, once the definition is
+complete. The module keeps the returned operator as the global ``COMPLEX``.
+
+.. note::
+
+ A ``static_init_block!`` can hold any registrations, types and operators
+ alike. The example keeps the type and the operators in separate blocks, but
+ one block can register both:
+
+ .. code-block:: rust
+
+ tvm_ffi::static_init_block! {
+ register_type::<ComplexType>()?;
+ OpDef::new("my_ext.complex", "Construct the complex number re + im *
i.")
+ .signature(...)
+ .set_infer_type(infer_complex);
+ }
-``infer_complex`` computes the type of a ``complex(re, im)`` call: a
+In both, the signature declares the arguments of the operator and of its Python
+and TVMScript builder, and the type inference function checks the arguments
+and computes the result type:
+
+.. tabs::
+
+ .. group-tab:: Rust
+
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin inference
+ :end-before: // end inference
+
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :pyobject: _infer_complex
+
+The function computes the type of a ``complex(re, im)`` call: a
``ComplexType`` whose ``dtype`` is the type of the two parts. It requires both
-parts to be scalar floats of the same type, and converts neither. It then uses
-both conversions of the reference: it builds a ``ComplexTypeObj``, turns it
into
-a ``ComplexType``, and returns it as a ``Type``. The other three operators
-follow the same pattern. ``cmul`` requires two complex arguments of the same
-type and returns it; ``real`` and ``imag`` require one and return its
-``dtype``. Passing an ordinary scalar to ``cmul`` therefore fails with
+parts to be scalar floats of the same type, and converts neither. The Rust
+version uses both conversions of the reference: it builds a
+``ComplexTypeObj``, turns it into a ``ComplexType``, and returns it as a
+``Type``; the Python version calls the class. The other three operators follow
+the same pattern. ``cmul`` requires two complex arguments of the same type and
+returns it; ``real`` and ``imag`` require one and return its ``dtype``. Passing
+an ordinary scalar to ``cmul`` therefore fails with
``TypeError: my_ext.cmul expects a complex value``.
+An operator's validator checks a call whose result type is given explicitly
+rather than inferred. In Rust, ``set_validator`` attaches one that repeats the
+checks of type inference. A Python operator has the validator that its
+signature generates, which checks the number of arguments; validators that also
+check types are native functions. The Python version therefore checks the types
+of such calls in its pass, before lowering them.
+
Writing a Pass
--------------
@@ -154,74 +267,137 @@ 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 looks up the operators by name, as C++
-does with ``Op::Get``. The operators are registered when the library loads, and
-``Op::get`` returns the registry's single object for a name.
+Only scalar arithmetic remains. To apply the rules, the pass tests whether a
+call's ``op`` is one of the four operators. The Rust pass looks them up by
name,
+as C++ does with ``Op::Get``. The operators are registered when the library
+loads, and ``Op::get`` returns the registry's single object for a name.
``call.op.same_as(&Op::get("my_ext.cmul")?)`` therefore tests whether a call's
-``op`` is ``my_ext.cmul`` by comparing pointers. ``Lowering`` implements the
-rewrites:
-
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin lowering
- :end-before: // end lowering
- :prepend: struct Lowering {
- bindings: Vec<(Var, Var, Var)>,
- }
- impl Lowering {
- fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> { ... }
- }
- #[dispatch(map, (AnyView<'_>,), Result<Any>)]
- impl Lowering {
- fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
- :append: }
+``op`` is ``my_ext.cmul`` by comparing pointers. The Python pass compares with
+the operators that ``Op.register`` returned, as in ``call.op.same_as(CMUL)``.
+The rewrites are:
+
+.. tabs::
+
+ .. group-tab:: Rust
+
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin lowering
+ :end-before: // end lowering
+ :prepend: struct Lowering {
+ bindings: Vec<(Var, Var, Var)>,
+ }
+ impl Lowering {
+ fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> { ... }
+ }
+ #[dispatch(map, (AnyView<'_>,), Result<Any>)]
+ impl Lowering {
+ fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
+ :append: }
+
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :pyobject: _Lowering
``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, (AnyView<'_>,), Result<Any>)]`` 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 callback passes each node to ``lowering.dispatch`` and
-returns ``Unchanged`` when no handler matches; these rewrites do not need the
-definition region kind.
+``complex(re, im)`` call or a variable already split. ``map_call`` first
+validates a call to one of the operators, then applies the three rules;
+``map_bind`` splits a complex binding, and ``map_sequence`` flattens the
+sequence that a split binding leaves inside another. In Rust,
+``#[dispatch(map, (AnyView<'_>,), Result<Any>)]`` sends each node to the
+``map_`` method of its type. In Python, the next excerpt pairs each method with
+the type of node it takes.
+
+.. seealso::
+
+ :ref:`rust-dispatch` and :ref:`rust-structural-utilities` in
+ :doc:`../ir_guides/working_with_ir_in_rust`.
The pass is built from an ordinary named function:
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin function callback
- :end-before: // end function callback
+.. tabs::
+
+ .. group-tab:: Rust
+
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin function callback
+ :end-before: // end function callback
+
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :pyobject: _lower_function
It receives a function, its module, and the pass context, and returns the
-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:
+rewritten function; in Rust, ``Some`` replaces the original. The first
+``structural_map`` walks the function with the rewrites. It visits children
+before their parent, ``WalkOrder::PostOrder`` in Rust and the default order in
+Python, so the operands of ``cmul`` are lowered before ``cmul`` itself. In
+Python, ``structural_map`` takes a list of ``(type, callback)`` entries and
+calls, on each node, the first callback whose type matches it.
-.. code-block:: rust
+The second ``structural_map`` only checks the result. It takes three functions,
+a tuple in Rust and a list of entries in Python, and calls each on the nodes of
+its type:
- /// Reject a complex type, wherever it appears: on a variable, a parameter
or an expression.
- fn check_type(ty: Type) -> Result<Type> { ... }
+.. tabs::
- /// Reject an expression whose type is complex.
- fn check_expr(expr: Expr) -> Result<Expr> { ... }
+ .. group-tab:: Rust
- /// Reject a call to one of the four complex operators that the rewrite did
not consume.
- fn check_call_lowered(call: Call) -> Result<Call> { ... }
+ .. code-block:: rust
+
+ /// Reject a complex type, wherever it appears: on a variable, a
parameter or an expression.
+ fn check_type(ty: Type) -> Result<Type> { ... }
+
+ /// Reject an expression whose type is complex.
+ 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> { ... }
+
+ .. group-tab:: Python
+
+ .. code-block:: python
+
+ def _check_type(ty):
+ """Reject a complex type, wherever it appears: on a variable, a
parameter or an expression."""
+
+ def _check_expr(expr):
+ """Reject an expression whose type is complex."""
+
+ def _check_call_lowered(call):
+ """Reject a call to one of the four complex operators that the
rewrite did not consume."""
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
before code generation.
-Finally, ``FunctionPass::new`` wraps the function as a pass, and
-``tvm_ffi_dll_export_typed_func!`` exports ``lower_complex``, which builds it,
-so that the Python package can call it from the loaded library:
+Finally, the function becomes a pass:
+
+.. tabs::
+
+ .. group-tab:: Rust
+
+ .. literalinclude:: ../../examples/extension/src/lib.rs
+ :language: rust
+ :start-after: // begin pass
+ :end-before: // end pass
-.. literalinclude:: ../../examples/extension/src/lib.rs
- :language: rust
- :start-after: // begin pass
- :end-before: // end pass
+ .. group-tab:: Python
+
+ .. literalinclude:: ../../examples/extension/my_ext_py/my_ext.py
+ :language: python
+ :pyobject: lower_complex
+
+In Rust, ``FunctionPass::new`` wraps the function as a pass, and
+``tvm_ffi_dll_export_typed_func!`` exports ``lower_complex``, which builds it,
+so that the Python package can call it from the loaded library. In Python,
+``tirx.transform.function_pass`` wraps the function, and the module defines
+``lower_complex`` directly.
Next, :doc:`using_ir_extensions` loads the library and compiles the function.
diff --git a/docs/advanced/extending_ir_in_rust.rst
b/docs/advanced/extending_ir_in_rust.rst
index 2a68bca4b..0de362e4b 100644
--- a/docs/advanced/extending_ir_in_rust.rst
+++ b/docs/advanced/extending_ir_in_rust.rst
@@ -25,45 +25,86 @@ Getting Started with Rust Extensions
TIRx is designed around the idea that "everything is a plugin": compiler
features can be provided by separately built libraries, using the same IR and
-registration interfaces as built-in features. This chapter puts that idea into
practice
-with a complex-number extension that loads into TIRx without modifying its
-core library.
-
-This tutorial shows how to extend TIRx from Rust: define new IR types and
-operators, write a pass that lowers them, and use them from Python like
-built-in features.
-
-On this page, we will first look at why extensions are written in Rust, then
-get ready for the rest of the tutorial: a Python package for the extension and
-the example program it will compile.
-
-Why Rust
---------
-
-TIRx is built on a stable C++ library for its IR, analysis, transformations,
-and code generation. In most cases, you work with it from Python: tensor
-programs are written as Python functions in TVMScript, then compiled and run
-directly.
-
-With agents, compiler development iterates quickly. Along the way, a project
may
-need new types and operators to describe its own concepts, such as a new number
-format or a hardware feature, and passes that lower them to what the backend
-supports. To support this, TIRx provides the ``tvm-tirx`` crate, with which you
-define new types, operators, and passes in Rust on the same IR objects as C++
-and Python.
-
-With the ``tvm-tirx`` crate, users can extend TIRx in separate Rust libraries.
-Each extension builds and loads on its own, without changing the C++ library,
so
-the C++ core stays stable while the extension iterates quickly. Once loaded,
its
-types and operators can be written in TVMScript and its passes can run in
Python
-compilation pipelines, just like built-in ones.
+registration interfaces as built-in features. This tutorial shows how to write
+such a library in Rust.
+
+Case Study: Complex Numbers
+---------------------------
+
+To show "everything is a plugin" in practice, we use complex numbers as an
+example. All complex-number support comes from an extension that loads into
+TIRx without modifying its core library.
+
+Our goal is that, once the extension is loaded, complex types and operators are
+written in TVMScript just like built-in ones, and functions that use them
+compile with TIRx. For example, the following function multiplies two complex
+tensors element by element. ``A`` and ``B`` hold the real and imaginary parts
+of the first tensor, ``C`` and ``D`` those of the second. We will compile and
+run it at the end of the tutorial.
+
+.. literalinclude:: ../../examples/extension/run_example.py
+ :language: python
+ :start-after: # begin kernel
+ :end-before: # end kernel
+
+``my_ext`` is the extension's Python package, which holds the builders of its
+type and operators; :doc:`using_ir_extensions` shows how it provides them.
+``z``, ``w``, and ``product`` are complex values that exist only in the IR.
+The function's inputs and outputs are ordinary ``float32`` tensors, so no
+runtime storage format for complex numbers is needed.
+
+Code generation has no complex type, so by the time the function reaches the
+backend, every complex value must be a pair of ``float32`` values. Expanding
+complex multiplication as
+
+.. math::
+
+ (a + bi)(c + di) = (ac - bd) + (ad + bc)i
+
+gives the loop we want the compiler to produce:
+
+.. code-block:: python
+
+ for i in range(N):
+ z_re: T.let[T.float32] = A[i]
+ z_im: T.let[T.float32] = B[i]
+ w_re: T.let[T.float32] = C[i]
+ w_im: T.let[T.float32] = D[i]
+ product_re: T.let[T.float32] = z_re * w_re - z_im * w_im
+ product_im: T.let[T.float32] = z_re * w_im + z_im * w_re
+ R[i] = product_re
+ Im[i] = product_im
+
+``z``, ``w``, and ``product`` have each become a pair of variables for the real
+and imaginary parts, and only scalar arithmetic remains.
+
+To support complex numbers as a custom type, we need:
+
+1. **A type.** For TIRx to recognize complex values, a complex type, the type
+ of ``z``, ``w``, and ``product``.
+2. **Operators.** For TIRx to compute on complex values, the operators
+ ``complex``, ``cmul``, ``real``, and ``imag``, which build complex values,
+ multiply them, and read their parts.
+3. **A lowering pass.** For the target to compile the function,
``LowerComplex``,
+ which lowers the complex type and operators to the scalar arithmetic above.
+
+.. note::
+
+ An extension can also add passes that optimize computations on its types,
+ for example to simplify a product with a known real factor. They are written
+ the same way as the lowering pass, so the example shows only the latter.
+
+The rest of this page sets up the extension's Python package.
+:doc:`defining_ir_extensions` then defines the type, operators, and pass, and
+:doc:`using_ir_extensions` compiles and runs the function.
Setting Up the Project
----------------------
-An extension is shipped as a Python package that contains its Rust library, so
-it is installed and imported like any other package. The example package,
-``my-ext``, has this layout:
+An extension can be shipped as a Python package that contains its Rust library
+(or written in pure Python, like the version shown beside Rust in
+:doc:`defining_ir_extensions`), so it can be installed and imported like any
+other package. The example package, ``my-ext``, has this layout:
.. code-block:: text
@@ -97,32 +138,4 @@ The package depends on ``apache-tvm-tirx`` at run time;
disabling build
isolation keeps its native build aligned with these sources instead of
installing a released package in a temporary build environment.
-Case Study: Complex Numbers
----------------------------
-
-To show how to define new types, operators, and passes, we will use complex
-numbers as a running example throughout this tutorial: a complex type,
operators
-that build and multiply complex values, and a pass that expands complex
-multiplication into real arithmetic:
-
-.. math::
-
- (a + bi)(c + di) = (ac - bd) + (ad + bc)i.
-
-The following function multiplies two tensors of complex numbers element by
-element. ``A`` and ``B`` hold the real and imaginary parts of the first
operand,
-``C`` and ``D`` those of the second:
-
-.. literalinclude:: ../../examples/extension/run_example.py
- :language: python
- :start-after: # begin kernel
- :end-before: # end kernel
-
-``my_ext`` is the extension's Python package, which holds the builders of its
-type and operators; :doc:`using_ir_extensions` shows how it provides them.
-``z``, ``w``, and ``product`` are complex values that exist only in the IR.
-The function's inputs and outputs are ordinary ``float32`` tensors, so no
-runtime storage format for complex numbers is needed.
-
-Next, :doc:`defining_ir_extensions` defines the type, operators, and pass this
-function uses, and :doc:`using_ir_extensions` compiles and runs it.
+Next, :doc:`defining_ir_extensions` walks through the extension's Rust library.
diff --git a/docs/advanced/using_ir_extensions.rst
b/docs/advanced/using_ir_extensions.rst
index 04e72f180..345768a5a 100644
--- a/docs/advanced/using_ir_extensions.rst
+++ b/docs/advanced/using_ir_extensions.rst
@@ -30,7 +30,7 @@ static init blocks, which run as soon as the library loads.
To write them in
TVMScript, Python also needs a module that holds their builders. Definitions
under ``tirx.`` go to ``T``, but an extension namespace such as ``my_ext`` is
published only into a module that the extension declares. The package's
-:download:`__init__.py <../../examples/extension/python/my_ext/__init__.py>`
+`__init__.py
<https://github.com/apache/tvm-tirx/blob/main/examples/extension/python/my_ext/__init__.py>`_
therefore declares itself as that module before it loads the library:
.. literalinclude:: ../../examples/extension/python/my_ext/__init__.py
@@ -38,12 +38,27 @@ therefore declares itself as that module before it loads
the library:
:start-after: # begin load
:end-before: # end load
-The two ``_init_exposed_*_api`` calls, internal to ``tvm_tirx.ir`` for now,
-publish the operators and the type into ``my_ext``; ``auto_emit_void_call``
-gives its builders the TVMScript behavior of the ``T`` builders.
-``load_lib_module`` finds the library that the package installs next to
-``__init__.py``. To parse printed IR that uses the extension, pass the package
-as well, as in ``extra_vars={"T": T, "my_ext": my_ext}``.
+The three lines do the following:
+
+- ``_init_exposed_op_api("my_ext", __name__, wrapper=auto_emit_void_call)``
+ declares the module ``__name__`` as the destination of the operators whose
+ names start with ``my_ext.``, and publishes each of them there without the
+ prefix, as ``my_ext.cmul``. ``auto_emit_void_call`` gives the builders the
+ TVMScript behavior of the ``T`` builders.
+- ``_init_exposed_type_api("my_ext", __name__)`` does the same for the types
+ whose keys start with ``my_ext.``, so ``my_ext.ComplexType`` is published as
+ ``my_ext.Complex``.
+- ``load_lib_module("my-ext", "my_ext", extra_lib_paths=[...])`` finds the
+ library ``my_ext`` (``libmy_ext.so`` on Linux) of the package ``my-ext``,
+ here next to ``__init__.py``, and loads it, which runs its static init
+ blocks. It returns the library as a module, whose exported functions, such as
+ ``lower_complex``, Python calls.
+
+Both declarations come before the library loads, so its registrations reach
+``my_ext`` as they happen. The two ``_init_exposed_*_api`` functions are
+internal to ``tvm_tirx.ir`` for now. To parse printed IR that uses the
+extension, pass the package as well, as in
+``extra_vars={"T": T, "my_ext": my_ext}``.
The registered type and operators let the IR express complex values;
``LowerComplex`` turns them into a form the backend can compile. The default
@@ -60,9 +75,9 @@ also named ``my_ext``, that runs it first:
Compile the Function
--------------------
-:download:`run_example.py <../../examples/extension/run_example.py>` imports
-``my_ext``, which loads the library and registers the pipeline, and compiles
-the function from :doc:`extending_ir_in_rust` with it:
+`run_example.py
<https://github.com/apache/tvm-tirx/blob/main/examples/extension/run_example.py>`_
+imports ``my_ext``, which loads the library and registers the pipeline, and
+compiles the function from :doc:`extending_ir_in_rust` with it:
.. literalinclude:: ../../examples/extension/run_example.py
:language: python
diff --git a/docs/conf.py b/docs/conf.py
index dc21ae85c..3c3392a82 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -148,6 +148,8 @@ html_theme_options = {
"repository_branch": "main",
"path_to_docs": "docs",
"use_repository_button": True,
+ # Replaced by a link to the page's source file, see
_add_source_link_button.
+ "use_download_button": False,
"show_toc_level": 2,
"max_navbar_depth": 4,
"collapse_navbar": True,
@@ -253,9 +255,40 @@ def _link_nested_class_to_enclosing_class(app, docname,
source):
)
+# Pages generated during the build, which have no source file in the
repository.
+_GENERATED_DOC_DIRS = ("reference/python/generated/",
"reference/cpp/generated/")
+
+
+def _add_source_link_button(app, pagename, templatename, context, doctree):
+ """Add a header button, next to the fullscreen one, that links the page's
source on GitHub."""
+ suffix = context.get("page_source_suffix")
+ if not suffix or pagename.startswith(_GENERATED_DOC_DIRS):
+ return
+ opts = html_theme_options
+ url = "/".join(
+ [opts["repository_url"], "blob", opts["repository_branch"],
opts["path_to_docs"]]
+ )
+ buttons = context["header_buttons"]
+ labels = [button.get("label") for button in buttons]
+ index = labels.index("fullscreen-button") if "fullscreen-button" in labels
else len(buttons)
+ buttons.insert(
+ index,
+ {
+ "type": "link",
+ "url": f"{url}/{pagename}{suffix}",
+ "tooltip": "Source file on GitHub",
+ "icon": "fas fa-file-code",
+ "label": "source-link-button",
+ "classes": "pst-navbar-icon",
+ },
+ )
+
+
def setup(app):
_prepare_stub_files()
_build_rust_docs()
app.connect("source-read", _mark_exhale_root_orphan)
app.connect("source-read", _link_nested_class_to_enclosing_class)
+ # After the theme's own header buttons, which it adds at priority 501.
+ app.connect("html-page-context", _add_source_link_button, priority=502)
app.connect("build-finished", _copy_rust_docs_to_output)
diff --git a/docs/get_started/quick_start.rst b/docs/get_started/quick_start.rst
index ffb5cfb3e..4741cce81 100644
--- a/docs/get_started/quick_start.rst
+++ b/docs/get_started/quick_start.rst
@@ -300,9 +300,15 @@ The same program can also be built with CMake. The
``CMakeLists.txt`` in
Further Reading
---------------
+- `Modern GPU Programming for MLSys
<https://mlc.ai/modern-gpu-programming-for-mlsys/>`_:
+ a book on writing fast GPU kernels with TIRx, from the GPU execution model
+ to GEMM and FlashAttention 4 on NVIDIA Blackwell (**recommended for learning
+ how to write programs with TIRx**).
- :doc:`/ir_guides/overview`: the IR objects in depth and how they fit
together.
- :doc:`Programming Guides </programming_guides/index>`: writing CUDA
kernels with TIRx.
+- :doc:`/advanced/extending_ir_in_rust`: extending the compiler with types,
+ operators, and passes of your own.
- :doc:`Python API </reference/python/index>`: the full Python API reference.
- :doc:`/dev/source_build`: building TIRx from source.
- `TVM-FFI documentation <https://tvm.apache.org/ffi/>`_: the FFI and ABI
diff --git a/docs/ir_guides/working_with_ir_in_rust.rst
b/docs/ir_guides/working_with_ir_in_rust.rst
index 069eed066..94b6a8786 100644
--- a/docs/ir_guides/working_with_ir_in_rust.rst
+++ b/docs/ir_guides/working_with_ir_in_rust.rst
@@ -304,6 +304,8 @@ types it knows, and ``_`` lets programs with new types pass
through.
A match looks at one node at a time; to process every node in a program
while keeping state, we dispatch.
+.. _rust-dispatch:
+
Dispatching on Node Types
~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -348,6 +350,8 @@ walk continues into their children, so new types pass
through. Here the
multiplication and the constants have no method, and the walk still
reaches ``x`` inside the multiplication.
+.. _rust-structural-utilities:
+
Structural Utilities
--------------------
diff --git a/examples/extension/README.md b/examples/extension/README.md
index ff1b32141..8c26307bd 100644
--- a/examples/extension/README.md
+++ b/examples/extension/README.md
@@ -27,6 +27,7 @@ pipeline. The tutorial in `docs/advanced/` walks through the
code.
- `CMakeLists.txt`, `pyproject.toml`: build the library and install it into
the package
- `python/my_ext/`: loads the library and registers the pipeline
- `run_example.py`: compiles a function that uses the extension and checks it
against NumPy
+- `my_ext_py/my_ext.py`: the same extension written in Python, a prototype
that needs no build
## Install and Run
@@ -43,3 +44,9 @@ The build uses this checkout's `tvm-tirx` crate and pinned
TVM FFI sources.
Disabling build isolation keeps the native libraries from the same checkout
available to cargo, instead of installing a released package in a temporary
build environment.
+
+To run the example with the Python prototype instead, without building
anything:
+
+```bash
+PYTHONPATH=examples/extension/my_ext_py python
examples/extension/run_example.py
+```
diff --git a/examples/extension/my_ext_py/my_ext.py
b/examples/extension/my_ext_py/my_ext.py
new file mode 100644
index 000000000..8574d9ea2
--- /dev/null
+++ b/examples/extension/my_ext_py/my_ext.py
@@ -0,0 +1,226 @@
+# 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.
+"""A Python prototype of my_ext: the complex type, its operators, and the pass
that lowers them.
+
+It defines the same names as the Rust package in ``python/my_ext``, so
``run_example.py`` runs
+with either; a process loads one of them.
+"""
+
+import tvm_ffi
+from tvm_ffi.dataclasses import py_class
+
+from tvm_tirx import ir, tirx
+from tvm_tirx.ir._expose_api import _init_exposed_op_api,
_init_exposed_type_api
+from tvm_tirx.script.ir_builder.base import auto_emit_void_call
+
+# begin load
+_init_exposed_op_api("my_ext", __name__, wrapper=auto_emit_void_call)
+_init_exposed_type_api("my_ext", __name__)
+# end load
+
+
+# begin type
+@py_class("my_ext.ComplexType")
+class ComplexType(ir.Type):
+ """A complex number whose real and imaginary parts share a scalar type."""
+
+ dtype: ir.PrimType
+
+
+tvm_ffi.get_global_func("ir.expose_type")("my_ext.ComplexType")
+# end type
+
+
+def _require(condition, message):
+ if not condition:
+ raise TypeError(message)
+
+
+def _check_call(call, count):
+ _require(len(call.args) == count, f"expected {count} value arguments")
+ _require(len(call.ty_args) == 0, "type arguments are not supported")
+ _require(call.attrs is None, "call attributes are not supported")
+
+
+def _float_part(ty):
+ """The type of a part of a complex number: a scalar float, such as float32
or bfloat16."""
+ _require(isinstance(ty, ir.PrimType), "expected a scalar float")
+ dtype = ty.dtype
+ _require(dtype in (f"float{dtype.bits}", f"bfloat{dtype.bits}"), "expected
a scalar float")
+ return ty
+
+
+def _infer_complex(call):
+ _check_call(call, 2)
+ elem = _float_part(call.args[0].ty)
+ im = _float_part(call.args[1].ty)
+ _require(elem.dtype == im.dtype, "my_ext.complex expects parts of the same
type")
+ return ComplexType(elem)
+
+
+def _complex_arg(call, index):
+ ty = call.args[index].ty
+ _require(isinstance(ty, ComplexType), f"{call.op.name} expects a complex
value")
+ _float_part(ty.dtype)
+ return ty
+
+
+def _infer_cmul(call):
+ _check_call(call, 2)
+ a, b = _complex_arg(call, 0), _complex_arg(call, 1)
+ _require(a.dtype.dtype == b.dtype.dtype, "my_ext.cmul expects complex
values of the same type")
+ return a
+
+
+def _infer_part(call):
+ _check_call(call, 1)
+ return _complex_arg(call, 0).dtype
+
+
+# begin constructor
+def _define(name, doc, args, infer_type):
+ op = ir.Op.register(f"my_ext.{name}", doc)
+ op.set_signature(args=args)
+ op.set_attr("FInferType", infer_type)
+ op.set_attr("TCallEffectKind", ir.CallEffectKind.Pure.value)
+ op.expose()
+ return op
+
+
+COMPLEX = _define(
+ "complex",
+ "Construct the complex number re + im * i.",
+ [("re", "The real part."), ("im", "The imaginary part.")],
+ _infer_complex,
+)
+# end constructor
+# begin operators
+CMUL = _define(
+ "cmul",
+ "Multiply two complex numbers.",
+ [("a", "The first complex value."), ("b", "The second complex value.")],
+ _infer_cmul,
+)
+REAL = _define(
+ "real", "Extract the real part of a complex number.", [("z", "The complex
value.")], _infer_part
+)
+IMAG = _define(
+ "imag",
+ "Extract the imaginary part of a complex number.",
+ [("z", "The complex value.")],
+ _infer_part,
+)
+_OPS = (COMPLEX, CMUL, REAL, IMAG)
+# end operators
+_INFER = {"my_ext.complex": _infer_complex, "my_ext.cmul": _infer_cmul}
+
+
+def _validate_call(call):
+ """Infer the type of a call to one of the operators again, and compare it
with its type."""
+ expected = _INFER.get(call.op.name, _infer_part)(call)
+ _require(
+ tvm_ffi.structural_equal(expected, call.ty), "incorrect result type
for complex operator"
+ )
+
+
+class _Lowering:
+ def __init__(self):
+ # Each complex local variable is replaced by two scalar local
variables.
+ self.bindings = {}
+
+ def parts(self, expr):
+ if isinstance(expr, ir.Call) and expr.op.same_as(COMPLEX):
+ return expr.args[0], expr.args[1]
+ if expr in self.bindings:
+ return self.bindings[expr]
+ raise ValueError("LowerComplex cannot resolve complex operand")
+
+ def map_call(self, call):
+ if any(call.op.same_as(op) for op in _OPS):
+ _validate_call(call)
+ if call.op.same_as(CMUL):
+ a, b = self.parts(call.args[0])
+ c, d = self.parts(call.args[1])
+ # (a + bi)(c + di) = (ac - bd) + (ad + bc)i
+ return ir.Call(COMPLEX, [a * c - b * d, a * d + b * c])
+ if call.op.same_as(REAL) or call.op.same_as(IMAG):
+ re, im = self.parts(call.args[0])
+ return re if call.op.same_as(REAL) else im
+ return call
+
+ def map_bind(self, bind):
+ ty = bind.var.ty
+ if isinstance(ty, ComplexType):
+ re, im = self.parts(bind.value)
+ re_var = ir.Var(f"{bind.var.name}_re", ty.dtype)
+ im_var = ir.Var(f"{bind.var.name}_im", ty.dtype)
+ self.bindings[bind.var] = (re_var, im_var)
+ return ir.SeqStmt([ir.Bind(re_var, re), ir.Bind(im_var, im)])
+ return bind
+
+ def map_sequence(self, seq):
+ return ir.SeqStmt(list(seq.seq))
+
+
+def _check_type(ty):
+ """Reject a complex type, wherever it appears: on a variable, a parameter
or an expression."""
+ if isinstance(ty, ComplexType):
+ raise ValueError("LowerComplex left an unsupported complex type")
+ return ty
+
+
+def _check_expr(expr):
+ """Reject an expression whose type is complex."""
+ _check_type(expr.ty)
+ return expr
+
+
+def _check_call_lowered(call):
+ """Reject a call to one of the four complex operators that the rewrite did
not consume."""
+ if any(call.op.same_as(op) for op in _OPS):
+ raise ValueError("LowerComplex left an unsupported complex call")
+ return call
+
+
+def _lower_function(func, _module, _ctx):
+ lowering = _Lowering()
+ rewritten = tvm_ffi.structural_map(
+ func,
+ [
+ (ir.Call, lowering.map_call),
+ (ir.Bind, lowering.map_bind),
+ (ir.SeqStmt, lowering.map_sequence),
+ ],
+ )
+ return tvm_ffi.structural_map(
+ rewritten, [(ir.Type, _check_type), (ir.Call, _check_call_lowered),
(ir.Expr, _check_expr)]
+ )
+
+
+def lower_complex() -> ir.transform.Pass:
+ """The pass that lowers complex values to scalar arithmetic."""
+ return tirx.transform.function_pass(_lower_function, opt_level=0,
name="my_ext.LowerComplex")
+
+
+# begin pipeline
+def _pipeline(*, target=None):
+ lower, finalize_host, finalize_device = tirx.get_tir_pipeline("default",
target=target)
+ return ir.transform.Sequential([lower_complex(), lower]), finalize_host,
finalize_device
+
+
+tirx.register_tir_pipeline("my_ext", _pipeline)
+# end pipeline
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index a8a54f970..5b60ad6a3 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -225,12 +225,6 @@ EXTENSION_CHECKS = """
expect_error(TypeError, "real expects a complex value", my_ext.real, a_var)
expect_error(TypeError, "imag expects a complex value", my_ext.imag, a_var)
- # Explicit result types must not let invalid calls escape validation.
- for name, args in (("complex", [a_var, b_var]), ("cmul", [z, z]), ("real",
[z]), ("imag", [z])):
- invalid = ir.Call(f"my_ext.{name}", args, ty="int32")
- expect_error(TypeError, "incorrect result type", invalid.validate)
- invalid = ir.Call("my_ext.cmul", [a_var, z], ty=z.ty)
- expect_error(TypeError, "expects a complex value", invalid.validate)
bad_type = my_ext.Complex(T.int32)
expect_error(TypeError, "expected a scalar float", my_ext.real,
ir.Var("bad", bad_type))
@@ -291,6 +285,32 @@ EXTENSION_CHECKS = """
"""
+# The validators of the Rust operators check a call whose result type is given
explicitly.
+VALIDATOR_CHECKS = """
+ for name, args in (("complex", [a_var, b_var]), ("cmul", [z, z]), ("real",
[z]), ("imag", [z])):
+ invalid = ir.Call(f"my_ext.{name}", args, ty="int32")
+ expect_error(TypeError, "incorrect result type", invalid.validate)
+ invalid = ir.Call("my_ext.cmul", [a_var, z], ty=z.ty)
+ expect_error(TypeError, "expects a complex value", invalid.validate)
+"""
+
+
+# The operators of the Python prototype validate the number of arguments, as
their signatures
+# declare; its pass checks the types of the calls it lowers.
+PROTOTYPE_CHECKS = """
+ invalid_calls = [
+ (ir.Call(f"my_ext.{name}", args, ty="int32"), "incorrect result type")
+ for name, args in (("complex", [a_var, b_var]), ("cmul", [z, z]),
("real", [z]))
+ ]
+ invalid_calls.append((ir.Call("my_ext.cmul", [a_var, z], ty=z.ty),
"expects a complex value"))
+ for invalid, message in invalid_calls:
+ invalid.validate()
+ invalid_func = tirx.Function([a_var, b_var],
ir.SeqStmt([ir.Evaluate(invalid)]))
+ expect_error(TypeError, message, lower, ir.IRModule({"main":
invalid_func}))
+ expect_error(TypeError, "argument", ir.Call("my_ext.real", [z, z],
ty="float32").validate)
+"""
+
+
def test_library_is_loaded_on_import():
"""``import tvm_tirx`` loads libtvm_tirx_rust, when the build has it, as a
backend."""
from tvm_ffi.libinfo import _find_library_by_basename # pylint:
disable=import-outside-toplevel
@@ -308,7 +328,17 @@ def test_library_is_loaded_on_import():
@pytest.mark.skipif(not tvm_tirx.runtime.enabled("llvm"), reason="LLVM is not
enabled")
def test_extension_example(extension_package):
script = EXTENSION_DIR / "run_example.py"
- run_python(script.read_text(), EXTENSION_CHECKS,
python_path=[extension_package])
+ run_python(
+ script.read_text(), EXTENSION_CHECKS, VALIDATOR_CHECKS,
python_path=[extension_package]
+ )
+
+
[email protected](not tvm_tirx.runtime.enabled("llvm"), reason="LLVM is not
enabled")
+def test_extension_prototype():
+ """The Python prototype of the extension runs the same example and
checks."""
+ script = EXTENSION_DIR / "run_example.py"
+ prototype = EXTENSION_DIR / "my_ext_py"
+ run_python(script.read_text(), EXTENSION_CHECKS, PROTOTYPE_CHECKS,
python_path=[prototype])
def test_rust_operators_are_registered_by_their_library(define_ir_library):
