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·· 10 小时前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 接收。
正文 · 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
摘要:现代机器学习内核随着新的优化和硬件特性而快速演进。机器学习编译器抽象必须适应更丰富的内存层次、更多形式的线程协作以及日益异构的后端,同时保留实现峰值性能所需的控制能力。我们提出 TIRx,一种面向 GPU 和 AI 加速器的机器学习编译器,通过显式接口暴露这些要素。TIRx 引入显式的执行作用域抽象,可在多个粒度上刻画参与计算的硬件团队。它还使用张量布局,通过命名轴和分片/副本/偏移语义将逻辑张量映射到异构硬件空间。最后,它引入 tile 原语调度,根据作用域、布局和后端将高层操作下沉为高效实现。该设计在保留原生级控制的同时,自动完成重复出现的地址计算、操作数准备和指令循环生成。我们的评估显示,在最新 GPU 和 AI 加速器上达到厂商级性能,而原生代码案例研究表明底层实现代码大幅减少。
| 评论: | 已被 ASPLOS 2027 接收 |
| 主题: | 分布式、并行与集群计算(cs.DC);人工智能(cs.AI);机器学习(cs.LG);编程语言(cs.PL) |
| 引用格式: | arXiv:2601.19092 [cs.DC] |
| (或本版本使用 arXiv:2601.19092v3 [cs.DC]) | |
| https://doi.org/10.48550/arXiv.2601.19092 由 DataCite 颁发 arXiv DOI |
提交历史
来自:Bohan Hou [查看邮箱]
[v1]
2026年1月27日 星期二 01:57:29 UTC(528 KB)
[v2]
2026年1月28日 星期三 22:37:40 UTC(508 KB)
[v3]
2026年10月30日 星期五 05:42:30 UTC(494 KB)
来源:arXiv cs.AI · arxiv.org