TensorIR Compiler

Early Release Status

TensorIR is being released early so the community can evaluate the project, provide feedback, and engage with the API and development direction.

This release is not yet intended as a performance benchmark or production-performance commitment. Performance may vary across workloads, configurations, hardware, and compiler versions as we continue to mature code generation, heuristics, autotuning, validation, and integration.

We welcome issue reports, feedback on usability and supported workflows, and contributions that help improve the project.

Building TensorIR Compiler

Prerequisites

Required:

• CMake 3.20.0 or later
• C++17 compatible compiler
• CUDA Toolkit 13.3 or later and a compatible NVIDIA driver
• Python 3.10 or later
• Ninja build system

CUDA Toolkit 13.1 can be used for kernels supported by the compatibility bytecode format by passing --bytecode-version=compatibility; the default compiler and test workflow requires CUDA Toolkit 13.3 or later.

Python bindings (enabled by default):

• Python development headers
• nanobind 2.9 or later

Pass -DTENSOR_IR_ENABLE_BINDINGS_PYTHON=OFF to build only the compiler without these Python binding dependencies.

Python DSL profiling (optional):

• cuda-python
• cupti-python

Testing/examples:

• pytest and PyTorch, validated with pytest 8.3.4 and PyTorch 2.10

Quick Start

TensorIR Compiler depends on CUDA Tile IR for GPU code generation. LLVM and MLIR are located through CMake config packages using find_package(MLIR CONFIG) and find_package(LLVM CONFIG). When TENSOR_IR_DOWNLOAD_LLVM=OFF, set MLIR_DIR to the directory containing MLIRConfig.cmake, or override package discovery using another method supported by CMake's find_package command.

cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
cmake --build build --target tensor_ir-compiler tensor_ir-opt tensor_ir_python_bindings --parallel 32

Build Options

To use an existing compatible LLVM installation:

cmake -S . -B build -G Ninja \
  -DTENSOR_IR_DOWNLOAD_LLVM=OFF \
  -DMLIR_DIR=/path/to/llvm/lib/cmake/mlir

The build produces:

• build/bin/tensor_ir-compiler
• build/bin/tensor_ir-opt
• build/python_packages/nv_tensor_ir (when Python bindings are enabled)

To use the Python bindings or DSL directly from the build tree, run from the repository root with:

export PYTHONPATH="$PWD/build/python_packages:$PYTHONPATH"

The Python package contains both the lower-level MLIR bindings under nv_tensor_ir._mlir and the DSL under nv_tensor_ir.dsl.

cmake --install build installs:

• bin/tensor_ir-compiler
• bin/tensor_ir-opt
• python_packages/nv_tensor_ir (when Python bindings are enabled)

Use cmake --install build --prefix <install-prefix> to install into a user-writable directory. Add <install-prefix>/python_packages to PYTHONPATH to import the installed Python package.

Usage

The compiler is compile-only by default: it does not build a reference graph or launch the kernel. --launch and --verify require an NVIDIA GPU compatible with the selected target architecture and a compatible driver; --verify launches the kernel and compares its results with the reference implementation.

Compile MLIR to GPU Kernel

build/bin/tensor_ir-compiler \
  test/Integration/Compiler/matmul_8x8x8.mlir \
  --verbose --print-ir-after-all

Compile Dynamic-Shape MLIR

build/bin/tensor_ir-compiler \
  test/Integration/Compiler/add_dynamic.mlir \
  --dynamic-dims=16,8 --dynamic-strides=8 --tile-size=8x8 \
  --verify

--dynamic-dims and --dynamic-strides provide runtime values for ? dimensions and strides by tensor dimension position, and the same values are reused across tensors. When a dynamic dimension position exceeds the supplied list, the last dimension value is reused. A non-empty dynamic-stride list must provide a value for every dynamic stride position; if --dynamic-strides is omitted, dynamic strides are inferred as packed strides. --tile-size controls compiler tiling; it is not a runtime problem-size option.

Command-Line Options

The options below cover the common compile and run workflow. For the complete option list and defaults, run build/bin/tensor_ir-compiler --help.

General:

• --verbose - Enable verbose output
• --dump-artifact=<path> - Write compiled CUDA Tile device code to <path> and metadata to <path>.meta

Compilation:

• --dump-ir=<path> - Dump lowered CUDA Tile dialect MLIR
• --dump-tileir-bc=<path> - Dump Tile IR bytecode after compilation
• --bytecode-version=<version> - Tile IR bytecode target: default, current, or compatibility. The default target is max(compatibility, 13.3).
• --codegen-strategy=<strategy> - TensorIR-to-CUDA-Tile lowering path: layout-propagation (default) or affine-map
• --host-target=<target> - Host execution: auto (default) selects native code for OSS CUDA Tile layout propagation and interpreter execution for other pipelines. Use interpreter to opt out, or native to require native support. Native compilation and loading failures are reported without interpreter retry.
• --kernel-target=<sm_XX[a|f]> - Target GPU architecture (for example, sm_100, sm_100a, or sm_100f); defaults to the SM100 family target sm_100f
• --tile-size=<MxN> or <MxNxK> - Tile sizes for kernel tiling
• --print-ir-after-all - Print IR after each pass; ignored when --print-ir-tree-dir is set
• --print-ir-tree-dir=<dir> - Dump per-pass IR snapshots to a directory; takes precedence over --print-ir-after-all
• --reproducer-dir=<dir> - Dump one MLIR reproducer for each compiler pipeline into a unique directory derived from <dir>
• --timing - Show pass timing information

For tool command lines, TIR_OVERRIDE_OPTIONS can add, remove, reorder, or replace command-line options before they are parsed. For example, TIR_OVERRIDE_OPTIONS=+dump-tileir-bc=/tmp/out.bc adds --dump-tileir-bc=/tmp/out.bc when it is absent. This environment override is not consulted by the compiler libraries or APIs.

Execution:

• With no execution option, compile without building or running a reference graph
• --launch - Execute the compiled kernel on GPU
• --verify - Launch and verify results against the reference implementation
• --iterations=<N> - Number of execution iterations for benchmarking
• --seed=<N> - Random seed for test data
• --tolerance=<value> - Numerical tolerance for verification

Dynamic shapes:

• --dynamic-dims=<d0,d1,...> - Runtime values for ? dimensions
• --dynamic-strides=<s0,s1,...> - Runtime values for ? strides

Example MLIR Inputs

Static-shape matmul:

module {
  nv_tensor_ir.graph @matmul_f32_static(
    %a: tensor<8x8xf32>,
    %b: tensor<8x8xf32>) -> (
      tensor<8x8xf32>) {
    %c = "nv_tensor_ir.matmul"(%a, %b) :
      (tensor<8x8xf32>,
       tensor<8x8xf32>) ->
      tensor<8x8xf32>
    results %c : tensor<8x8xf32>
  }
}

Dynamic-shape pointwise add:

module {
  nv_tensor_ir.graph @add_dynamic_shape(
    %a: tensor<?x?xf32> {nv_tensor_ir.stride = "(?,1)"},
    %b: tensor<?x?xf32> {nv_tensor_ir.stride = "(?,1)"}
  ) -> (tensor<?x?xf32> {nv_tensor_ir.stride = "(?,1)"}) {
    %add = add %a, %b : tensor<?x?xf32>
    results %add : tensor<?x?xf32>
  }
}

More runnable MLIR examples are available under test/Integration/Compiler/.

Python Bindings

The OSS build packages TensorIR MLIR Python bindings under build/python_packages/nv_tensor_ir. These bindings can register the TensorIR dialect, build or parse TensorIR MLIR modules, and compile an MLIR module operation through the OSS TensorIR extension. Runtime launch accepts DLPack-compatible tensors; the example below uses PyTorch only as one tensor provider.

At launch, callers must provide CUDA-accessible DLPack tensors that match the compiled argument shapes, data types, and layouts. The runtime does not validate every mismatch before launching the kernel.

import torch

from nv_tensor_ir._mlir import ir
from nv_tensor_ir._mlir.dialects import nv_tensor_ir

mlir_text = """
module {
  nv_tensor_ir.graph @add_op(
    %a: tensor<8x8xf32> {nv_tensor_ir.stride = "(8,1)"},
    %b: tensor<8x8xf32> {nv_tensor_ir.stride = "(8,1)"}
  ) -> (tensor<8x8xf32> {nv_tensor_ir.stride = "(8,1)"}) {
    %out = add %a, %b : tensor<8x8xf32>
    results %out : tensor<8x8xf32>
  }
}
"""

with ir.Context() as ctx, ir.Location.unknown(ctx):
    nv_tensor_ir.register_dialect(ctx, load=True)
    module = ir.Module.parse(mlir_text)

    assert nv_tensor_ir.can_compile(module, tile_sizes=[8])
    with nv_tensor_ir.compile(module, tile_sizes=[8]) as artifact:
        assert artifact.get_bytecode()

        a = torch.randn((8, 8), device="cuda", dtype=torch.float32)
        b = torch.randn((8, 8), device="cuda", dtype=torch.float32)
        out = torch.empty_like(a)

        with artifact.load() as kernel:
            kernel.launch(a, b, out)

By default, TensorIR targets max(compatibility, 13.4). To request the compatibility bytecode target explicitly, pass a CompileOptions object:

options = nv_tensor_ir.CompileOptions()
options.tile_sizes = [8]
options.bytecode_version = nv_tensor_ir.BytecodeVersion.compatibility()
artifact = nv_tensor_ir.compile(module, options=options)

Compilation and deserialization return an Artifact, the Python counterpart of C++ TensorIRArtifact. They do not load host or device code. Calling artifact.load() creates an independently owned RuntimeKernel and initializes its host and device state before returning. Loading may compile device bytecode or fail if the artifact targets an incompatible GPU. Use a context manager or kernel.destroy() to unload it. Destroying the artifact does not invalidate kernels already loaded from it.

A compiled artifact can be serialized and written to disk, so a separate runtime process can load and launch it later without depending on the compiler at all:

# Host/compile-time process: compile once, persist the result.
with open("add_op.tirartifact", "wb") as f:
    f.write(artifact.serialize())

# Runtime process: deserialize the artifact, then explicitly load its code.
with open("add_op.tirartifact", "rb") as f:
    artifact = nv_tensor_ir.Artifact.deserialize(f.read())
with artifact.load() as kernel:
    kernel.launch(a, b, out)

Python DSL

The OSS build also packages a lightweight Python DSL under nv_tensor_ir.dsl. The DSL traces a Python kernel into TensorIR MLIR, compiles it through the same OSS compiler path, and runs it with DLPack-compatible tensors. Use tir.TensorSpec with tir.DataType to compile from metadata without allocating framework tensors.

import torch

from nv_tensor_ir import dsl as tir


USE_GELU = True


@tir.kernel
def fused_gemm_epilogue(a, b, bias, residual):
    x = a @ b + bias
    if USE_GELU:  # Static control flow, resolved while tracing.
        x = tir.gelu_approx_tanh(x)
    return x + residual


a = torch.randn((128, 64), device="cuda")
b = torch.randn((64, 128), device="cuda")
bias = torch.randn((128, 128), device="cuda")
residual = torch.randn_like(bias)
output = torch.empty_like(bias)

compiled = tir.compile(
    fused_gemm_epilogue,
    a,
    b,
    bias,
    residual,
    output=output,
    tile_sizes=(64, 64),
)
compiled.run(a, b, bias, residual, output=output)

To reuse compiled artifacts across equivalent DSL compilations, create a ArtifactCache and pass it to tir.compile. The cache is explicit: omitting artifact_cache disables cache lookup and insertion.

artifact_cache = tir.ArtifactCache()
compiled = tir.compile(
    fused_gemm_epilogue,
    a,
    b,
    bias,
    residual,
    output=output,
    tile_sizes=(64, 64),
    artifact_cache=artifact_cache,
)
# Alternatively, use tir.calculate_cache_key(compiled.module, ...) to get the
# cache key used for this compilation when managing the cache manually.

# ... reuse artifact_cache in later tir.compile calls ...

artifact_cache.flush()

The Python DSL provides a simple built-in profiling utility. Run the compiled kernel using the run_profile method, and it will print the kernel execution-time statistics after the run completes.

compiled.run_profile(
    a, b, bias, residual, output=output,
    warmup=2,
    iterations=50,
)

Will output:

[TENSOR_IR_PROFILE] kernel_fused_gemm_epilogue: calls=50, min=28.320 us, max=28.992 us, median=28.672 us, avg=28.667 us, outliers=0

Testing

After configuring the generated OSS tree, run the compiler sample MLIR, Python binding, and DSL smoke tests with:

cmake --build build --target check-tensor-ir

The target sets the test PYTHONPATH and points the tests at the build-tree tensor_ir-compiler. The smoke tests launch kernels and require a compatible NVIDIA GPU and driver.

License

TensorIR-owned source files are distributed under Apache-2.0 WITH LLVM-exception, as identified by the SPDX headers in each file. CUDA Tile is fetched from its upstream repository and distributed under its own license.