Crate tvm_tirx

Rust bindings for the TIRx compiler objects, over TVM FFI.

The bindings are generated by tvm-ffi-stubgen --target rust from the reflection metadata of libtvm_tirx_compiler. The modules follow the Python package tvm_tirx:

• ir: the shared IR (expressions, statements, types, functions, modules and operators), with the primitive expressions in ir::prim and the pass infrastructure in ir::transform.
• tirx: tensor functions and types, index maps, layouts, thread scopes, tile instructions and the TIRx passes.
• sym: symbolic analysis: the analyzer, integer sets and iterator maps.
• target: compilation targets and target kinds.
• backend: backend-specific objects, such as the attributes of CUDA instructions.
• runtime: the runtime objects the IR refers to, tensor allocation and runtime queries.
• codegen: compilation of IR modules into runtime modules, and the LLVM backend queries.
• script: printing IR objects as TVMScript.

Using the bindings

The crate links libtvm_tirx_compiler, so its type keys and global functions are registered when a program starts. The build script finds the library through tvm-tirx-config --libdir, as tvm-ffi-sys finds libtvm_ffi through tvm-ffi-config --libdir: both helpers come with their Python packages (apache-tvm-tirx and apache-tvm-ffi) and must be on PATH when the crate is built. The native global functions have Rust wrappers, as methods of the handle types and as functions of the modules, such as target::Target::new for target.Target. Native functions that have no Rust wrapper are called through tvm_ffi::Function::get_global.

The crate uses the library only through the TVM FFI registry, so a linker that drops unreferenced libraries (--as-needed, which rustc passes on Linux) keeps it only because the crate holds a reference to a symbol of it, the link anchor, that is linked along with the crate. A program that calls at least one wrapper therefore has the library as a dependency, at no run-time cost. A program that reaches the registry only through tvm_ffi itself, without calling a wrapper, must keep the library by other means, such as the linker arguments tvm-tirx-config --libs compiler prints.

use tvm_tirx::target::Target;

let target = Target::new("llvm")?;
assert_eq!(target.kind()?.name()?.as_str(), "llvm");

A handle such as ir::Expr is a reference-counted pointer to a native object, so cloning it does not copy the IR, and its *Obj type (ir::ExprObj) is the object. Each binding says whether its layout is complete or opaque. A complete *Obj reproduces the native object byte for byte and exposes its reflected fields directly. The layout of an opaque one cannot be rebuilt from the reflection (it has unreflected members, or no layout metadata), so its fields are read through fallible getters over the C ABI. Use tvm_ffi::ObjectRefCast for checked casts between handles.

A complete object can also be allocated from Rust. Its new takes every field and stores it as given, without the checks, defaults or derived fields of the native constructor. Where the native constructor does more than that (ir::prim::Add checks its operand types, tirx::Function derives its type from the signature), new is reserved for a constructor with those semantics, and the allocation of the fields as given is from_fields. The IR constructors are written in Rust after the C++ ones, with the same checks, defaults, derived fields and errors; they call into C++ only for the nodes C++ shares (UnknownLoc, MissingType and the PrimType of each dtype, fetched once per thread), for the hooks an operator or type registers (such as FInferType), and for printing in an error message. A constructor returns a Result when its C++ counterpart can throw, and the node itself otherwise: ir::Var::new or ir::UnknownLoc::new cannot fail, because the native calls they make are to functions the linked library registers, and a failure there is a mismatch between the crate and the library, which panics.

The common types and constants have the short constructors C++ has: ir::PrimType::int and its siblings for PrimType::Int(32), ir::IntImm::int32, ir::IntImm::bool and ir::FloatImm::float32 for the constants, ir::prim::make_const for a constant of any type, and the conversions of i32, i64, bool, f32 and f64 into an ir::Expr that C++ makes implicitly. ir::dtype spells the DLDataType of each (dtype::int(32)) without parsing a string. The unknown location is ir::UNKNOWN_LOC as a constructor argument and ir::Location::unknown as a value.

The operators are written in Rust too, after src/ir/prim/op.cc: ir::prim::op::add and its siblings match their operand types and fold constants as C++ add does, and the math functions (ir::prim::op::floor, ir::prim::op::exp, …) build the same calls. They take expressions or literals typed as C++ types them and build at the unknown location (the _at variants take one), and an ir::prim::op::ExprResult gives them operator syntax: floordiv(&x + 1, 4)? is the C++ floordiv(x + 1, 4). An ExprResult passes into any function or constructor that takes an ir::IntoExpr, as any expression does.

Operators and types are defined from Rust in a tvm_ffi::static_init_block!, as C++ defines them in a TVM_FFI_STATIC_INIT_BLOCK. An operator is an ir::OpDef chain, the C++ OpDef, with the signature (ir::sig, whose descriptors generate the validator), the attributes, by the keys of ir::op_attr and tirx::op_attr, and the type inference, validator and region hooks as Rust functions, which C++ calls natively; a type, written outside the block as any other (Defining a type in Rust), is registered by ir::register_type:

use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
use tvm_tirx::ir::sig::{self, IntExpr};
use tvm_tirx::ir::{register_type, CallEffectKind, CallObj, Op, OpDef, Type};
use tvm_tirx::tvm_ffi::{self, Result};

fn index_type(call: &CallObj) -> Result<Type> {
    Ok(call.args.get(0)?.ty.clone())
}

tvm_ffi::static_init_block! {
    OpDef::new("my_ext.clamp", "Clamp an index to the range `[0, bound)`.")
        .signature((
            sig::arg::<IntExpr>("index", "The index."),
            sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
        ))
        .set_infer_type(index_type)
        .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
    register_type::<PairType>()?;
}

let clamp = Op::get("my_ext.clamp")?;

A block runs when its program or library loads, as a C++ one does, wherever in it the block is. ir::Op::get then returns the operator, as C++ Op::Get does: the registry keeps one object per name, so call.op.same_as(&clamp) tests whether a call is to it. Every library built on the crate exports tvm_ffi::run_static_init_blocks, 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 and types the crate defines by the time the import returns, as the C++ ones are, and calls tirx_register_ir to raise the error of a block that failed. Python uses the operators as it uses the C++ ones: Op.get finds them, a Call built in Python runs their hooks, and the TVMScript printer prints their calls. Completing an operator, when its ir::OpDef chain ends, or a type with ir::register_type requests API exposure through the ir.expose_op or ir.expose_type FFI receiver. Python installs these receivers before loading native libraries and publishes the selected definitions to registered destinations, including destinations initialized later. Op API destinations construct raw calls; script destinations add builder behavior. Type destinations publish constructors, such as T.<Name> for tirx.<Name>Type. Handwritten Python bindings take precedence in each destination. Without receivers, native registration proceeds without API exposure.

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

[tvm_ffi::rs_class] defines a new IR type, a type of expressions as ir::PrimType and ir::TupleType are, from its object struct, whose first field base is the ir::TypeObj it extends, which holds the location. It registers the type under ir.Type with the reflection of every other field, so that the C++ structural equality and hashing, deep copy, repr, the TVMScript printer and Python see the fields as they see a C++ type’s. [tvm_ffi::def_object_ref] defines the reference, which converts into ir::Type. The object’s new takes the TypeObj, then the fields in order:

use tvm_tirx::ir::{Location, PrimType, Type, TypeObj};
use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class, ObjectRefCore};

/// The type of a pair of values.
#[rs_class("my_ext.PairType", final)]
pub struct PairTypeObj {
    base: TypeObj,
    pub first: Type,
    pub second: Type,
    /// Not part of structural equality or hashing.
    #[reflect(seq_hash_ignore)]
    pub note: tvm_ffi::String,
}

def_object_ref!(pub PairType, Type, PairTypeObj);

impl PairType {
    pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
        PairTypeObj::new(TypeObj::new(Location::unknown()), first, second, "").into()
    }
}

fn main() {
    let pair = PairType::new(PrimType::int(32), PrimType::float(32));
    let ty: Type = pair.into(); // an `ir.Type` that C++ sees as `my_ext.PairType`
    assert!(ty.downcast_ref::<PairTypeObj>().is_some());
}

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 and UnsafeInit; the bindings generate them for a complete type that is not final, as ir.Type. Registered by ir::register_type in a tvm_ffi::static_init_block!, with the operators, the type is also constructed by Python from its fields and parsed by TVMScript, as a C++ type.

Generated and hand-written code

The code between a // tvm-ffi-stubgen(begin) and its // tvm-ffi-stubgen(end) is generated and rewritten in place. Everything outside the blocks is hand-written, including the tvm-ffi-stubgen(...) directives that steer the generator.

The blocks of the objects a Python module defines are in the Rust file at the same path: python/tvm_tirx/ir/prim/expr.py is src/ir/prim/expr.rs. target and runtime, whose objects come from a single Python module each, are one file. A type is in the module where Python defines it, which is not always its type-key namespace: tirx.cuda.* is in backend::cuda and tirx.tensor.* in tirx::tile. A type Python has no class for is next to the types it belongs to, and its file says so. A parent module re-exports its files, so ir::Expr names ir::expr::Expr. It also holds the prefix directive of its namespace: an object registered later that no file has a block for yet is generated there, to be moved to its file.

Regenerate from the repository root after rebuilding libtvm_tirx_compiler (add --check to only report stale blocks):

tvm-ffi-stubgen rust/src --target rust --dlls build/lib/libtvm_tirx_compiler.so

tvm-ffi-stubgen comes with the tvm_ffi Python package, which must be the one of 3rdparty/tvm-ffi: the blocks are what its generator emits. The blocks are formatted as rustfmt formats them, so cargo fmt --check and tvm-ffi-stubgen --check agree.

Tests

Build libtvm_tirx_compiler first. The build scripts need tvm-ffi-config and tvm-tirx-config on PATH, which uv run provides from the .venv after the editable install, where tvm-tirx-config --libdir is build/lib:

uv run --no-sync cargo test --manifest-path rust/Cargo.toml

Set TVM_TIRX_LIBRARY_PATH to the directory of another libtvm_tirx_compiler to build and test against that one; tvm-tirx-config puts it first.

The tests are laid out like the modules. Each top-level module is one test binary, tests/<module>/main.rs, and the tests of src/<path>.rs are in tests/<path>.rs: tests/tirx/transform/function_pass.rs tests src/tirx/transform/function_pass.rs. Run the tests of one module with --test <module>, such as --test tirx. tests/common/ holds the helpers the binaries share.

Re-exports

• pub use tvm_ffi;

Modules

• backend
  Backend-specific objects.
• codegen
  Code generation: IR modules compiled into runtime modules.
• ir
  The shared IR: expressions, statements, types, functions, modules and operators.
• runtime
  The runtime objects the IR refers to.
• script
  TVMScript: IR objects printed as the Python script that builds them.
• sym
  Symbolic analysis of integer expressions.
• target
  Compilation targets.
• tirx
  Tensor programs: functions, tensor types, index maps, layouts and thread scopes.

Macros

• def_expr_op_over_loads
  Add primitive-expression operators to an existing reference wrapper.