I8MM: INT8 行列積向け ARM 拡張

INT8の行列積の計算を行うことができるI8MM の使い方の紹介です。

I8MMの概要

AI 推論に必要なデータの精度は、多くの場合 8 ビット整数 (INT8) で用が足りると知られていますが、 CPU や GPU などのプロセッサは最近になるまで INT8 の計算を集中的に行うワークロードは想定して設計されていませんでした。近年の AI ブームに合わせて INT8 の計算処理がこれらのプロセッサに追加されるトレンドがあり、例えば x86 における VNNI 拡張や、 Nvidia の GPU における Tensor Core と呼ばれる演算器は、このトレンドに則った新機能として高い評価を得ています。

ARM プロセッサにおいてもこのトレンドに追随すべく、 I8MM と呼ばれる INT8 の行列積の計算に便利な拡張命令が発表され、搭載製品の出荷も始まっています。

今回は、 AI 推論の高速化に便利なこの I8MM の使い方の紹介をしようと思います。

I8MM の基本仕様

UMMLA/SMMLA/USMMLA/USDOT の 4 命令からなり、 C 言語の vmmlaq_u32/vmmlaq_s32/vusmmlaq_s32/vusdotq_s32 のイントリンシックが対応します。このうち USDOT は仕様が独特で、第 1 オペランドは INT32 x 4 のベクトル、第 2 オペランドは UINT8 x 16 のベクトル、第 3 オペランドは INT8 x 16 のベクトルで、各オペランドを C[0:3] 、 A[0:15] 、 B[0:15] と表現すると、

C[0] += dot(A[0:3], B[0:3])
C[1] += dot(A[4:7], B[4:7])
C[2] += dot(A[8:11], B[8:11])
C[3] += dot(A[12:15], B[12:15])

の計算を行います。要は x86 の VNNI のような動作です。一見これで用が足りそうなのですが、 UMMLA/SMMLA/USMMLA は後述のようにより効率的な命令ですので、 x86 との互換などの限られた用途のみで使われることになると思います。

UMMLA/SMMLA/USMMLA の動作は基本的に同じで、各オペランドの整数の解釈の符号の有無のみの違いになっています。

基本的な動作は、第 1 オペランドの 32 ビット整数に対して、第 2 オペランドの 8 ビット整数の 8 要素と第 3 オペランドの 8 ビット整数の 8 要素のドット積を取って足しこむというものです。各オペランドを C[0:3] 、 A[0:15] 、 B{0:15] と表現すると、

C[0] += dot(A[0:7], B[0:7])
C[1] += dot(A[0:7], B[8:15])
C[2] += dot(A[8:15], B[0:7])
C[3] += dot(A[8:15], B[8:15])

の計算を行います。 USDOT が 1 命令あたり 16 加算 + 16 乗算の計算を行うのに対して、これらの命令は 1 命令当たり 32 加算 + 32 乗算と効率が 2 倍になっています。

各オペランドの符号は、

UMMLA: C: UINT32, A: UINT8, B: UINT8
SMMLA: C: INT32, A: INT8, B: INT8
USMMLA: C: INT32, A: UINT8, B: INT8

です。

コンパイラのサポート

Clang コンパイラでは、 -march で指定するアーキテクチャの後ろに +i8mm の文字列を付け加えることで I8MM 使用を伝えることができます。

+i8mm と併用できるアーキテクチャは armv8.2-a から armv8.5-a の範囲で、例えば armv8.2-a に I8MM 対応を付け加えたいときは、 -march=armv8.2-a+i8mm とオプションを指定してやります。

armv8.1-a 以下は I8MM 非対応で、 armv8.6-a 以上は I8MM 対応をデフォルトで含みます。

I8MM 対応が有効になっている場合は __ARM_FEATURE_MATMUL_INT8 マクロに 1 が設定されます。

ax株式会社はAIを実用化する会社として、クロスプラットフォームでGPUを使用した高速な推論を行うことができるailia SDKを開発しています。ax株式会社ではコンサルティングからモデル作成、SDKの提供、AIを利用したアプリ・システム開発、サポートまで、 AIに関するトータルソリューションを提供していますのでお気軽にお問い合わせください。