arXiv cs.AI· Bohan Hou, Hongyi Jin, Guanjie Wang, Shushi Hong, Jinqi Chen, Yaxing Cai, Lijie Yang, Zihao Ye, Yaoyao Ding, Ruihang Lai, Todd Mowry, Tianqi Chen·· 11 小时前AI 评分34
TIRx：跨异构后端的统一 Tile 原语 ML 编译器
TIRx: A Unified Tile Primitive ML Compiler Across Heterogeneous Backends
AI 导读
TIRx 是面向 GPU 与 AI 加速器的 ML 编译器，通过显式执行作用域抽象、张量布局和 tile-primitive dispatch，将高层操作按作用域、布局与后端下沉为高效实现。它在最新 GPU 和 AI 加速器上达到厂商级性能，原生代码案例显示底层实现代码量大幅减少。该工作已被 ASPLOS 2027 接收。
正文 · 原文
Authors:Bohan Hou, Hongyi Jin, Guanjie Wang, Shushi Hong, Jinqi Chen, Yaxing Cai, Lijie Yang, Zihao Ye, Yaoyao Ding, Ruihang Lai, Todd Mowry, Tianqi Chen
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.
| Comments: | Accepted to ASPLOS 2027 |
| Subjects: | Distributed, Parallel, and Cluster Computing (cs.DC); Artificial Intelligence (cs.AI); Machine Learning (cs.LG); Programming Languages (cs.PL) |
| Cite as: | arXiv:2601.19092 [cs.DC] |
| (or arXiv:2601.19092v3 [cs.DC] for this version) | |
| https://doi.org/10.48550/arXiv.2601.19092 arXiv-issued DOI via DataCite |
Submission history
From: Bohan Hou [view email]
[v1]
Tue, 27 Jan 2026 01:57:29 UTC (528 KB)
[v2]
Wed, 28 Jan 2026 22:37:40 UTC (508 KB)
[v3]
Fri, 30 Oct 2026 05:42:30 UTC (494 KB)
来源：arXiv cs.AI · arxiv.org