A Multi-Level Superoptimizer for Tensor Programs

要約

テンソル プログラム用の初のマルチレベル スーパーオプティマイザーである Mirage を紹介します。
Mirage の重要なアイデアは、GPU 計算階層のカーネル、スレッド ブロック、スレッド レベルでのテンソル プログラムの統一表現である $\mu$Graphs です。
$\mu$Graphs を使用すると、Mirage は代数変換、スケジュール変換、新しいカスタム カーネルの生成を組み合わせた新しい最適化を発見できます。
大規模な検索スペースをナビゲートするために、Mirage は抽象化に基づいたプルーニング手法を導入し、検索スペースを大幅に削減し、一定の最適性を保証します。
最適化された $\mu$Graph が入力プログラムと同等であることを保証するために、Mirage は強力な理論的保証を備えた確率的同等性検証手順を導入します。
私たちの評価では、広く使用され高度に最適化されている DNN であっても、Mirage が既存のアプローチよりも最大 3.5$\times$ 優れていることがわかりました。
Mirage は https://github.com/mirage-project/mirage で公開されています。

要約(オリジナル)

We introduce Mirage, the first multi-level superoptimizer for tensor programs. A key idea in Mirage is $\mu$Graphs, a uniform representation of tensor programs at the kernel, thread block, and thread levels of the GPU compute hierarchy. $\mu$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 $\mu$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.5$\times$ even for DNNs that are widely used and heavily optimized. Mirage is publicly available at https://github.com/mirage-project/mirage.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google