arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2405.05751cs.LGcs.AIcs.PL
收藏

Mirage:面向张量程序的多级超级优化器

Mirage: A Multi-Level Superoptimizer for Tensor Programs

Carnegie Mellon University(卡内基梅隆大学);Peking University(北京大学);Pennsylvania State University(宾夕法尼亚州立大学);Purdue University(Purdue 大学);Weizmann Institute of Science(魏茨曼科学研究院)

Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chunan Shi, Jianan Ji, Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, Zhihao Jia

AI总结:

本文提出张量程序多级超级优化器 Mirage,利用跨 GPU 三层级的 μGraphs 统一表示、抽象剪枝和概率等价验证,自动联合搜索代数、调度与自定义 kernel 优化,最高较现有方法提速 3.3 倍。

摘要:

我们提出 Mirage,这是首个面向张量程序的多级超级优化器。Mirage 的一个关键思想是 μGraphs,即在 GPU 计算层次结构的 kernel、thread block 和 thread 三个层级上对张量程序进行统一表示。μGraphs 使 Mirage 能够发现结合代数变换、调度变换以及生成新的自定义 kernel 的新颖优化。为了在庞大的搜索空间中导航,Mirage 引入了一种基于抽象的剪枝技术,可显著缩减搜索空间并提供一定的最优性保证。为确保优化后的 μGraph 与输入程序等价,Mirage 引入了具有强理论保证的概率等价性验证过程。我们的评估表明,即使对于被广泛使用且经过深度优化的 DNN,Mirage 也比现有方法最高快 3.3 倍。Mirage 已公开发布于 https://github.com/mirage-project/mirage。

英文摘要:

We introduce Mirage, the first multi-level superoptimizer for tensor programs. A key idea in Mirage is $μ$Graphs, a uniform representation of tensor programs at the kernel, thread block, and thread levels of the GPU compute hierarchy. $μ$Graphs enable Mirage to discover novel optimizations that combine algebraic transformations, schedule transformations, and generation of new custom kernels. To navigate the large search space, Mirage introduces a pruning technique based on abstraction that significantly reduces the search space and provides a certain optimality guarantee. To ensure that the optimized $μ$Graph is equivalent to the input program, Mirage introduces a probabilistic equivalence verification procedure with strong theoretical guarantees. Our evaluation shows that Mirage outperforms existing approaches by up to 3.3$\times$ even for DNNs that are widely used and heavily optimized. Mirage is publicly available at https://github.com/mirage-project/mirage.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑