2024-12-24

Tenstorrent WormholeでFFT (DFT)

9月くらいにTenstorrentのWormhole n300sというPCIeカード型のAIアクセラレータを購入したが、ジェット機並みの騒音の冷却ファンに辟易してあまり触れていなかった。屋外に設置する用のPCを準備できた、ということで開帳して年末休みに色々試していく。

Tenstorrent Wormholeとは？

Tenstorrent Wormholeは、RISC-Vマルチコアプロセッサ (Tensix-Core) を64基ほど搭載する、PCIeカード型AIアクセラレータ。
MIMD型メニーコアに分類され、特徴としてはスクラッチパッドメモリ、2Dトーラス状インターコネクトを備える点だろう。中間結果のストアにはスクラッチパッドメモリを使用し、コア間のデータ共有には相互結合網を用いることでメモリ律速を回避することに全力を注いだプロセッサという印象。 tenstorrent.com

プロセッサアーキテクチャの詳細に関しては、以下の記事が詳しい
Tenstorrent Wormhole Series Part 1: Physicalities

プログラムインターフェース

Tenstorrentは抽象度の異なる３つのSDKをいずれもOSSとして公開している。

tt-budaはPyTorchやTensorflowで既成のモデルを用いる場合と似たPython環境を提供する。 tt-nnはnumpyほどの抽象度を提供するPython環境で、深層学習モデルの記述を行う場合を想定している。
tt-metalは低レベル開発環境で、スクラッチパッドメモリからインターコネクト、個々のコアまでをプログラムする環境を提供する。 RISC-Vコア上で動作するプログラムはC/C++で記述し、行列演算器やSIMD演算器は独自のAPIでプログラムする仕組みである。

ファンのノイズについて

ファンのノイズは掃除機と同じかそれ以上の音量といったところ。付属するブロアファンクーラーは2線式なのだが、実はクーラーのケースの一部を切断し基板を露出させると4線式のPWM信号、回転パルス信号の信号線が切断されていることが確認できる。 PWM信号線 (入力) は内部でプルアップされており、（意図的に）常に100%の回転数となるよう出荷時に設定しているようだ。温度を監視する限り、チップ温度はアイドル状態であれば余裕があるので、信号線をはんだ付けしてマザーボード内蔵のファンコンから回転数を制御できるよう改修した。常に計算を実行する想定では無いので問題はないだろう。
これに関しても、後で記事にしようと思う。

FFT実装方針

大域フィルタ層 (Global Layer Filter) は大域的な特徴の抽出を目的とした畳み込み層の一種で、入力テンソルと同じ形状のフィルタを用いる点が特徴だ。このように大型のフィルタを用いることは、TransformerのAttention機構と同様、大域的特徴の学習を可能とする一方で $O(N^{2})$ の計算量を要求する。この計算量の削減手法として、入力データとフィルタを周波数空間上に変換し、要素積を求め、逆変換する手法が提案されている。 FFTで周波数空間に変換・逆変換する場合、その計算量はどちらも $O(N\log N)$ であり、 $N$ が十分に大きければ計算量の削減が可能だ。一方で、現在のGPUは密行列の演算能力は高い一方、FFTを比較的苦手とすることから効果は限定されるのが現状である[1]。
Wormhole上のインターコネクトやスクラッチパッドメモリを活かした効果的なFFT実装を提案出来れば価値はあるだろう。
なお今回、はじめにPE1基のFFT実装を実施し、その後複数PEの並列処理を実装してく予定だ。　 FFTのアルゴリズムとしては、基本的な周波数間引きCooley-Turkeyのアルゴリズムを踏襲する。 FFTをN=2,4まで分割する実装が有名で、この場合、N=2や4の計算ステップ以降は、乗算を符号反転や複素共役で代替できるトリックが採用される。 (このトリックを指して"バタフライ演算"と呼ぶこともあるが、この記事では便宜上計算フロー上の交差を総称するものとする)
このトリックは、乗算器が限定される通常のプロセッサであれば有効だが、今回使用するWormholeは乗算器を大量に集積しており、 8x16 @ 16x16 の密行列積を1サイクルで計算できる。 tt-metalのAPI上基本となる行列サイズの $32\times 32$ に換算すると、16サイクルで計算可能だ。その為、ある程度のNでFFTを打ち切って、通常のDFTに切り替える複合的なアプローチを採用した方が高速かもしれない、という推測が立つ.

さて、どこまでのNで打ち切れば最も演算に必要なサイクル数を抑えられるか初めに考察しよう.

DFTに行列乗算器を使用した場合の必要サイクル数

[cycle]	Matmul & Eltwise-FMA (8x16)@(16x16)	Sum
N=64	128	128
N=32	16	16
N=16	4	4
N=8	1	1

FFTに行列乗算を使用した場合の必要サイクル数

[cycle]	Matmul & Eltwise-FMA (8x16)@(16x16)	Eltwise Mul (8x16)	Sum
N=1024	16	8	24
N=64	16	8	24
N=32	16	8	24
N=16	4	2	6
N=8	1	0	1

この結果から、 $N=32$ でDFTに切り替えた方が計算に必要なサイクル数を最小限に抑えることがわかる。

FFTの行列演算実装

Wormholeのアーキテクチャに沿ったFFTの計算方法を考えていく。

サンプリング点 $t_k$ の各要素を行列 $X$ に以下のように格納する。ここで $d$ は行列の行・列の長さを示す。

$X = \begin{bmatrix} x_0 & x_{1} & \cdots & x_{d-1} \\ x_{d} & x_{d+1} & \cdots & \vdots \\ \vdots & \vdots & \ddots & \vdots \\ x_{d\cdot (d-1)} & \cdots & \cdots & x_{d\cdot d-1} \end{bmatrix}$

バタフライ演算のたすき掛け状のデータフローを行列乗算で実現する。シフト量 $s$ に応じたバタフライ演算行列 $B_N$ を右から掛けることで実現する. この行列の要素 $b_{ij}$ は以下である.

$\begin{eqnarray} B_N &=& [\, b_{ij} \,] \\ b_{ij} &=& \begin{cases} 1 & \text{if } i = j, \\ 1 & \text{if } i - s = j, \\ -1 & \text{if } i = j - s , \\ 0 & \text{otherwise}. \end{cases} \end{eqnarray}$

このときのシフト量は以下である。ここで $d$ は行列の要素数の平方根である。

$s = \begin{cases} N / d & \text{if } N > d, \\ N & \text{if } N \leq d. \end{cases}$

回転因子を要素として持つ行列 $W_N$ も定義する.

$\begin{eqnarray} W_N &=& [\ w_{ij}\ ] \\ w_{ij} &=& \begin{cases} w^{i\cdot d + j \mod N} & \text{if } \lfloor\frac{i\cdot d + j}{N}\rfloor \mod 2 \equiv 1, \\ 1 & \text{otherwise}. \end{cases} \\ w &=& e^{-j\frac{2\pi}{2N}} \end{eqnarray}$

以上で定義した行列を用いて、FFTは次のように計算できる。

$\begin{equation} F(X) = F_{k=1} \end{equation} \\ F_k = \begin{cases} X & \text{if } k = N, \\ F_{2k}^T B_k \odot W_k & \text{if } k \geq d, \\ F_{2k} B_k \odot W_k & \text{if } k < d. \end{cases}$

16点FFTにおける、計算グラフと行列演算の対応を次の画像に示す.

このアルゴリズムの注意点としては、 $k\geq d$ の結果に対してFFTを実施する場合、計算結果に対して転置が必要となることだ. 高速な行列レジスタ上で終始完結する処理とはいえ、tt-metalのtranspose APIはどの程度のサイクル数が必要かは不明なため、注意が必要だ。
通常のFFT同様、計算結果の系列は入力系列の順序と異なることからbit-revertでインデックスを変換する必要がある。

DFTの行列演算実装

通常のDFT同様、GEMM ( $32\times 32=1024$ 点を一斉に計算) で計算する。

$\begin{eqnarray} F(X) = W_{32}X &=& \begin{bmatrix} w^{0\cdot 0}_{32} & w^{0\cdot 1}_{32} & \cdots & w^{0\cdot 31}_{32} \\ w^{1\cdot 0}_{32} & w^{1\cdot 1}_{32} & \cdots & \vdots \\ \vdots & \vdots & w^{k\cdot t}_{32} & \vdots \\ w^{31\cdot 0}_{32} & w^{31\cdot 1}_{32} & \cdots & w^{31\cdot 31}_{32} \end{bmatrix} \begin{bmatrix} x_0 & x_{32} & \cdots & x_{32\cdot 31} \\ x_1 & x_{33} & \cdots & \vdots \\ \vdots & \vdots & \ddots & \vdots \\ x_{31} & x_{63} & \cdots & x_{1023} \end{bmatrix} \\ &=& \begin{bmatrix} F_{t_0\rightarrow t_{31}}(0) & F_{t_{32}\rightarrow t_{63}}(0) & \cdots & F_{t_{32\cdot 31}\rightarrow t_{32\cdot 32 - 1}}(0) \\ F_{t_0\rightarrow t_{31}}(1) & F_{t_{32}\rightarrow t_{63}}(1) & \ddots & F_{t_{32\cdot 31}\rightarrow t_{32\cdot 32 - 1}}(1) \\ \vdots & \ddots & F_{t_{32j}\rightarrow t_{32(j+1)-1}}(i) & \vdots \\ F_{t_0\rightarrow t_{31}}(31) & \cdots & \cdots & F_{t_{32\cdot 31}\rightarrow t_{32\cdot 32 - 1}}(31) \end{bmatrix} \end{eqnarray}$

回転因子は次式で示される.

$w_{32} = e^{-j\frac{2\pi}{32}}$

この回転因子行列 $W_N$ に関しては、再計算する必要もない定数なのでコンパイル時算出するパラメータとして与える。

スクラッチパッドメモリ使用戦略

キャッシュメモリと異なり、PEのL1はプログラマ自身が制御しなければならないスクラッチパッドメモリだ。計算に必要な行列と必要な容量について整理しておく.

行列名	形状	サイズ	内容
$X_{1024}$	$32\times 32$	4KiB	1024点の入力データ
$W_{32}$	$32\times 32$	4KiB	N=32の回転因子行列
$W_{64}$	$32\times 32$	4KiB	N=64の回転因子行列
$W_{128}$	$32\times 32$	4KiB	N=128の回転因子行列
$W_{256}$	$32\times 32$	4KiB	N=256の回転因子行列
$W_{512}$	$32\times 32$	4KiB	N=512の回転因子行列
$W_{1024}$	$32\times 32$	4KiB	N=1024の回転因子行列
$B_{64}$	$32\times 32$	4KiB	N=64のバタフライ演算行列
$B_{128}$	$32\times 32$	4KiB	N=128のバタフライ演算行列
$B_{256}$	$32\times 32$	4KiB	N=256のバタフライ演算行列
$B_{512}$	$32\times 32$	4KiB	N=512のバタフライ演算行列
$B_{1024}$	$32\times 32$	4KiB	N=1024のバタフライ演算行列
$Y_{1024}$	$32\times 32$	4KiB	1024点の結果・中間結果格納行列

Wormholeは、半精度浮動小数点数を基本として扱うため、数値１つにつき2Byteが必要だ。加えて、複素数であることを考慮すると4Byte必要である。 $13\times 2\times 2$ KiB (52KiB)程度の容量をL1に一度に格納できればプログラムが格段に組みやすくなる。 WormholeはPE1基あたり1MiB程度のL1を持つので、かなり余裕がある。使用メモリ量に対して16倍ほどの空きがあることから、おおよそ1基のPEあたり16384点のFFTに対応する余地があるだろう。

今日のまとめ

アーキテクチャへの落とし込みは一通り目途がついた。次回以降実際に実装する。

参考図書

[1] 深層ニューラルネットワークの高速化 ML Systems

2020-12-01

vai_q_tensorflowをソースからビルドできない問題

vai_q_tensorflowはGithubにリポジトリが公開されており,ビルド方法もREADME.mdに記されているが,実行したところBUILDファイルが見つからないと怒られる.vai_q_tensorflowはTensorflow r1.15からのフォークであり,これと同様にbazelを用いてビルドするがレシピであるBUILDファイルがなぜか見当たらない..gitignoreを参照したところ,どうやら間違って一部のBUILDファイルを含めないよう設定されているようだ(適当すぎる...).そこでフォーク元のtensorflowからBUILDファイルを取得し配置する必要があった.その他にも問題が複数発生したので対処法を記す.
↓修正済みのリポジトリ github.com

手順1

tensorflowをgit pull.vai_q_tensorflowと同階層のディレクトリに配置する.その後tensorflowのブランチをr1.15にcheckoutしておく.

手順2

次のシェルを実行する

refer my blog

手順3

私の環境では次のエラーも発生した.これはフォーク元のtensorflow側の問題の様だ.それぞれTensorflowのGithubに同様のissueが存在した.

#34197

github.com 3つのソースを書き替える必要がある.

#41061

github.com 上記のissueで述べられているが,次の変更を加えることでコンパイルが成功する. github.com

2020-07-03

RT(TTU)コアの処理内容・内部構造の調査

patents.google.com

図などは上記の明細書を参考のこと
TTU - Tree Traversal Unit (nVidiaのTuringコアアーキテクチャから採用された"RTコア"の開発名)
レイトレーシング（Ray Tracing)において、1つの光線毎に、 $n$ のプリミティブ（オブジェクトの最小単位,たいてい三角形）への衝突判定を行う必要があるが、愚直に計算するのは非効率.
処理数の削減を行うためにBVH(Bounding Volume Hierarchy)と呼ばれる木構造が導入される.下記のサイトを参照.
shinjiogaki.github.io TTUは,BVHを辿る処理(AABB (Axis Aligned Bounding Box) への衝突判定を含む,数千命令に相当)をハードウェアで実現する.

TTUの具体的処理

用語

SM - Streaming Multi-processor
Compression Block - BVHの部分木に対応するデータブロック

処理

SMから,対象のRay,Compression Blockへのポインタを含んだ命令を受け取る.
Schudulerに命令本体が転送される.Ray,Compression Blockへのポインタはそれぞれ別のLocal Storageに保管される.
Setup Unitに命令本体とCompression Blockへのポインタが転送される.
Setup Unitは、受け取った命令を基にLocal Storageに保存されたRayの情報と,受け取ったCompression Blockへのポインタを基にメモリから取得(キャッシュがヒットすれば即座に取得できる)したCompression Blockのデータを取得する.
Setup Unitは、各Traversal UnitにRayとCompression Blockを送信しCompression Block毎に並列に処理を実行する.疑似コードのOuter Loopに相当するLoopを回す.
Traversal Unitは、処理が完了するとStack Management Unit内部のResultキューにRayが衝突する可能性のあるプリミティブへのポインタを入れる.疑似コードのInner Loopに相当するLoopを回す.後述するがTraversal UnitはAABBへの衝突判定は行うがプリミティブへの衝突判定は行わない.プリミティブへの衝突判定は別のハードウェアで実現する.Resultキューに入ったプリミティブ（へのポインタ）は,プリミティブ専用の衝突判定処理を待機する.
Stack Management Unitは、処理が完了したTraversal Unitが有り,Local StorageにCompression Blockが存在するならばSchedulerにTraversal Unitが再度処理可能であることを通知する.

疑似コード

procedure outerTraversal(Ray, *bvhRoot)
  traversalStack
  traversalStack.push(bvhRoot)
  while traversalStack is not empty do
    *node = traversalStack.pop
    intersectedExternalNodes = innerTraversal(ray,node)
    traversalStack.push(intersectedExternalNodes)
  end do
end

procedure innerTraversal(Ray, *blockRoot)
  intersectedExternalNodes
  localStack.push(blockRoot)
  while localStack is not empty do
    *node = localStack.pop()
    if ray intersects node then
      addToResultQueue(node)
    end
    else if node is TransitionNode then
      intersectedExternalNodes.add(node.externalNodePointer)
    end
    else then
      childNodes = node.childNodes
      sort childNodes
      localStack.push(childNodes)
    end
  return intersectedExternalNodes
  end do
end

内部データ構造

やや具体的な話、実装によって多少変わってくると思う.

Compression Blockのアラインメントの例

struct Compression_Block{ // is restricted to Cache Size
  struct Node nodes[N];
};

Nodeのアラインメントの例

// double fixedは2ワード長の固定小数点型を想定
struct Node{ // Root Node
  char type_id : 2;
  char AABB_Bitmask : 6;
  double x_1; // 2word floating point data
  double x_2;
  double y_1;
  double y_2;
  double z_1;
  double z_2;
  struct Compression_Block *child_ptr;
};
struct Node{ // Internal/Leaf Node
  char type_id : 2;
  char AABB_Bitmask : 6;
  char x_1; // 1Byte=256 step data
  char y_1;
  char y_2;
  struct Compression_Block *child_ptr;
};

nodes 各ノードのデータ,連続配置される.順番は深さ優先,幅優先どちらでもよい.

type_id ノードの種別（root,leaf,transition）を表す

transition Compression Blockの境界となるノード,一方のleafノードであり,もう一方のrootとなるようなノード
leaf 葉ノード
root 根ノード

AABB_Bitmask 後述

x_1,x_2,y_1...等 座標の数値

child_ptr 葉ノードでなければ子ノードへのポインタ,葉ノードであればプリミティブへのポインタ(メモリ上を指す)

各ノードはAABBの情報（3次元直交座標系でおいては2つのベクトル,つまり6つの数値）を持つ.このうち,BVH全体のrootノードは全体座標系(Global Coordinate)における6つの数値を必ず持つ.しかしその子孫であれば,親ノードからの相対座標(Local Coordinate)で指定できる.さらに親ノードの数値を継承する場合、省略することが出来る.これらの工夫により数値の表現に必要なビット数が節約できる.上の構造体を模した疑似コードにおいて,AABB_Bitmaskは親ノードの数値を引き継ぐので省略する座標をビットマスクで表現している.

疲れたのでここまで