SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

要約

スパース テンソルは、最新のディープ ラーニング ワークロードの重要なコンポーネントになりつつあります。
ただし、高性能のスパース オペレーターの開発は困難で退屈な場合があり、既存のベンダー ライブラリでは、新しいオペレーターからのエスカレートする要求を満たすことができません。
スパース tensor コンパイラは演算子の開発を簡素化しますが、単一のスパース形式ではハードウェア効率を最大化できず、シングルショット コンパイラは最新のハードウェアとシステムの進歩に追いつけないため、ディープ ラーニングのための効率的なスパース コンパイルは依然として困難です。
このホワイト ペーパーでは、これらの両方の課題に対処するための鍵は、構成可能な形式と構成可能な変換を活用することであると述べています。
SparseTIR は、ディープ ラーニング ワークロードに構成可能な形式と構成可能な変換を提供するスパース テンソル コンパイルの抽象化です。
SparseTIR は、パフォーマンス チューニングのために、これらの構成可能なコンポーネントに対して検索空間を構築します。
これらの改善により、SparseTIR は単一オペレーターの GPU でベンダー ライブラリと比較して一貫したパフォーマンス スピードアップを実現します。
また、SparseTIR は、GraphSAGE トレーニングで 1.08 ～ 1.52 倍、RGCN 推論で 4.20 ～ 40.18 倍、エンドツーエンドの GNN を高速化します。

要約(オリジナル)

Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and existing vendor libraries cannot satisfy the escalating demands from new operators. Sparse tensor compilers simplify the development of operators, but efficient sparse compilation for deep learning remains challenging because a single sparse format cannot maximize hardware efficiency, and single-shot compilers cannot keep up with latest hardware and system advances. In this paper, we observe that the key to addressing both these challenges is to leverage composable formats and composable transformations. We propose SparseTIR, a sparse tensor compilation abstraction that offers composable formats and composable transformations for deep learning workloads. SparseTIR constructs a search space over these composable components for performance tuning. With these improvements, SparseTIR obtains consistent performance speedups vs vendor libraries on GPUs for single operators: 1.20-2.34x for GNN operators, 1.05-2.98x for sparse attention operators, and 0.56-7.45x for sparse convolution operators. SparseTIR also accelerates end-to-end GNNs by 1.08-1.52x for GraphSAGE training, and 4.20-40.18x for RGCN inference.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google