tvm_tirx/

lib.rs

/*
 * 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.
 */

//! The TIRx compiler in Rust: its IR objects over TVM FFI, and the definition of new types,
//! operators and passes that extend it.
//!
//! The IR 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.
//!
//! ```no_run
//! use tvm_tirx::target::Target;
//!
//! let target = Target::new("llvm")?;
//! assert_eq!(target.kind()?.name()?.as_str(), "llvm");
//! # Ok::<(), tvm_tirx::tvm_ffi::Error>(())
//! ```
//!
//! 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](#defining-a-type-in-rust)), is registered by
//! [`ir::register_type`]:
//!
//! ```no_run
//! 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};
//! # use tvm_tirx::ir::TypeObj;
//! # use tvm_tirx::tvm_ffi::{def_object_ref, rs_class, ObjectRefCore};
//! # #[rs_class("my_ext.PairType", final)]
//! # pub struct PairTypeObj { base: TypeObj, pub first: Type, pub second: Type }
//! # def_object_ref!(pub PairType, Type, PairTypeObj);
//!
//! 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")?;
//! # Ok::<(), tvm_ffi::Error>(())
//! ```
//!
//! 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:
//!
//! ```no_run
//! 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):
//!
//! ```bash
//! 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`:
//!
//! ```bash
//! 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.

// The crate reaches `libtvm_tirx_compiler` through the FFI registry, so a linker that drops
// unreferenced libraries (`--as-needed`) would drop the library the build script links, and its
// registrations would never run. Each TIRx shared library exports a `<library>_link_anchor`
// function that does nothing, defined in the `link_anchor.cc` of its sources (for the compiler,
// `src/ir/link_anchor.cc`) and declared by whoever references it. The static below holds the
// anchor's address: a relocation against the symbol in the object of this module, which costs
// nothing at run time and is what makes the linker keep the library.
//
// A program pulls in only the objects of the crate it references, so the static keeps the
// library only when the object of this module is pulled. It is, because the objects of the crate
// reference one another densely and `sym::analyzer` references `ExitOnDrop::new` below. The test
// binaries rely on the same mechanism and fail at their first registry lookup if that changes. A
// program that uses 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.
extern "C" {
    fn tvm_tirx_compiler_link_anchor() -> i32;
}

#[used]
static TVM_TIRX_COMPILER_LINK_ANCHOR: unsafe extern "C" fn() -> i32 = tvm_tirx_compiler_link_anchor;

/// Makes a scope guard neither `Send` nor `Sync`: the native scopes keep their state per thread,
/// or in an object that is not synchronized, so a scope must end on the thread that opened it.
pub(crate) type ThreadBound = std::marker::PhantomData<*const ()>;

/// Calls a native exit function when dropped: the end of a scope a native enter function opened.
pub(crate) struct ExitOnDrop(tvm_ffi::Function, ThreadBound);

impl ExitOnDrop {
    pub(crate) fn new(exit: tvm_ffi::Function) -> Self {
        Self(exit, ThreadBound::default())
    }
}

impl Drop for ExitOnDrop {
    fn drop(&mut self) {
        // A destructor cannot report the error, and the native exit functions do not fail.
        let _ = self.0.call_packed(&[]);
    }
}

pub mod backend;
pub mod codegen;
pub mod ir;
mod r#macro;
pub mod runtime;
pub mod script;
pub mod sym;
pub mod target;
pub mod tirx;

pub use tvm_ffi;

// The errors of the static init blocks of every library built on the crate, which
// `import tvm_tirx` raises for `libtvm_tirx_rust`: `__tvm_ffi_tirx_register_ir`, the TVM FFI
// function `tirx_register_ir` of the library as `tvm_ffi.load_module` loads it. The module keeps
// the symbol, which `#[no_mangle]` exports from wherever it is, out of the crate's documentation.
#[doc(hidden)]
pub mod exports {
    tvm_ffi::tvm_ffi_dll_export_typed_func!(tirx_register_ir, tvm_ffi::run_static_init_blocks);
}