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 5b297717d [REFACTOR][RUST] Load libtvm_tirx_rust through the backend
autoload (#51)
5b297717d is described below
commit 5b297717d18e38caa17d92b5ab1c29c1d61383ce
Author: Linzhang Li <[email protected]>
AuthorDate: Sat Oct 24 20:11:57 2026 -0400
[REFACTOR][RUST] Load libtvm_tirx_rust through the backend autoload (#51)
`import tvm_tirx` no longer loads `libtvm_tirx_rust` from `base.py`,
right after the compiler
library. A built-in backend, `tvm_tirx.backend.rust`, loads it instead,
through the backend
autoload, as the CUDA, Metal and ROCm backends load their runtime
libraries. The library is still
built by the TIRx build (`TVM_TIRX_USE_RUST`); moving it into a package
of its own is left for
later.
---
cmake/modules/Rust.cmake | 5 +-
docs/reference/python/backend.rst | 1 +
.../python/{backend.rst => backend.rust.rst} | 32 ++----------
docs/reference/python/index.rst | 3 ++
python/tvm_tirx/backend/_autoload_backends.py | 1 +
python/tvm_tirx/backend/rust/__init__.py | 45 +++++++++++++++++
python/tvm_tirx/base.py | 22 ---------
rust/examples/define_ops.rs | 4 +-
rust/src/lib.rs | 15 +++---
tests/python/ir/test_op_rust.py | 57 ++++++++++++++++------
10 files changed, 108 insertions(+), 77 deletions(-)
diff --git a/cmake/modules/Rust.cmake b/cmake/modules/Rust.cmake
index cc28707a0..66e260220 100644
--- a/cmake/modules/Rust.cmake
+++ b/cmake/modules/Rust.cmake
@@ -16,8 +16,9 @@
# under the License.
# The crate in rust/, built by cargo as a shared library and installed as
libtvm_tirx_rust next to
-# libtvm_tirx_compiler. `import tvm_tirx` loads it after the compiler library
and registers the
-# operators it defines (python/tvm_tirx/base.py), as loading the compiler
registers the C++ ones.
+# libtvm_tirx_compiler. `import tvm_tirx` loads it through the backend
autoload and registers the
+# operators it defines (python/tvm_tirx/backend/rust), as loading the compiler
registers the C++
+# ones.
#
# TVM_TIRX_USE_RUST is ON, OFF or AUTO. AUTO builds the library when it can:
not on Windows, with
# the crate's sources (a source distribution has none), with a cargo at least
as new as the
diff --git a/docs/reference/python/backend.rst
b/docs/reference/python/backend.rst
index d5a7352bc..40a7f6da3 100644
--- a/docs/reference/python/backend.rst
+++ b/docs/reference/python/backend.rst
@@ -32,6 +32,7 @@ Namespaces
* :doc:`tvm_tirx.backend.metal <backend.metal>`
* :doc:`tvm_tirx.backend.opencl <backend.opencl>`
* :doc:`tvm_tirx.backend.rocm <backend.rocm>`
+* :doc:`tvm_tirx.backend.rust <backend.rust>`
* :doc:`tvm_tirx.backend.trn <backend.trn>`
* :doc:`tvm_tirx.backend.vulkan <backend.vulkan>`
* :doc:`tvm_tirx.backend.webgpu <backend.webgpu>`
diff --git a/docs/reference/python/backend.rst
b/docs/reference/python/backend.rust.rst
similarity index 55%
copy from docs/reference/python/backend.rst
copy to docs/reference/python/backend.rust.rst
index d5a7352bc..549a6b24d 100644
--- a/docs/reference/python/backend.rst
+++ b/docs/reference/python/backend.rust.rst
@@ -15,34 +15,13 @@
specific language governing permissions and limitations
under the License.
-tvm_tirx.backend
-================
+tvm_tirx.backend.rust
+=====================
-.. automodule:: tvm_tirx.backend
+.. automodule:: tvm_tirx.backend.rust
:no-members:
-.. currentmodule:: tvm_tirx.backend
-
-Namespaces
-----------
-
-* :doc:`tvm_tirx.backend.cuda <backend.cuda>`
-* :doc:`tvm_tirx.backend.adreno <backend.adreno>`
-* :doc:`tvm_tirx.backend.hexagon <backend.hexagon>`
-* :doc:`tvm_tirx.backend.metal <backend.metal>`
-* :doc:`tvm_tirx.backend.opencl <backend.opencl>`
-* :doc:`tvm_tirx.backend.rocm <backend.rocm>`
-* :doc:`tvm_tirx.backend.trn <backend.trn>`
-* :doc:`tvm_tirx.backend.vulkan <backend.vulkan>`
-* :doc:`tvm_tirx.backend.webgpu <backend.webgpu>`
-* :doc:`tvm_tirx.backend.config <backend.config>`
-
-.. toctree::
- :hidden:
- :maxdepth: 1
- :titlesonly:
-
- backend.config
+.. currentmodule:: tvm_tirx.backend.rust
Functions
---------
@@ -50,5 +29,4 @@ Functions
.. autosummary::
:toctree: generated/
- is_loaded
- load
+ register_backend
diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst
index 162d3a8d3..9b5d0a919 100644
--- a/docs/reference/python/index.rst
+++ b/docs/reference/python/index.rst
@@ -102,6 +102,8 @@ Namespaces
- Trainium NKI intrinsic wrappers.
* - :doc:`tvm_tirx.backend.trn.transform <backend.trn.transform>`
- Trainium-specific lowering passes.
+ * - :doc:`tvm_tirx.backend.rust <backend.rust>`
+ - Operators written in Rust.
* - :doc:`tvm_tirx.support <support>`
- Build tools and system utilities.
* - :doc:`tvm_tirx.testing <testing>`
@@ -147,6 +149,7 @@ Namespaces
tvm_tirx.backend.trn <backend.trn>
tvm_tirx.backend.trn.op <backend.trn.op>
tvm_tirx.backend.trn.transform <backend.trn.transform>
+ tvm_tirx.backend.rust <backend.rust>
tvm_tirx.support <support>
tvm_tirx.testing <testing>
diff --git a/python/tvm_tirx/backend/_autoload_backends.py
b/python/tvm_tirx/backend/_autoload_backends.py
index ad333064e..5b8e2d2f8 100644
--- a/python/tvm_tirx/backend/_autoload_backends.py
+++ b/python/tvm_tirx/backend/_autoload_backends.py
@@ -34,6 +34,7 @@ _BUILTIN_BACKENDS = (
"webgpu",
"hexagon",
"adreno",
+ "rust",
)
_AUTO_LOAD_DONE = False
diff --git a/python/tvm_tirx/backend/rust/__init__.py
b/python/tvm_tirx/backend/rust/__init__.py
new file mode 100644
index 000000000..13aebfd49
--- /dev/null
+++ b/python/tvm_tirx/backend/rust/__init__.py
@@ -0,0 +1,45 @@
+# 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.
+"""Rust-owned backend hooks."""
+
+from pathlib import Path
+
+import tvm_ffi
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm_tirx.base import _LOADED_LIBS
+
+
+def register_backend():
+ """Register Rust-owned Python semantics."""
+ runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+ try:
+ _LOADED_LIBS["tvm_tirx_rust"] = load_lib_ctypes(
+ package="tvm_tirx",
+ target_name="tvm_tirx_rust",
+ extra_lib_paths=[runtime_dir],
+ mode="RTLD_LOCAL",
+ )
+ except (OSError, FileNotFoundError, RuntimeError):
+ pass
+ else:
+ # Unlike a runtime sidecar, the library registers its operators only
when asked to.
+
tvm_ffi.load_module(_LOADED_LIBS["tvm_tirx_rust"]._name).tirx_register_ops()
+ return None
+
+
+__all__ = ["register_backend"]
diff --git a/python/tvm_tirx/base.py b/python/tvm_tirx/base.py
index 1cfdc9249..af8665125 100644
--- a/python/tvm_tirx/base.py
+++ b/python/tvm_tirx/base.py
@@ -18,12 +18,9 @@
# pylint: disable=invalid-name, import-outside-toplevel
"""Base library for TVM."""
-import ctypes
import os
-import sys
from pathlib import Path
-import tvm_ffi
from tvm_ffi import libinfo as tvm_ffi_libinfo
from tvm_ffi.libinfo import load_lib_ctypes
@@ -64,25 +61,6 @@ if not _RUNTIME_ONLY:
"tvm_tirx", "tvm_tirx_compiler", "RTLD_LOCAL",
extra_lib_paths=_lib_paths
)
-# The operators written in Rust are in libtvm_tirx_rust, when the build has it
-# (``TVM_TIRX_USE_RUST``). They are registered now, as loading the compiler
library registered the
-# C++ ones, before any module looks for them; ``TVM_TIRX_RUST_LIBRARY`` names
another build of the
-# library to load instead, such as ``rust/target/debug/libtvm_tirx.so``.
Otherwise only the one next
-# to the compiler library is loaded, which the same build made: one found
elsewhere, such as on
-# ``LD_LIBRARY_PATH``, may be another build's.
-if not _RUNTIME_ONLY:
- _rust_lib_path = os.environ.get("TVM_TIRX_RUST_LIBRARY")
- if not _rust_lib_path:
- _rust_lib_path =
Path(_LOADED_LIBS["tvm_tirx_compiler"]._name).resolve().parent / (
- "libtvm_tirx_rust.dylib" if sys.platform == "darwin" else
"libtvm_tirx_rust.so"
- )
- if not _rust_lib_path.is_file():
- _rust_lib_path = None
- if _rust_lib_path:
- # The handle keeps the library, and the hooks it registers, loaded for
the process.
- _LOADED_LIBS["tvm_tirx_rust"] = ctypes.CDLL(str(_rust_lib_path),
ctypes.RTLD_LOCAL)
- tvm_ffi.load_module(str(_rust_lib_path)).tirx_register_ops()
-
if _RUNTIME_ONLY:
from tvm_ffi import registry as _tvm_ffi_registry
diff --git a/rust/examples/define_ops.rs b/rust/examples/define_ops.rs
index 814a02516..7bea93fc3 100644
--- a/rust/examples/define_ops.rs
+++ b/rust/examples/define_ops.rs
@@ -21,8 +21,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 in place of `libtvm_tirx_rust`
-//! (`TVM_TIRX_RUST_LIBRARY`), so `import tvm_tirx` registers these operators.
As a C++ operator
+//! `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
//! `with T.example_loop(n) as j:`.
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index c4d377153..1706935ef 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -138,14 +138,13 @@
//!
//! 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
-//! after the compiler library and calls `tirx_register_ops`, so the operators
the crate defines
-//! are registered before any Python module looks for them, 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. Set `TVM_TIRX_RUST_LIBRARY` to the path of
another build of the
-//! library, such as `rust/target/debug/libtvm_tirx.so` from `cargo build`, to
load that one
-//! instead. `examples/define_ops.rs` builds a library of its own the same way.
+//! `rust-version` (`TVM_TIRX_USE_RUST`, `AUTO` by default). `import tvm_tirx`
loads it through
+//! 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.
//!
//! # Defining a type in Rust
//!
diff --git a/tests/python/ir/test_op_rust.py b/tests/python/ir/test_op_rust.py
index f740e4c1a..a40da764a 100644
--- a/tests/python/ir/test_op_rust.py
+++ b/tests/python/ir/test_op_rust.py
@@ -14,11 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Operators defined in Rust are registered on ``import tvm_tirx``, as C++
ones are.
+"""Operators defined in Rust are registered when their library is loaded, and
used as C++ ones are.
-``rust/examples/define_ops.rs`` is built by cargo and loaded in place of
libtvm_tirx_rust
-(``TVM_TIRX_RUST_LIBRARY``) by a fresh interpreter, which then uses its
operators as it uses C++
-ones: their TVMScript builders are Python functions written for them.
+``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.
"""
import os
@@ -26,6 +26,7 @@ import re
import shutil
import subprocess
import sys
+import tempfile
import textwrap
from pathlib import Path
@@ -65,26 +66,49 @@ def define_ops_library():
[cargo, "build", "--manifest-path", str(manifest), "--example",
"define_ops"], check=True
)
target_dir = Path(os.environ.get("CARGO_TARGET_DIR", RUST_DIR / "target"))
- name = "libdefine_ops.dylib" if sys.platform == "darwin" else
"libdefine_ops.so"
+ name = {"darwin": "libdefine_ops.dylib", "win32": "define_ops.dll"}.get(
+ sys.platform, "libdefine_ops.so"
+ )
return target_dir / "debug" / "examples" / name
def run_python(code, library):
- """Run `code` in a fresh interpreter that loads `library` as
libtvm_tirx_rust."""
- env = dict(os.environ, TVM_TIRX_RUST_LIBRARY=str(library))
- result = subprocess.run(
- [sys.executable, "-c", textwrap.dedent(code)],
- env=env,
- capture_output=True,
- text=True,
- check=False,
- )
+ """Run `code` in a fresh interpreter, with the path of `library` as its
argument.
+
+ The code runs from a file: TVMScript reads the source of the functions it
parses, which Python
+ keeps for code passed with `-c` only from 3.13 on.
+ """
+ with tempfile.TemporaryDirectory() as tmp:
+ script = Path(tmp) / "script.py"
+ script.write_text(textwrap.dedent(code))
+ result = subprocess.run(
+ [sys.executable, str(script), str(library)],
+ capture_output=True,
+ text=True,
+ check=False,
+ )
assert result.returncode == 0, result.stderr
-def test_rust_operators_are_registered_on_import(define_ops_library):
+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
+
+ from tvm_tirx.base import _LOADED_LIBS # pylint:
disable=import-outside-toplevel
+
+ runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+ try:
+ _find_library_by_basename("tvm_tirx", "tvm_tirx_rust",
extra_lib_paths=[runtime_dir])
+ except RuntimeError:
+ pytest.skip("the build has no libtvm_tirx_rust (TVM_TIRX_USE_RUST)")
+ assert "tvm_tirx_rust" in _LOADED_LIBS
+
+
+def test_rust_operators_are_registered_by_their_library(define_ops_library):
run_python(
"""
+ import sys
+
import tvm_ffi
import tvm_tirx
from tvm_tirx import ir
@@ -92,7 +116,8 @@ def
test_rust_operators_are_registered_on_import(define_ops_library):
from tvm_tirx.script.ir_builder.base import auto_emit_void_call
from tvm_tirx.tirx.script import ir_builder
- # Registered on import, with the doc comment and signature the Rust
definition gives.
+ # Registered by the library, with the doc comment and signature the
Rust definition gives.
+ tvm_ffi.load_module(sys.argv[1]).tirx_register_ops()
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"]
