TIRx: An Open Compiler Stack for Evolving Frontier ML Kernels
10 hours ago
- TIRx is an open-source, hardware-native DSL and compiler for ML kernels built on Apache TVM, targeting GPUs and specialized AI accelerators.
- It provides a low-level boundary that keeps orchestration, synchronization, and backend intrinsics in source code for expert control.
- Recurring tile primitives are exposed to the compiler for reusability across kernels and backends.
- New hardware support is introduced first as intrinsics, then promoted to tile primitives as patterns stabilize.
- The programming model uses execution scope, tensor layout, and tile primitive dispatch to manage hardware-native operations.
- Tensor layout is a storage-first interface mapping logical coordinates to physical hardware axes with shard, replica, and offset components.
- The compiler backend focuses on resolving tile primitives locally, avoiding heavy optimization passes that could hinder new kernel patterns.
- Performance evaluations on NVIDIA B200 show TIRx achieves competitive TFLOPS for dense GEMM, block-scaled low-precision GEMM, and FlashAttention-4.
- TIRx supports megakernels by enabling composable tile tasks that can be stitched together at the IR level.
- Agentic kernel programming is supported through a toolable compiler stack and structured search spaces for optimization.
- The ecosystem includes a PyPI package, kernel library, benchmarks, and an open course on modern GPU programming.