T-Headが提案しているRISC-VのMatrix Extensionについて、マニュアルを読みながら理解していこうと思う。

とりあえずマニュアルで、どのようなレジスタが存在しているのかを理解していく。プログラミングモデルとサンプルコードも読み進めていきたい。

この文書は、RISC-Vのマトリックス拡張の提案である。

この拡張は、柔軟性と実装を考慮し、ベクタ拡張から切り離されたアーキテクチャを選択している。マトリクス拡張をサポートする各Hartsは、RLEN定数パラメータを定義する。RLENは、セクション「行列レジスタ」で説明されているように、行列レジスタの行のビット数である。プログラマのモデルは、異なるRLEN値を持つハーツ上でバイナリの移植性を達成するように設計されている。

この拡張はRISC-V Vector拡張に強くインスパイアされている。

プログラマのモデル

マトリックス拡張は、8つの2次元マトリックス・レジスタと6つのCSRを追加する。

アドレス	特権	名前	説明
xxxx	URW	xmrstart	行列の開始行位置
xxxx	URW	xmcsr	行列の制御と状態レジスタ
xxxx	URW	xmsize	行列のサイズのコンフィグレーション、行列のサイズ
xxxx	URO	xmlenb	行列レジスタのバイトサイズ、mrows*xrlenb
xxxx	URO	xrlenb	RLEN (バイトサイズ)
xxxx	URO	xmisa	行列命令のサブセット

行列レジスタ

行列拡張が実装されている場合、M0、M1、M2、M3、M4、M5、M6、M7という8つの2次元行列レジスタがアーキテクチャ状態に追加される。

RLENは各レジスタ行のビット長である。RLEN はどのような実装においても定数であり、演算によって生成または消費される行列要素の最大ビットサイズよりも小さく、2のべき乗でなければならず、 $2^{16}$ を超えてはならない。行列レジスタの行はRLEN/32、たとえば4/8/16である。したがって、各行列レジスタは [M x K] MSEW 要素で構成される。なお、M=RLEN/32, K=RLEN/MSEW, MSEWは要素サイズである。

注： $M\times K\times \text{MSEW} = \text{RLEN}/32\times\text{RLEN}/\text{MSEW}\times\text{MSEW} = \text{RLEN}\times\text{RLEN} / 32 = \text{MLEN}$ と呼ぶ

MSEWが一定の場合、RLENが2倍になると行も列も2倍になり、結果的に4倍のビット数になる。

RLENが一定の場合、MSEWが2倍になると、行は変わらず、列は半分になる。

MLEN : the size in bits for each register corresponding MLEN values 512/2048/8192/32768...
RLEN : the length in bits for each register row supported RLEN values 128/256/512/1024...
EW: element width supported EW values 4/8/16/32/64...

行列レジスタのサイズ例は以下のように変化する。

size[2:0]	operand datawidth	MLEN(bit)	M	K	Matrix size in bits
100	4bit	128	4	32	512
		256	8	64	2048
		512	16	128	8192
000	8bit	128	4	16	512
		256	8	32	2048
		512	16	64	8192
001	16bit	128	4	8	512
		256	8	16	2048
		512	16	32	8192
010	32bit	128	4	4	512
		256	8	8	2048
		512	16	16	8192
011	64bit	128	4	2	512
		256	8	4	2048
		512	16	8	8192

この表を見てもよくわからないのだが、図を見るとわかる。実装依存で行列レジスタの行サイズ(RLEN)がハードウェア固定となっている。実際の行列レジスタは $\text{RLEN}\times\text{RLEN}/32$ の大きさなのだが、まずはRLENの使い方から。

行列の行の大きさは常にRLEN/32となるように設計されており、たとえばRLEN=128であれば常に行の要素数はRLEN/32=4で固定される。

その代わりに、行数が変わるということになる。RLEN=128の場合、データサイズを32ビットにすると、行が消費するのは $4\times32=$ 128-bitsで4行となる。データサイズが16ビットに縮小されると、行の要素数が変わらない代わりに、行数が8行と2倍に増える。一方でデータサイズが64ビットに増加すると、行の要素数が変わらない代わりに、行数が2行と1/2に縮小されるという寸法だ。

行列サイズの構成

行列サイズ構成は XLEN ビットの WARL リードライトレジスタで、行列構成命令によってのみ更新される。行列サイズ・レジスタは sizeK、sizeN、sizeM の 3 つのフィールドを持つ。ビット[XLEN-1:32]は予約されている。

bits	Name	Description
XLEN-1:XLEN-32	0	reserved if non-zero
31:16	sizeK[15:0]	column of Matrix A or Matrix B, バイト単位
15:8	sizeN[7:0]	row of Matrix B
7:0	sizeM[7:0]	row of Matrix A

sizeM/sizeN/sizeKフィールドは、行列命令によって必要とされるソース要素と更新されるデスティネーション要素を指定する符号なし整数を保持する。sizeKが要素サイズのバイト数倍でない場合、不正命令例外が発生する。

$C[M$ [N]+=A[M][K]*B^T[N][K]] を計算する行列乗算命令では、3つのソースオペランドと1つのデスティネーションオペランドがある。サイズM×サイズNの要素だけが更新され、他の要素はゼロで設定される。ソース・オペランドの次元は以下のように定義される：

行列 A: sizeMx (sizeK/ 要素サイズ)
行列 B: sizeNx (sizeK/ 要素サイズ)
行列C： sizeMx sizeN

このように、行列レジスタによる行列形状の制約がある。

sizeK≤ xrlenb
sizeM≤ RLEN/32
sizeN ≤ RLEN/32、fmmacc.hの場合 sizeN≤ 2*(RLEN/32)

RLEN=128 の 32 ビット行列乗算を例にとると、sizeM=2 / sizeK=12 / sizeN=2 の構成は、 $\text{MatrixA}(2\times 3)\times \text{MatrixB}^T(2\times 3)+\text{MatrixC}(2\times 2)$ を示し、緑色のブロック要素のみが命令によって使用または更新される。

ポイントワイズ命令とロード/ストア命令では、sizeMとsizeKで指定された行列の形状が実行中に保持される。sizeM×sizeN要素のみが更新され、他の要素はゼロに設定される。サイズの制限は以下の通りである：

sizeM≤ RLEN/32
sizeK≤ max_colb

Int32 行列の加算を例として，sizeM=2/sizeK=12 の構成は， $\text{MatrixA}(2\times 3) +\text{MatrixB}(2x3)=\text{MatrixC}(2\times 3)$ を示し，緑色のブロック要素のみが命令によって使用または更新される。

行列の制御とステータス

xmcsr CSR は WARL リードライトレジスタである。ビット[XLEN-1:3]は予約されており、ゼロを書き込む必要がある。行列制御およびステータス・レジスタのレイアウトは以下の通りである：

bits	name	description
XLEN-1:3	0	reserved if non-zero
2	xmsat	Fixed-point accrued saturation flag
1:0	xmxrm	Fixed-point rounding mode

行列固定小数点丸めモード

行列固定小数点丸めモード(xmxrm)は行列制御およびステータス・レジスタの[3:2]ビットで定義される。xmxrmは以下のように vxrm[1:0]と同じ符号化と丸めアルゴリズムを使用する。丸め前の結果が v であり、その結果の d ビットが丸められるとする。丸め後の結果は、(v >> d) + rとなり、rは以下の表で指定される丸めモードに依存する。

vxrm[1]	vxrm[0]	rounding mode	rounding increment r
0	0	rnu round-to-nearest-up (add +0.5 LSB)	v[d-1]
0	1	rne round-to-nearest-even	v[d-1] & (v[d-2:0]≠0 \| v[d])
1	0	rdn round-down (truncate)	0
1	1	rod round-to-odd (OR bits into LSB, aka "jam")	!v[d] & v[d-1:0]≠0

以下の命令説明では、この演算を表すために丸め関数を使用している：

roundoff_unsigned(v, d) = (unsigned(v) >> d) + r
roundoff_signed(v, d) = (signed(v) >> d) + r

行列固定小数点飽和フラグ

xmxsatフィールドは、固定小数点命令が出力先のフォーマットに適合するために出力値を飽和させなければならなかったかどうかを示す。

行列レジスタ情報

行列レジスタ情報には、2つの読み出し専用XLENビット・レジスタがあり、これらはどの実装でも一定である。

xrlenb：各行列レジスタ行の RLEN ビットの状態を示すバイトの RLEN
xmlenb: 行列レジスタサイズ（バイト単位）， mrows*xrlenb, mrows=RLEN/32

行列開始行

xmrstart読み書きレジスタは、行列ロード/ストア命令で実行される最初の行列行インデックスを示す。通常xmrstartは、行列ロード/ストア命令のトラップ時にのみハードウェアによって書き込まれ、レジスタの符号なし値は、再開可能なトラップが処理された後に実行が再開されるべき行を指定する。

mcfg/mcfgiを含むすべてのマトリックス命令は、xmrstart CSRをゼロにリセットする。

xmrstart CSR は、最大行インデックス（最大行より 1 つ小さい）または log2(RLEN/32) を保持するのに十分な書き込み可能ビット数のみを持つように定義されている。xmrstart CSRの上位ビットはゼロにハードワイヤリングされている（読み出しはゼロ、書き込みは無視される）。

例えば、xmrstartは0から3までの行インデックスを表す2ビットを持つ。