多级张量代数超级优化器

成天评科技文化 2024-10-03 21:04:33

Mirage:为 PyTorch 程序自动生成快速 GPU 内核

Mirage 是一款通过超级优化技术自动为 PyTorch 程序生成快速 GPU 内核的工具。例如,要获得用于注意力的快速 GPU 内核,用户只需编写几行 Python 代码来描述注意力的计算。

对于给定的 PyTorch 程序,Mirage 会自动搜索功能上与输入程序等效的潜在 GPU 内核空间,并发现高度优化的内核候选者。这种方法使 Mirage 能够找到性能优于现有专家设计的内核的新自定义内核。

Mirage 是一个张量代数超级优化器,可自动发现针对 DNN 的高度优化的张量程序。Mirage 可自动识别和验证复杂的优化,其中许多优化需要在 GPU 计算层次结构的内核、线程块和线程级别进行联合优化。对于输入 DNN,Mirage 会搜索功能上与给定 DNN 等效的潜在张量程序空间,以发现高度优化的候选程序。这种方法使 Mirage 能够找到性能优于现有专家设计的内核的新自定义内核。

论文链接https://arxiv.org/abs/2405.05751

项目链接https://github.com/mirage-project/mirage

0 阅读:3