シストリックアレイ — TPUの心臓部を1から

掛け算器を並べるだけでは速くならないのは、データを運ぶ手間が演算より重いからです。TPUの中心にある格子＝シストリックアレイが、行列積の三重ループを縦・横・時間にどう割り当てているのかを、2×2の手計算からサイクル単位のコード、バッチと次元の実務まで前提知識ゼロで解きほぐします。

In-Datacenter Performance Analysis of a Tensor Processing Unit

井戸まで走るか、手渡しでつなぐか

100人に「井戸から水を汲んで畑に撒け」と命じたとします。各自がバケツを持って往復するより、井戸から畑まで一列に並んで手渡しするほうが速い。誰も走らず、各自は隣に渡すだけ。それでも水は途切れずに届きます。

計算機の中でも事情は同じです。掛け算器はシリコン上ではとても小さく、いくらでも並べられます。ところがその1個が1回の積和をするには、入力を2つ読み、結果を1つ書かなければならない。相手がレジスタファイルやSRAMだと、演算そのものより出し入れのほうが時間も電力も食います。演算器を1000個に増やしても、そのぶんのデータを毎サイクル運べなければ残りは手を止めて待つだけ。この「運ぶ側が足を引っ張る」構図はメモリの壁で扱った通りで、AIチップ設計の出発点になっています。

シストリックアレイ（systolic array）は、これにバケツリレーで答えた設計です。演算器を格子状に敷き詰め、各演算器は隣接する演算器としかデータをやり取りしない。一度チップに入ったデータはメモリへ戻ることなく格子を端から端まで流れ、通り道の全演算器に仕事をさせてから出ていきます。systolic とは心臓の収縮のこと。データが鼓動のように一拍ずつ格子を進む様子から、1970年代末にH.T. Kung と Charles Leiserson がこう名付けました。TPUは、この40年前のアイデアを行列積という一点に絞って復活させたチップです。

格子の1マスがやっていること

格子の1マス——PE（Processing Element、処理要素）——の中身は驚くほど単純です。掛け算器が1個、足し算器が1個、値を1拍だけ抱えておくレジスタが数個。それだけです。

入り口は3つあります。左の隣から流れてくる活性化 $a$、上の隣から降りてくる部分和 $p$、そして自分の中に置きっぱなしの重み $w$。1サイクルごとに、PEは次の2つを同時に行います。

つまり言い換えれば、「上から来た数に、左から来た数と自分の重みの積を足して下へ落とす。左から来た数はそのまま右へ受け流す」だけです。$p_{\text{in}}$ は上のPEが出した途中経過、$a$ は左から届いた入力の値、$w$ は自分が抱えている重みを指します。

嬉しい点が3つあります。PEは隣としか配線されていないのでチップを横断する長い配線が要らず、高いクロックで回せる。重み $w$ は一度置いたら動かないので、重みの読み出しが消える。部分和が手渡しで降りていくので、途中経過をメモリへ書き戻さずに済む。バケツリレーで誰も井戸へ走らないのと同じ理屈です。

2つ目の「置いたら動かない」は文字通りの意味です。計算を始める前に重みを1回だけ格子へ流し込んでPEのレジスタに焼き付け、そのあとは入力を何千本流そうが重みには一切触れません。ニューラルネットの推論では同じ重みを大量の入力に当て続けるので、この前提が素直に成立します。逆に言えば、重みを置き替える頻度が高い使い方ではこの利点が消えます。

PEが毎拍刻んでいるこの「掛けて足す」は、内積の1項そのものです。

行列積を、格子と時間に割り当てる

本題です。行列積 $XW$ をこの格子でどう計算するのか。$X$ は入力を縦に $T$ 本並べた $T \times N$ の行列、$W$ は $N \times M$ の重み行列とします。

つまり、出力の $t$ 行 $j$ 列にある1個の数は、入力ベクトル1本と重み行列の1列を先頭から順に1組ずつ掛け、それを全部足し上げた合計だ、ということです。$t$ は何本目の入力ベクトルか、$j$ は出力の何番目の要素か、$i$ は足し合わせる中身の番号です。プログラムで書けば三重ループになります。

シストリックアレイの核心は、この三重ループの3本を「縦」「横」「時間」に割り振ったことにあります。$i$（足し算の中身）を格子の縦へ、$j$（出力の要素番号）を格子の横へ、$t$（何本目の入力か）を時間へ。