Quick Start#

Note

All the code in this guide is under examples/quick_start in the repository.

This guide walks through a small add_one program that computes B = A + 1 on CPU and CUDA with TIRx, a minimal, open ML compiler foundation:

Programs in Python: tensor programs are written as Python functions and compiled directly.
Transparent program structure: every element of a program, from functions to statements and types, is accessible from Python and other languages, so developers and agents can inspect and transform it.
Built for extension: all compiler APIs are open, so developers and agents can quickly iterate on extensions in separate packages that link to TIRx dynamically, without rebuilding one monolithic bundle.

This guide shows the first two in action and gives a first taste of the third.

Prerequisite

Python 3.10 or newer
PyTorch, to create and check tensors
CUDA 13 or newer environment, for the CUDA example

TIRx, installed with:

pip install --upgrade --force-reinstall apache-tvm-tirx
# add the cuda-bindings dependency:
# pip install --upgrade --force-reinstall "apache-tvm-tirx[cuda]"

Write a Simple `add_one`#

A TIRx program is an ordinary Python function marked with @T.function. Here add_one takes an input tensor A and an output tensor B, each holding n float32 values. We create n with T.dynamic instead of picking a fixed number, so the same program works for tensors of any length.

import torch

import tvm_tirx
from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[i] = A[i] + 1

import torch

import tvm_tirx
from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, 256), block=256))
    i = T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x")
    if i < n:
        B[i] = A[i] + 1

The CPU version is a loop that walks over all n elements. The CUDA version spreads the work across many GPU threads instead: T.device_entry starts the GPU code and launches it in blocks of 256 threads, with enough blocks to cover all n elements. Each thread works out the element i it owns and updates only that one. The last block may have more threads than elements left, so the if i < n check keeps those extra threads from going past the end of the tensors.

Compile and Run#

With the program written, we compile it for a target, "llvm" for the CPU or "cuda" for an NVIDIA GPU, using tvm_tirx.compile. The result is a compiled program that we call like a function, passing PyTorch tensors directly, so we create the inputs with torch and check the output with torch.testing.assert_close:

add_one_compiled = tvm_tirx.compile(add_one, target="llvm")
a = torch.arange(10, dtype=torch.float32)
b = torch.empty_like(a)
add_one_compiled(a, b)
torch.testing.assert_close(b, a + 1)
print(b)

add_one_compiled = tvm_tirx.compile(add_one, target="cuda")
a = torch.arange(10, dtype=torch.float32, device="cuda")
b = torch.empty_like(a)
add_one_compiled(a, b)
torch.testing.assert_close(b, a + 1)
print(b)

The script prints the result, each input value plus one (the CUDA version also shows device='cuda:0'):

tensor([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.])

Inspect and Transform the Program#

Besides compiling and running add_one, we can also look inside it. The add_one returned by @T.function is a tree of program elements, commonly called an intermediate representation (IR). To see how the program is structured, we can use format_tree to display its IR as a tree:

import tvm_tirx

print(tvm_tirx.ir.format_tree(add_one))

Function
├── params[0]: A = Var(ty=TensorType(shape=[n], dtype=float32))
├── params[1]: B = Var(ty=TensorType(shape=[n], dtype=float32))
└── body: SeqStmt
    └── seq[0]: For(begin=IntImm(value=0), end=n)
        ├── var: i = Var(ty=int32)
        └── body: SeqStmt
            └── seq[0]: TensorStore(dest=B, indices=[i])
                └── value: Add
                    ├── a: TensorLoad(source=A, indices=[i])
                    └── b: FloatImm(value=1.0)

Every element is accessible from Python and other languages, so we can walk this tree directly in Python:

A, B = add_one.params       # parameters are Vars
n = A.ty.shape[0]           # the symbolic length n, also a Var
loop = add_one.body[0]      # for i in range(n)
print(loop.body[0].value)   # A[i] + T.float32(1.0)

Named values in the program, such as the parameters A and B, the length n, and the loop index i, are variables represented by Var objects. Note that the constant 1 was converted to T.float32(1.0) to match the element type.

These elements are minimal TVM-FFI objects, which share one representation across languages. On top of them, TIRx provides reusable APIs to build analyses and transformations. For example, we can specialize add_one for tensors of exactly 128 elements by replacing every use of n with 128:

add_one_specialized = tvm_ffi.structural_map(
    add_one,
    (tvm_tirx.ir.Var, lambda v: T.int32(128) if v.same_as(n) else v),
)
add_one_specialized.show()

@T.function
def add_one(A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)):
    for i in range(128):
        B[i] = A[i] + T.float32(1.0)

tvm_ffi.structural_map walks every element of the program and calls our function on each Var, replacing it with whatever we return. The result, add_one_specialized, is a new program specialized to tensors of 128 elements. We can compile it with tvm_tirx.compile like any other program, while the original add_one stays unchanged.

Export and Load the Compiled Program#

So far, we have compiled and run add_one in the same Python process. In practice, we often want to compile a program once and run it somewhere else: in another process, on another machine, or inside an application that is not written in Python.

We support this by exporting the compiled program as a shared library. Under the hood, tvm_tirx.compile returns an Executable. It holds the compiled module, which contains the function add_one; calling the executable, as we did earlier, calls that function. The executable’s export_library method writes this module to a shared library:

add_one_compiled.export_library("./add_one_cpu.so")

add_one_compiled.export_library("./add_one_cuda.so")

The exported library is a regular TVM-FFI module, so any language or framework that TVM-FFI supports, such as C++, Rust, and JAX, can load and call it. Running it does not need the compiler, only the minimal TVM runtime libraries, libtvm_runtime.so and libtvm_runtime_cuda.so (for CUDA), each under 1 MB.

Load in Python#

Later, for example in a new Python process, tvm_ffi.load_module loads the module back. Importing tvm_tirx loads the runtime libraries the module needs. The loaded module exposes its functions by name, so we call add_one as an attribute of the module:

import torch
import tvm_ffi

# Load the runtime libraries.
import tvm_tirx  # noqa: F401

loaded = tvm_ffi.load_module("./add_one_cpu.so")
a = torch.arange(10, dtype=torch.float32)
b = torch.empty_like(a)
loaded.add_one(a, b)
torch.testing.assert_close(b, a + 1)

import torch
import tvm_ffi

# Load the runtime libraries.
import tvm_tirx  # noqa: F401

loaded = tvm_ffi.load_module("./add_one_cuda.so")
a = torch.arange(10, dtype=torch.float32, device="cuda")
b = torch.empty_like(a)
loaded.add_one(a, b)
torch.testing.assert_close(b, a + 1)

Load from C++#

The following C++ program loads the CPU build, add_one_cpu.so, and calls add_one through the TVM-FFI C++ API. TIRx kernels require 64-byte aligned tensor data for performance, so the program uses the runtime tensor allocator to provide the required alignment on each platform:

#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/tvm_ffi.h>
#include <tvm/runtime/tensor.h>

#include <iostream>

namespace ffi = tvm::ffi;

int main() {
  DLDataType f32{kDLFloat, 32, 1};
  DLDevice cpu{kDLCPU, 0};
  // The runtime allocator provides the alignment required by TIRx kernels.
  ffi::Tensor a = tvm::runtime::Empty({10}, f32, cpu);
  ffi::Tensor b = tvm::runtime::Empty({10}, f32, cpu);
  float* a_data = static_cast<float*>(a.data_ptr());
  for (int i = 0; i < 10; ++i) a_data[i] = static_cast<float>(i);

  ffi::Module mod = ffi::Module::LoadFromFile("./add_one_cpu.so");
  ffi::Function add_one = mod->GetFunction("add_one").value();
  add_one(a, b);

  const float* b_data = static_cast<const float*>(b.data_ptr());
  for (int i = 0; i < 10; ++i) std::cout << b_data[i] << " ";
  std::cout << std::endl;
  return 0;
}

To build it, we link against only the runtime libraries, which tvm-tirx-config finds for us, so the program runs without the compiler. Run it from the directory that contains add_one_cpu.so:

g++ -O2 load_cpp.cc \
    $(tvm-tirx-config --cxxflags --ldflags --libs runtime) \
    -Wl,-rpath,$(tvm-tirx-config --libdir) \
    -Wl,-rpath,$(tvm-ffi-config --libdir) \
    -o load_cpp
./load_cpp

The same program can also be built with CMake. The CMakeLists.txt in examples/quick_start finds the TIRx and TVM-FFI CMake packages with python -m tvm_tirx.config --cmakedir and python -m tvm_ffi.config --cmakedir, then links load_cpp against tvm_runtime::shared and tvm_ffi::shared:

cmake . -B build -DCMAKE_BUILD_TYPE=RelWithDebInfo
cmake --build build --config RelWithDebInfo
build/load_cpp

Troubleshooting#

If compiling for CUDA fails with ValueError: CUDA compilation requires an explicit backend_config arch or Target arch when no CUDA device is available even though the machine has a GPU, TIRx could not find the CUDA runtime library libcudart.so.13. Install and import a CUDA 13 build of PyTorch, which ships this library, or install the CUDA 13 toolkit and add its lib64 directory to LD_LIBRARY_PATH.

Quick Start

Contents

Quick Start#

Write a Simple `add_one`#

Compile and Run#

Inspect and Transform the Program#

Export and Load the Compiled Program#

Load in Python#

Load from C++#

Troubleshooting#

Further Reading#

Quick Start

Contents

Quick Start#

Write a Simple add_one#

Compile and Run#

Inspect and Transform the Program#

Export and Load the Compiled Program#

Load in Python#

Load from C++#

Troubleshooting#

Further Reading#

Write a Simple `add_one`#