NVIDIA/tensor-ir

NVIDIA/tensor-ir

Description: TensorIR is a lightweight NVIDIA-owned MLIR compiler frontend for expressing tensor computations and lowering them to NVIDIA CUDA Tile IR.

Language: C++

License: Apache-2.0

Stars: 10

Forks: 1

Open issues: 0

Created: 2026-08-04T15:22:14Z

Pushed: 2026-08-17T22:01:03Z

Default branch: main

Fork: no

Archived: no

README:

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.

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

| Option | Description | |--------|-------------| | -DTENSOR_IR_CUDA_TILE_SOURCE_URL= | Override the pinned CUDA Tile source archive | | -DTENSOR_IR_DOWNLOAD_LLVM=ON/OFF | Download the pinned LLVM source (default: ON for top-level builds) | | -DTENSOR_IR_LLVM_SOURCE_URL= | LLVM source archive URL; accepts a cached file:// URL | | -DMLIR_DIR= | MLIR package directory used when LLVM downloading is disabled | | -DTENSOR_IR_ENABLE_BINDINGS_PYTHON=ON/OFF | Build the TensorIR Python bindings (default: ON) | | -DTENSOR_IR_INCLUDE_TESTS=ON/OFF | Build test-only passes and configure the lit suite (default: ON) | | -DTENSOR_IR_CUDART_LINKAGE=DYNAMIC/STATIC | Dynamically load cudart or link CUDA::cudart_static (default: DYNAMIC) | | -DCMAKE_BUILD_TYPE=Release | Build type (Release/Debug/RelWithDebInfo) |

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 to install into a user-writable directory. Add /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 position exceeds the supplied list, the last value is reused. 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= - Write compiled CUDA Tile device code to ``

and metadata to .meta

Compilation:

• --dump-ir= - Dump lowered CUDA Tile dialect MLIR
• --dump-tileir-bc= - Dump Tile IR bytecode after compilation
• --bytecode-version= - Tile IR bytecode target: default,

current, or compatibility. The default target is max(compatibility, 13.3).

• --codegen-strategy= - TensorIR-to-CUDA-Tile lowering path:

layout-propagation (default) or affine-map

• --target-sm= - Target GPU architecture (for example,

sm_100, sm_100a, or sm_100f); defaults to the SM100 family target sm_100f

• --tile-size= or `` - 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= - Dump per-pass IR snapshots to a directory;

takes precedence over --print-ir-after-all

• --timing - Show pass timing information
• --uniform-signature - Pass all sizes and strides as kernel arguments,

including static sizes and strides

The IR...