TIRx: A Unified Tile Primitive ML Compiler Across Heterogeneous Backends

Bohan Hou; Hongyi Jin; Guanjie Wang; Shushi Hong; Jinqi Chen; Yaxing Cai; Lijie Yang; Zihao Ye

cs.DCJan 27, 2026

TIRx: A Unified Tile Primitive ML Compiler Across Heterogeneous Backends

Authors: Bohan Hou, Hongyi Jin, Guanjie Wang, Shushi Hong, Jinqi Chen, Yaxing Cai, Lijie Yang, Zihao Ye,

Organizations: Carnegie Mellon University · Shanghai Jiao Tong University · NVIDIA · Princeton University ·

Abstract

Modern ML kernels evolve rapidly with new optimizations and hardware features. ML compiler abstractions must accommodate richer memory hierarchies, more forms of thread cooperation, and increasingly heterogeneous backends while retaining the control required for peak performance. We present TIRx, an ML compiler for GPUs and AI accelerators that exposes these ingredients through explicit interfaces. TIRx introduces an explicit execution-scope abstraction that captures participating hardware teams at multiple granularities. It also uses tensor layouts to map logical tensors to heterogeneous hardware spaces through named axes and shard/replica/offset semantics. Finally, it introduces tile-primitive dispatch, which lowers high-level operations to efficient implementations according to scope, layout, and backend. This design preserves native-level control while automating recurring address calculation, operand preparation, and instruction-loop generation. Our evaluation shows vendor-level performance across the latest GPUs and AI accelerators, while native-code case studies show substantial reductions in low-level implementation code.

Figures & tables

Figure 1. TIRx programming model on an NVFP4 GEMM kernel. The kernel retains native-level control flow and scheduling. Tensor layouts describe logical-to-physical placement, execution scopes identify participating hardware teams, and tile primitives invoke reusable operations under these annotations. § 3 , § 4 , and § 5 describe the three interfaces. Overview of a \sysname{} NVFP4 GEMM kernel showing native-level kernel code together with tensor-layout, execution-scope, and tile-primitive interfaces.

Figure 2. Elements of TIRx Layout. An iter specifies a triple

(\text{extent},\text{stride},\text{axis})

and defines a linear, strided access on that axis. A list of iters forms the shard part

\mathbf{D}

, a set of iters forms the replica part

\mathbf{R}

, and

\mathbf{O}

is a fixed offset. Diagram showing the three components of a \sysname{} layout: shard D as a list of iters, replica R as a set of iters, and offset O as a fixed coordinate.

Figure 3. TIRx layout case studies. The figures in each column show the logical shape of the tensor and each cell shows the physical location(s) mapped from the corresponding logical element. Axes are color coded. Left. Mapping a logical $8\times 16$ tile to 4 GPU warps with 32 lanes and 2 registers each; 2 warps sharded and 2 warps replicated. Middle. Distributed sharding of a $64\times 128$ matrix across 4 GPUs; the top uses full sharding across 4 GPUs, the bottom uses a $2\times 2$ mesh with shards and replicas. Right. Native hardware memories; the top depicts an AI accelerator 2D-partitioned scratchpad SRAM, the bottom shows NVIDIA Blackwell 2D tensor memory. Three layout examples: a GPU tensor core tile mapped to warps, lanes, and registers; distributed sharding across a 4-GPU mesh; and native hardware memory layouts for an AI accelerator scratchpad and Blackwell tensor memory.

Figure 4. TIRx Layout Python API. If some stride is not paired with an axis, the axis m is used by default. Code snippet showing the \sysname{} Layout Python API for defining layouts with shard, replica, and offset components.

Figure 5. Example distributed tensor signature for reduce scatter. Diagram showing a distributed tensor signature for a reduce-scatter operation with input and output sharding specifications.

Axis	Backend	Extent	Use
m	all	buffer-defined	linear memory coordinate
laneid	NVIDIA	32	lane within a warp
warpid	NVIDIA	launch-defined	warp within a CTA
wid_in_wg	NVIDIA	4	warp within a warpgroup
tid_in_wg	NVIDIA	128	thread within a warpgroup
cbx/cby/cbz	NVIDIA	launch-defined	CTA coordinate in a cluster

Table 1. Representative hardware axes used by TIRx layouts.

Figure 6. Each operation has a distinct team kind (cluster, warpgroup, thread) and involved tensor layout, while the active thread set is independently controlled by filtering. Diagram showing the \sysname{} DSL execution scope slice API with thread filtering and team switching constructs.

Figure 7. Layout grouping and permutation expose a larger legal box within the selected TMA family, reducing the number of issued instructions. Diagram showing how grouping and permutation help discover a larger TMA box.

Primitive	Specialized family	Fallback
copy	ldmatrix/stmatrix , NKI copy	vector/scalar load/store
copy_async	TMA, tcgen05.cp/ld/st , cp.async	none
gemm	mma.sync.aligned , nki.matmul	scalar FMA
gemm_async	tcgen05.mma	none
reduce	warp shuffle, NKI tensor reduce	vector/scalar arithmetic
pointwise	NKI/vector arithmetic	scalar arithmetic

Table 2. Representative tile primitives, specialized instruction families, and registered fallbacks.

Figure 8. FP16 GEMM and FP8, NVFP4 block-scaled GEMM throughput (TFLOP/s) at batch size 8192 across different weight shapes. The dashed line marks 50% of device peak (for FP16, FP8 and NVFP4, respectively). TIRx is within 4% of the vendor-tuned cuBLAS on average for FP16 GEMMs while significantly outperforming Triton. Each

x

-axis label denotes the weight shape

N

K

(with

M{=}8192

). Higher is better. Bar charts comparing FP16, FP8 block-scaled, and NVFP4 GEMM throughput for \sysname{} against cuBLAS, Triton, DeepGEMM, and cuBLASLt across different weight shapes.

Figure 9. Forward multi-head attention (MHA) throughput on B200 (head dimension 128) for non-causal and causal settings across sequence lengths from 1K to 32K. Higher is better. Bar charts comparing forward MHA throughput on B200 for \sysname{}, FlashAttention-4, cuDNN, and Triton/Gluon across sequence lengths.

Figure 10. FP16 GEMM + Reduce-Scatter latency (ms) evaluated across different weight shapes. Lower is better. Bar chart comparing FP16 GEMM plus Reduce-Scatter latency for \sysname{}, cuBLAS plus NCCL, and Triton-distributed across weight shapes.

Figure 11. FP16 GEMM and multi-head attention (MHA) test results. FP16 GEMM is tested on square shapes. MHA is tested with varying input lengths with no causal mask. Bar chart comparing FP16 GEMM and MHA performance on Trainium1 for \sysname{}, NKI, and Neuron compiler.

GPU Kernel Generation

GPU (graphics processing unit) kernel generation automates writing and optimizing device code for faster computation across accelerators.

See all 70 papers

RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust
Sep 29, 2026Eugene Hauptmann, Nataliya KosmynaEfficient Neural Network Inference
Massachusetts Institute of Technology Cambridge, MA, USA
Axon: A Synthesizing Superoptimizer for Tensor Programs
Jun 24, 2026Akash Kothari, Shaowei Zhu, Daniel Kroening +1AI Accelerator Inference High-Performance Computing
University of Illinois at Urbana-Champaign, USA · Amazon Web Services, USA
AI as a Compiler: Compiling Triton kernels without the Triton compiler
Sep 29, 2026François Costa, Charly Castes, Thomas Bourgeat +1GPU Kernel Optimization Large Language Model-Guided Optimization
Stanford University · EPFL

Sep 29, 2026J/K move · Enter open · S save

RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust

Eugene Hauptmann, Nataliya Kosmyna

Massachusetts Institute of Technology Cambridge, MA, USA

Efficient Neural Network Inference

Open paper AskPDF arXiv

Production machine learning (ML) stacks often split graph compilation and kernel execution across different layers and languages, making backend behavior, deployment guarantees, and performance fallbacks hard to reason about end-to-end. RLX addresses this gap with a single Rust codebase that combines compiler and runtime roles around one primitive-level, three-level intermediate representation (IR), plus a transparent dispatch contract that resolves each operator to native, common-IR, or rewritten lowering and fails compilation when legalization is not possible. The same IR targets fourteen runtime devices (cpu, metal, mlx, ane, cuda, rocm, oneapi, tpu, hexagon, gpu, vulkan, opengl, directx, webgpu) and two specialty codegen paths (Cortex-M INT8 and FPGA), ingests safetensors, GGUF, ONNX, and rten formats, supports F16/BF16/F64/C64 and quantized INT4/INT8 flows with AMP/PTQ/QAT, and scales via tensor-/pipeline-parallel collectives over TCP and RDMA transports. Beyond neural workloads, RLX also extends to scientific/physics-style domains through sparse and dense linear algebra extensions (e.g., CSR LU/CG/matvec and LAPACK- backed factorizations) and 3D Gaussian splatting operators. We evaluate RLX against PyTorch, TensorFlow, JAX, candle, burn, tch, rten, MLX, CoreML, IREE, Glow, TensorRT, and tinygrad under identical input generation and p50 measurement methodology on one host. On all-MiniLM-L6-v2, RLX-Metal is fastest at every batch (e.g., 16.6 ms at batch 32 vs. PyTorch-MPS 26.7 ms). In the MNIST training table, RLX also has the top-throughput entry (graph-fused MLP: 946,487 img/s), above NumPy+BLAS (787,349 img/s), while retaining 100% top-1 parity on reference checks (e.g., Qwen3).

Abstract

GPU Kernel Generation

Explore similar work

RLX: A Unified Multi-Backend Tensor Compiler and Distributed Runtime in Rust

Axon: A Synthesizing Superoptimizer for Tensor Programs

AI as a Compiler: Compiling Triton kernels without the Triton compiler