[論文レビュー] Automatic Kernel Generation for Volta Tensor Cores

Somashekaracharya G. Bhaskaracharya, Julien Demouth+1 more|arXiv (Cornell University)|Jun 22, 2020

Parallel Computing and Optimization Techniques参考文献 19被引用数 13

ひとこと要約

本稿では、NVIDIA Voltaテンソルコアを対象とした、行列積に畳み込み演算（ReLU やバイアス加算など）を融合させた高度に最適化されたCUDAカーネルを自動生成する多面体コンパイラー手法を提示する。mma.sync.m8n8k4 PTX 指令とマクロ-MMA 組み合わせを活用することで、手動で最適化されたライブラリ実装よりも最大2.55倍の高速化を達成し、テンソルコアワークロードにおける自動カーネル生成の有効性を示している。

ABSTRACT

A commonly occurring computation idiom in neural networks is to perform some pointwise operations on the result of a matrix multiplication. Such a sequence of operations is typically represented as a computation graph in deep learning compilers. When compiling to a GPU target, these computations can be individually mapped to manually tuned implementations provided by libraries such as cuBLAS and cuDNN. These libraries also provide off-the-shelf support for targeting tensor cores in NVIDIA GPUs, which can lead to huge performance boosts through their specialized support for mixed-precision matrix math. Alternatively, tensor cores can be programmed directly using CUDA APIs or inline assembly instructions, which opens up the possibility of generating efficient CUDA kernels automatically for such computations. Automatic kernel generation is particularly crucial when it is beneficial to generate efficient code for an entire computation graph by fusing several operations into a single device function instead of invoking a separate kernel for each of them. Polyhedral compilation techniques provide a systematic approach for the analysis and transformation of a sequence of affine loop-nests. In this paper, we describe a polyhedral approach to generate efficient CUDA kernels for matrix multiplication using inline assembly instructions for programming tensor cores on NVIDIA Volta GPUs. Furthermore, we build on this approach to generate fused kernels for computation sequences involving matrix multiplication and pointwise operations such as bias addition, ReLU activation etc. Experimental evaluation of these techniques show that automatically generated kernels can provide significantly better performance than manually tuned library implementations, with speedups ranging up to 2.55X.

研究の動機と目的

• 各演算ごとに別個のカーネルを実行する深層学習計算グラフ処理において、繰り返し発生するグローバルメモリアクセスによる性能ボトルネックを解消すること。
• 混合精度行列積と畳み込み演算を融合した、Voltaテンソルコアを効果的に活用する高効率なCUDAカーネルの自動生成を可能にすること。
• 多面体コンパイラー技術を拡張し、細粒度なテンソルコア使用制御を可能にする mma.sync.m8n8k4 PTX 指令を直接ターゲットにすること。
• 自動生成されたカーネルが、cuBLAS や cuDNN などの手動で最適化されたライブラリ実装を上回ることを、融合計算シナリオにおいて実証すること。
• テンソルコア加速を前提とした深層学習ワークロード向けに、高パフォーマンスカーネルを体系的かつ再利用可能に生成するフレームワークを提供すること。

提案手法

• 多面体フレームワークを用いて、行列積および畳み込み演算を表すアフィンループネストを解析・変換する。
• 8×8×8 の行列断片をスレッドクワッド上で処理する mma.sync.m8n8k4 PTX 指令を直接使用して、Voltaテンソルコアをプログラミングするカーネルを生成する。
• 低レベルの mma.sync.m8n8k4 操作を組み合わせて、16×16×8 および 32×32×8 の高レベルなマクロ-MMA 操作を構築し、メモリ帯域幅の利用効率を向上させる。
• ReLU、シグモイド、バイアス加算などの上流（プロローグ）および下流（エピローグ）の畳み込み演算と行列積を融合可能にする。
• Diesel DSL コンパイラーと統合し、テンソルコアカーネルのタイルサイズとレジスタ使用量を最適化するカスタムスケジューリングを採用する。
• 高レベルな計算グラフを最適化されたCUDAカーネルにマッピングする体系的なコード生成パイプラインを採用し、最小限の手動チューニングで実現する。

実験結果

リサーチクエスチョン

• RQ1 多面体コンパイラー技術は、Voltaテンソルコア向けに高パフォーマンスCUDAカーネルを効果的に生成するために適用可能か？
• RQ2 自動カーネル生成は、行列積と畳み込み演算を融合したシナリオにおいて、cuBLAS や cuDNN などの手動最適化ライブラリをどの程度上回るか？
• RQ3 mma.sync.m8n8k4 PTX 指令の使用は、融合カーネルにおける高い算術強度とメモリ帯域幅効率を達成するためにどの程度有効か？
• RQ4 さまざまなワークロード条件下で、異なるタイルサイズや統合戦略の性能的トレードオフは何か？
• RQ5 16×16×8 および 32×32×8 のマクロ-MMA 組み合わせを、多面体解析を用いて自動的に生成・最適化できるか？

主な発見

• 提案手法は、特に大きな問題サイズにおいて、cuBLAS や cuDNN に対して最大2.55倍の高速化を達成した。
• 単一の行列積カーネルにおいても、cuBLAS に対して1.46倍、cuDNN に対して1.35倍の高速化を達成し、融合なしでも有効性が示された。
• 複数の行列積を含む融合カーネルでは、レジスタ圧力が高まり、タイルサイズを小さくする必要が生じたが、性能向上は依然として得られた。
• Halide や TVM といった既存のフレームワークに比べ、直接 mma.sync.m8n8k4 を使用し、マクロ-MMA 組み合わせを実装することで、特にテンソルコアカーネル生成において優れた性能を発揮した。
• 広範な問題サイズにわたり顕著なパフォーマンス向上を達成し、特に行列積と ReLU やバイアス加算を含む融合計算グラフで最大の利得が得られた。
• 多面体コンパイラー技術が低レベルのテンソルコア命令をターゲットに拡張可能であり、高度に最適化されたカーネルの自動生成を可能にすることを実証した。