タイル、ハロー、部分和

記憶に収まるようテンソルを分割する

1タイルの出力高さTh、幅Tw、入力チャネル数Tc、出力チャネル数Tmを選びます。入力ハローは出力タイルの外まで広がります。ストライドS、ダイレーションD、カーネルRなら必要入力高さは(Th−1)S+(R−1)D+1で、幅にも対応する水平式を使います。境界パディングは取得せず生成できる場合があります。

入力チャネルを分割すると、全チャネルタイルにわたって各出力部分和を保持する必要があります。各タイルを早く丸めると演算が変わります。最初のチャネルタイルで1回だけバイアスを加え、その和を続けます。出力チャネルタイルを変えると重みとバイアスも変わります。

記憶予算の計算例

3×3、ストライド1、ダイレーション1でTh=Tw=2,Tc=2,Tm=2とします。ハローは4×4でint8入力32バイトです。重みは2×2×3×3=36バイトです。8個の32ビット部分和は32バイト、2個の32ビットバイアスは8で、計108バイトです。スケール、アドレス、有効性、制御記憶は除外し、別の許容量が必要です。

入力と重みの二重バッファは2×(32+36)+32+8=176バイトで、128バイト予算を超えます。Tcを1に減らすと入力16、重み18バイトです。二重バッファ記憶は2×34+32+8=108バイトです。2つの入力チャネルタイルを通して部分和を保持します。バッファが収まっても転送と計算の重ね合わせを証明しません。スケジュールを別に確認します。

ストライド2では同じ2×2出力タイルの入力範囲が5×5へ広がります。隣接タイルはハロー値を共有し得ますが、独立再ロードでは通信が重複します。下限となる重複なしテンソルバイト数と、実際のタイル転送を区別します。

M=5、Tm=2なら最後のタイルは有効出力チャネル1つです。余分な重みと出力アクセスをマスクします。最終空間タイルと入力チャネルタイルにも同等のマスクを適用します。NVIDIA CUTLASS資料は全展開行列なしでタイルを生成する実際のimplicit-GEMM方式を説明します。ここではそのCUDA実装を実行せず、記憶と部分和を計算します。