Skip to content
Published on

シストリックアレイとデータフローアーキテクチャ — TPUの心臓原理

🇰🇷KO🇺🇸EN🇯🇵JA
Split View필사 모드
Authors

はじめに

今日の巨大な言語モデルを学習し推論する作業は、結局のところ膨大な量の行列積（matrix multiplication）に帰着します。Transformerのアテンションも、MLP層も、畳み込み（convolution）も、内部的には積和（multiply-accumulate, MAC）演算の巨大な山です。

ところがこのMAC演算そのものは、実はそれほど高価ではありません。本当に高価なのは、データをメモリから演算器へ運ぶ作業です。一回の乗算に要するエネルギーよりも、その乗算に必要な数値をDRAMから読み出すエネルギーが数百倍大きいことも珍しくありません。これがいわゆるメモリウォール（memory wall）問題です。

シストリックアレイは、まさにこの問題に正面から挑むために考案された構造です。データを一度読み込んだら、そのデータがチップ内部の演算器の格子（grid）に沿って流れながら何度も再利用されるようにします。GoogleのTPU（Tensor Processing Unit）がこの原理の最も有名な実装であり、NVIDIA GPUのテンソルコアも似た発想を共有しています。

本記事では、シストリックアレイとは何でどう動くのか、weight-stationaryやoutput-stationaryといったデータフロー戦略が何を決めるのか、そしてコンパイラが大きな行列をどのように小さな格子へマッピングするのかを、段階的に見ていきます。

なぜ行列積がAIの中核演算なのか

まず基準点を定めましょう。サイズがM×Kの行列Aと、K×Nの行列Bを掛けてM×Nの行列Cを得る演算を考えます。

C[m][n] = sum over k of  A[m][k] * B[k][n]

この一行に必要な積和の回数はM×N×Kです。たとえば1024×1024の行列を二つ掛けると、約10億（10^9）に近いMAC演算が起こります。巨大モデルでは、こうした行列積が毎秒数千回繰り返されます。

核心的な洞察は次の通りです。同じデータが何度も使われるという点です。

  • Aの一行 A[m][:] は、Cの一行全体（N個の要素）を計算するのに繰り返し使われます。
  • Bの一列 B[:][n] は、Cの一列全体（M個の要素）を計算するのに繰り返し使われます。

素朴に毎回メモリから読み直すと、データ移動量が爆発します。だから一度読んだデータをチップ内に留め、できる限り再利用しなければなりません。シストリックアレイは、この再利用をハードウェア構造そのものに刻み込んだ設計です。

シストリックアレイとは何か

この名前は、1978年にH. T. KungとCharles Leisersonが提案した概念に由来します。systolicは、心臓が規則的に拍動して血液を押し出す（systole）様子から取った言葉です。データがクロックごとに一マスずつ、まるで心拍のように規則的に格子を流れていくからです。

構造は単純です。小さな演算器（processing element, PE）を二次元の格子に密に配置します。各PEは次のことだけを行います。

  1. 上（または左）から入ってきた値を受け取る。
  2. 自分が持つ値と掛ける。
  3. アキュムレータに足す。
  4. 受け取った値を下（または右）の隣のPEへそのまま渡す。

PEは外部メモリに直接アクセスしません。隣のPEとだけ会話します。データは格子の縁から入り、一マスずつ前進して内部のすべてのPEを通ります。一度入ったデータが格子全体を横切る間に、無数のMACで再利用されるのです。

        B の重みが上から下へ流れる
        b00   b01   b02
         |     |     |
   a0 → [PE]→[PE]→[PE]→   (A の入力が左から右へ流れる)
         |     |     |
   a1 → [PE]→[PE]→[PE]→
         |     |     |
   a2 → [PE]→[PE]→[PE]→
         |     |     |
        部分和が下へ累積して抜けていく

各PE内部を擬似コードで表すと次の通りです。

PEがクロックサイクルごとに行うこと:
  in_left    = 左の隣から受け取った活性値
  in_top     = 上の隣から受け取った部分和
  product    = in_left * weight_held_here
  out_bottom = in_top + product
  out_right  = in_left          (活性値を右へそのまま伝える)

ここで決定的なのは、weight_held_here がPEの中に留まっているという点です。重みを一度ロードしておけば、流れ込むすべての活性値に対して同じ重みが再利用されます。これが次節で扱うweight-stationary戦略です。

データが流れながら掛けられる — 動作の追跡

概念を手で追ってみましょう。2×2の行列を二つ掛けます。

A = | a00  a01 |      B = | b00  b01 |
    | a10  a11 |          | b10  b11 |

求める結果:
C[0][0] = a00*b00 + a01*b10
C[0][1] = a00*b01 + a01*b11
C[1][0] = a10*b00 + a11*b10
C[1][1] = a10*b01 + a11*b11

2×2のPE格子を使い、Bの重みを各PEに固定ロードします。

PE位置ごとにロードされた重み:
  [PE00=b00] [PE01=b01]
  [PE10=b10] [PE11=b11]

ここでAの行を左の縁から流し込みます。重要なのは、データを斜め（skewed）に入れる点です。行ごとに一クロックずつ遅延させると累積のタイミングが合います。

時間 →
入力ストリーム (左の縁):
  行 0 (PE00, PE01 へ):  a00, a01
  行 1 (PE10, PE11 へ):  a10, a11  (一サイクル遅れて進入)

PE00をクロックごとに追うと次の通りです。

サイクル 1: a00 到着 → 累積 = a00*b00
サイクル 2: a01 到着 → 累積 = a00*b00 + a01*b10  ← C[0][0] 完成

PE00の上を流れていた部分和が下方向へ合わさり、片側の縁から抜け出すとき、私たちが求めたCの各要素が完成します。核心は、a00やa01といった値をメモリから一度だけ読み、b00やb10といった重みは一切読み直さずPE内に常駐させた点です。

この小さな例を256×256の格子へ拡張すると、一度に256×256 = 65,536個のMACをクロックごとに同時に実行します。TPUが行列積で圧倒的なスループットを出す理由がまさにこれです。

TPUのシストリック設計

Googleが2016年の論文で公開した第1世代TPUの心臓は、256×256のシストリック行列積ユニット（Matrix Multiply Unit, MXU）でした。65,536個の8ビット整数MACユニットが格子をなし、単一命令で巨大な行列タイルを処理しました。

TPU設計の核心的特徴を整理すると次の通りです。

  • 巨大な単一MXU: 多数の小さなコアではなく、一つの大きな格子に演算を集中させます。制御回路のオーバーヘッドが小さく、チップ面積の大部分を実際の演算器に使えます。
  • オンチップのデータ再利用: 重みと活性値が格子内を流れて再利用されるため、DRAMアクセス回数が劇的に減ります。
  • 決定論的な実行: キャッシュミスや分岐予測のような非決定的要素がほとんどなく、性能を正確に予測しコンパイル段階でスケジューリングできます。

世代を経てTPUは進化しました。2026年現在、Googleは第6世代Trillium（TPU v6）を前世代比でピーク性能約4.7倍まで引き上げ、推論に特化した第7世代Ironwoodと併用しています。しかし、格子状のシストリック演算ユニットにデータを流して再利用するという根本原理は、第1世代から変わっていません。

Weight-Stationary と Output-Stationary

シストリックアレイを設計する際に最も重要な決定は、何をPE内に固定（stationary）するかです。三つのデータ（重み、活性値、部分和）のうち何を留めるかで、データフロー戦略の名前が分かれます。

Weight-Stationary（重み固定）

重みをPEにロードし、活性値を流し込み、部分和を格子の外へ出します。

PE 内に留まる:   重み (W)
流れ込む:        活性値 (A)
流れ出る:        部分和 (psum)
  • 利点: 同じ重みを複数の入力バッチに再利用でき、バッチが大きいときに重み再利用の効率が最高になります。
  • 適する場合: 推論のように重みは固定で入力だけが変わり続けるワークロード。

Output-Stationary（出力固定）

各PEが出力要素一つのアキュムレータを最後まで保持します。重みと活性値の両方が流れ込みます。

PE 内に留まる:   部分和 (psum) — 出力要素
流れ込む:        重み (W) と活性値 (A)
流れ出る:        完成した出力 (最後に一度)
  • 利点: 部分和をあちこち動かさないので累積精度の損失がなく、psum移動のエネルギーがゼロに近づきます。
  • 適する場合: 累積の深さ（K）が非常に長く、部分和の移動コストが負担になる場合。

Input-Stationary（入力固定）

活性値をPEに固定し、重みを流します。同じ入力が複数の出力チャネルに再利用される畳み込みパターンなどで有用です。

三つの戦略を一覧で比較すると次の通りです。

戦略固定データ流れるデータ強み弱い状況
Weight-stationary重み活性値, 部分和重み再利用を最大化, 推論向き部分和移動のエネルギーが発生
Output-stationary部分和重み, 活性値累積精度, psum移動が最小重みを毎回流し直す必要
Input-stationary活性値重み, 部分和入力再利用の大きい畳み込みに有利出力チャネルが多いと負担

実際のアクセラレータは、ワークロードに応じてこれらの戦略を選択的に、あるいは混在させて使います。どの戦略もすべての場合に最善ではなく、行列の形（M, N, Kの比率）とバッチサイズが最適な選択を左右します。

データ再利用とエネルギー

なぜデータフロー戦略がこれほど重要なのでしょうか。答えはエネルギーにあります。さまざまな測定で共通して現れる大まかな傾向は、一回の演算より一回のデータ移動のほうがはるかに高価だということです。

おおよそのエネルギーコスト比較 (相対値, 傾向):
  レジスタ/PE内部アクセス      : 1倍
  オンチップSRAMアクセス       : 数〜数十倍
  チップ外DRAMアクセス         : 数百〜数千倍

正確な数値はプロセスや設計で異なりますが、方向は一貫しています。DRAMまで往復するコストが圧倒的に大きいのです。したがってアクセラレータ設計の目標は、単に乗算器をたくさん積むことではなく、一度読んだデータでできるだけ多くのMACを行うことです。

この効率を定量化する概念が算術強度（arithmetic intensity）です。

算術強度 = 実行した演算数 / 移動したバイト数  (単位: FLOP/byte)

算術強度が高いほどメモリ帯域に縛られにくく、演算器の利用率が上がります。シストリックアレイが狙うのはまさにこの指標の最大化です。重みを一つ読んで格子内にロードした後、数百個の活性値と掛ければ、その重みの算術強度は数百FLOP/byteまで跳ね上がります。

データフローアーキテクチャ一般論

シストリックアレイは、より大きな概念であるデータフローアーキテクチャの一例です。伝統的なフォン・ノイマン（von Neumann）構造は、命令がデータを引き寄せる制御主導（control-driven）方式です。一方データフローは、データが準備できると演算が発火（fire）するデータ主導方式です。

フォン・ノイマン:  PCが命令を指す → データを取得 → 実行 → 結果を保存
データフロー    :  オペランドが到着 → 即座に発火 → 結果を次のノードへ伝える

データフローの魅力は次の通りです。

  • 明示的なプログラムカウンタや複雑な制御フローが減ります。
  • データ依存が自然に並列性を露わにします。
  • データが演算器の間を直接流れるので、中間結果をメモリに保存して読み直す無駄が減ります。

現代のAIアクセラレータは、程度の差はあれこのデータフローの哲学を共有します。計算グラフ（computation graph）をチップ上に空間的に展開し、テンソルをその上に流すのが核心的な発想です。

GPUのテンソルコアとの比較

NVIDIA GPUは伝統的に、数千個の小さなコアがSIMT方式で動作する構造でした。ところが2017年のVolta世代から、テンソルコア（Tensor Core）という専用の行列積ユニットが追加されました。テンソルコアは、小さな行列タイル（たとえば4×4や16×16）を一命令で積和するハードウェアです。

シストリックアレイとテンソルコアを比較すると次の通りです。

側面TPU シストリックアレイGPU テンソルコア
格子規模一つの巨大な格子(例: 256×256)小さなタイルユニット多数をコアに分散
制御方式コンパイル時の決定論的スケジュール実行時のスレッド/ワープスケジューリング
柔軟性行列積に高度に特化汎用GPU演算と共存, 柔軟
データ再利用格子の流れでハードウェアが強制レジスタ/共有メモリの活用に依存

大まかにまとめると、TPUは行列積という一つの仕事のためにチップ全体を一塊として設計したのに対し、GPUは汎用性を保ちながら行列積の加速ユニットを差し込んだ形です。2026年現在、NVIDIA Blackwell世代と次世代Vera Rubinは第2世代Transformer Engineとより大きなテンソルコアを搭載し、低精度演算とメモリ帯域で大きな進展を見せています。

どちらが優れているというより、ワークロードと運用環境によってトレードオフが異なります。大規模な学習/推論のためのデータセンター環境では、両方のアプローチが生き残り競争しています。

利用率とタイリング

シストリックアレイが理論上クロックごとに65,536個のMACを実行できるからといって、実際に常にそれだけ働くわけではありません。実際の利用率（utilization）はしばしばそれより低くなります。理由は、行列の形が格子サイズにきれいに割り切れないからです。

たとえば256×256の格子に9×9の行列を流し込むと、格子のごく一部だけが働き、残りは遊びます。また格子を満たすには、データが流れ込み流れ出るパイプラインの充填/排出（fill/drain）区間があり、この区間の間は一部のPEがアイドル状態になります。

この問題に対処する技法がタイリング（tiling）です。大きな行列を格子サイズに合う小さなタイルに分け、順に流し込みます。

大きな行列 C (1024 x 1024) を
格子サイズ 256 x 256 のタイルに分割:

  +------+------+------+------+
  | T00  | T01  | T02  | T03  |
  +------+------+------+------+
  | T10  | T11  | T12  | T13  |
  +------+------+------+------+
  |  ... 4 x 4 = 16個のタイル ... |
  +------+------+------+------+

各タイルは格子を満たすので利用率が高い。

タイルサイズを格子に揃え、パイプラインが十分深く満たされるよう十分多くのタイルを連続して流し込むことが、利用率を引き上げる鍵です。行列の次元が格子サイズの倍数のとき、利用率が最も良くなります。

コンパイラのマッピング

開発者が直接PE一つひとつにデータを流し込むわけではありません。それはコンパイラが行います。TPUの場合はXLA、GPUの場合はcuBLAS/cuDNNやTritonといったツールが、高水準のテンソル演算をハードウェアの格子にマッピングします。

コンパイラの核心的な作業は次の通りです。

1. 計算グラフから行列積ノードを見つける。
2. 行列の次元(M, N, K)を格子サイズに合うタイルに分割する。
3. データフロー戦略(weight-stationary など)を選ぶ。
4. タイルを流し込む順序をスケジューリングする。
5. データロードと演算が重なるようパイプラインを構成する。

この過程の品質が最終的な性能を左右します。同じハードウェアでも、コンパイラがタイルサイズと流れをどれだけうまく決めるかで利用率が大きく変わります。だからアクセラレータ企業は、ハードウェアと同じくらいコンパイラスタックに莫大な投資をします。

開発者にとっての実用的な示唆は、行列の次元をなるべくハードウェアに優しい値（たとえば8の倍数、128の倍数）に揃えると、コンパイラがより良いマッピングを見つけやすいということです。

長所と短所

シストリックアレイのアプローチの長所と短所を整理します。

長所:

  • 行列積に対して非常に高いスループットとエネルギー効率を出します。
  • データ移動をハードウェア構造が強制的に減らし、メモリウォール問題を緩和します。
  • 制御が単純で、チップ面積の大部分を実際の演算器に割り当てられます。
  • 決定論的なので性能予測とスケジューリングが容易です。

短所:

  • 行列積ではない不規則な演算には不向きです。
  • 行列の形が格子と合わないと利用率が下がります。
  • パイプラインの充填/排出のオーバーヘッドが小さな行列で目立ちます。
  • 格子サイズが固定なので柔軟性がGPUより低いです。

このトレードオフのため、シストリックアレイは万能の解ではなく、行列積が圧倒的に支配的なディープラーニングのワークロードに特化したツールとして理解するのが正確です。

開発者への示唆

ハードウェアを直接設計しない開発者にとっても、この原理は実務的に有用です。

  • バッチサイズを大きくすると、weight-stationaryのアクセラレータで重み再利用が増え効率が良くなります。推論サービングでバッチングが重要な理由の一つです。
  • 行列の次元を揃えると（例: 128の倍数にパディング）、コンパイラが格子を満たすタイルを作りやすく利用率が上がります。
  • 低精度演算(FP8, FP4 など)を活用すると、同じ格子でより多くのMACを処理でき、データ移動のバイト数が減って算術強度が上がります。
  • メモリアクセスパターンを意識すると、テンソルを連続的で整列したレイアウトに保つことがデータの流れに有利です。

要するに「乗算は安く、データ移動は高い」という原則を頭に置けば、モデル構造とサービング設定をハードウェアに優しく整える直感が育ちます。

おわりに

シストリックアレイは、単純な発想を最後まで押し進めた結果物です。データを一度読んだら最大限再利用せよという原則を、演算器を格子に配置しデータをその上に流す構造へそのまま刻み込みました。心臓が拍動するようにクロックごとにデータが一マスずつ前進して掛けられ累積される、この優雅な流れが、TPUの圧倒的な行列積性能を支える心臓原理です。

データフロー戦略の選択、タイリングと利用率、コンパイラのマッピングという三つの軸を理解すれば、なぜ特定のアクセラレータが特定のワークロードで輝くのか、そして自分のモデルをどうハードウェアに優しく整えるかについて、確かな直感を得られます。AIハードウェアの発展は結局メモリウォールとの終わりなき戦いであり、シストリックアレイはその戦いで最も長く検証されてきた武器の一つです。

参考資料

Discuss on Twitter • View on GitHub