シストリックアレイとは

← 07THEME 08 デバイスと回路09 →

シストリックアレイとは、単純な演算セルを規則正しく格子状に配置し、計算の途中結果を隣のセルへ直接手渡しながら行列演算を進める仕組みです。演算器そのものの速さよりも、途中結果の送り先まで含めて設計した配置といえます。乗算の答えを隣へ送るので、行列乗算の途中でメモリへ戻りません。そのかわり升目の大きさは固定で、性能は計算を升目の形へどれだけ合わせきれるかで決まります。余った升目は、ゼロを掛けて電力と時間を使うだけです。

心臓の拍動になぞらえた、隣とだけつながる並列計算機

systolicは「心臓収縮の」を意味する形容詞で、心臓が収縮を繰り返しながら血液を全身へ送り、また汲み上げる動きになぞらえて名づけられました¹。血液にあたるのがデータで、心臓にあたるのが1つ1つの演算セルです。

シストリックアレイは、KungとLeisersonが提案したVLSI向きの並列計算機の仕組みで、単純なプロセッサを多数、決まった規則で組み合わせた並列計算のシステムです。特徴は3つあります。各セルが左右に隣接するセルとだけ直接つながる局所結合、すべてのセルが単純で同じ構造を持ち実数の四則演算などを1ステップでこなす均一構造、そしてアレイへの入出力を逐次的に、アレイ上での計算を並列に行い、それらがパイプライン化されて重なりながら進むことです。

1992年の時点で想定されていた使い方は、ノイマン型計算機のシステムバスへつなぐ特定用途向けのハードウェア付加装置で、その内部には簡単なプロセッサが数千個から数万個、規則正しく並びます。この配置は、いまのAIアクセラレータがチップの中で占める位置とよく重なります。

CPUとGPUはメモリとの往復を繰り返し、TPUは隣の乗算アキュムレータへ答えを送る

CPUは、メモリから値を取り出し、演算し、その結果をふたたびメモリへ書き込む往復を繰り返します。この出入りの遅さがチップ全体の演算量に上限を作る現象が、フォンノイマンボトルネックです。GPUは演算器を数千個備えて演算量を稼ぎますが、汎用プロセッサとして数多くのソフトウェアを支えるため、何千もの演算器で行う演算のたびにレジスタや共有メモリへ出入りする点はCPUと共通です。

TPUでは、数千個の乗算アキュムレータが互いに直につながって1つの大きな格子をなしていて、Google Cloudはこの構造をシストリックアレイアーキテクチャと呼んでいます。乗算が1回済むごとにその答えが次の乗算アキュムレータへ送られ、データとパラメータの乗算結果をすべて足し合わせたものが出力になります。そのため、行列乗算の計算はメモリアクセスを挟まずに進みます。

メモリへ戻るか、隣へ渡すか

汎用プロセッサは演算のたびにレジスタや共有メモリへ出入りし、シストリックアレイは乗算の答えを隣の乗算アキュムレータへ手渡すので、行列乗算の途中でメモリへ戻りません。

TPUのMXUは256×256の升目で、1サイクルに16,000回の乗累算をこなす

TPUチップには1つ以上のTensorCoreが入り、そのTensorCoreが行列乗算ユニット（MXU）と、ベクトル用・スカラー用のユニットを備えます²。MXUは格子状に並べた乗算アキュムレータの集まりで、大きさはTPU v6eとTPU7xで256×256、それより前の世代で128×128です。辺が2倍なので、升目の面積は4倍になります。1つのMXUが1サイクルで16,000回の乗累算をこなし、乗算はbfloat16の入力を取り、累積はFP32で行います。

升目の大きさ

同じ縮尺で並べると、TPU v6eとTPU7xの256×256の升目は前の世代の128×128の4倍の面積で、その中には32×32の試作LSIの升目が縦横8つずつ入ります。

演算の中身そのものは素朴です。積和演算に特化した演算ユニットでは、大量の積和演算器による並列演算でニューラルネットワークの高速な計算が可能になります³。掛けて足す、それだけの回路をひたすら敷き詰めた升目が、シストリックアレイの実体です。

升目を埋めきれないと、ゼロで埋めた分だけ利用率が下がる

大きさの決まった正方形の升目には、裏返しの弱点が伴います。高い性能に届くCloud TPUプログラムは、密な計算を128×128のかたまりへタイル化できるものです。埋まりきらずに余った部分はコンパイラがテンソルをゼロで埋めるので、その埋め合わせの分だけTPUコアの利用率が下がります。ゼロを掛けた升目が費やす電力と時間は、そのまま損になります。

タイル化とゼロ埋め

計算を128×128のかたまりへ割り付けると、端で余った升目はコンパイラがゼロで埋め、その分だけTPUコアの利用率が下がります。

XLAコンパイラは最初のバッチで作った計算グラフを使い回すため、テンソルの形が毎回そろっている計算がTPUに向きます。升目の形に計算のほうを合わせる作業が、性能を左右します。

国内では、32×32の並列演算アレイを持つ試作LSIが作られた

NEDOと東京工業大学は2021年8月23日に、エッジ機器でのCNN推論向けに、入力データの平面シフトを受け持つ整形の仕組みと、直積型の並列演算アレイとを組み合わせたアーキテクチャー⁴を提案しました。出発点の課題は、演算を省いてモデルを小さくするとメモリの読み出し位置が飛び飛びになり、データを使い回しにくくなって演算器の空き時間が増えるため、並列処理の計算効率が落ちることです。升目を敷き詰める設計では、稼働率が共通の関心事になります。

試作LSIは並列演算アレイのサイズを32×32とし、活性値と係数値に4ビット固定小数点（INT4）の量子化を採用して、TSMCの40nmプロセスで製作されました。両者は、電源1.1V・最大534MHzという条件での消費電力を400mW以内と実測し、カーネル要素数を9分の1まで絞って残した要素だけを演算した条件で、この値が実効効率26.5TOPS／Wに相当するとしています。

よくある質問（FAQ）

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

単純な演算セルを規則正しく格子状に配置し、計算の途中結果を隣のセルへ直接手渡しながら行列演算を進める仕組みです。KungとLeisersonが、VLSI向きの並列計算機の仕組みとして提案しました。

名前の由来は何ですか？

systolicは「心臓収縮の」を意味する形容詞です。心臓が収縮を繰り返しながら血液を全身へ送り、また汲み上げる動きになぞらえて名づけられました。データが血液、演算セルが心臓にあたります。

なぜ行列演算が速く進むのですか？

途中結果の行き先が近いためです。乗算が1回済むごとにその答えが次の乗算アキュムレータへ送られるので、行列乗算の計算はメモリアクセスを挟まずに進みます。

どれくらいの大きさの升目ですか？

TPUのMXUは格子状に並んだ乗算アキュムレータの集まりで、大きさはTPU v6eとTPU7xで256×256、それより前の世代で128×128です。1つのMXUが1サイクルで16,000回の乗累算をこなします。

GPUとどこが異なりますか？

途中結果の行き先です。GPUは汎用プロセッサとして数多くのソフトウェアを支えるため、何千もの演算器で行う演算のたびにレジスタや共有メモリへ出入りします。

弱点はありますか？

升目を埋めきれるかどうかに性能が左右されます。高い性能に届くCloud TPUプログラムは密な計算を128×128のかたまりへタイル化できるもので、余った部分はコンパイラがゼロで埋めるため、TPUコアの利用率が下がります。

国内にも演算器を格子状に配置した例はありますか？

NEDOと東京工業大学は2021年8月に、CNN推論向けに直積型の並列演算アレイを中核へ据えたアーキテクチャーを提案し、アレイサイズを32×32とした試作LSIで最大26.5TOPS／Wの実効効率を実測したとしています。

比較・違いを学ぶ

AIアクセラレータとは── シストリックアレイを内側に持つ側の、チップ内の一区画
メモリ階層とは── シストリックアレイが往復を減らそうとしている相手側

出典

5件を表示