TIRx: An Open Compiler Stack for Evolving Frontier ML Kernels
10 hours ago
- TIRx 是一个基于 Apache TVM 的开源、硬件原生 DSL 和编译器，用于 ML 内核，面向 GPU 和专用 AI 加速器。
- 它提供了一个低级边界，将编排、同步和后端内建函数保留在源代码中，供专家控制。
- 重复出现的图块原语暴露给编译器，以便在内核和后端之间实现可重用性。
- 新的硬件支持首先作为内建函数引入，然后随着模式稳定而提升为图块原语。
- 编程模型使用执行作用域、张量布局和图块原语调度来管理硬件原生操作。
- 张量布局是一种存储优先接口，将逻辑坐标映射到物理硬件轴，包含分片、副本和偏移组件。
- 编译器后端专注于局部解析图块原语，避免可能阻碍新内核模式的繁重优化过程。
- 在 NVIDIA B200 上的性能评估显示，TIRx 在密集 GEMM、块缩放低精度 GEMM 和 FlashAttention-4 上实现了具有竞争力的 TFLOPS。
- TIRx 通过支持可组合的图块任务来支持超内核，这些任务可以在 IR 级别拼接在一起。
- 通过可工具化的编译器栈和结构化的优化搜索空间，支持代理内核编程。
- 生态系统包括一个 PyPI 包、内核库、基准测试以及一个关于现代 GPU 编程的开放课程。