跳到正文
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

查看 PDF HTML(实验性)

摘要:现代机器学习内核随着新的优化和硬件特性而快速演进。机器学习编译器抽象必须适应更丰富的内存层次、更多形式的线程协作以及日益异构的后端,同时保留实现峰值性能所需的控制能力。我们提出 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