Apple独自開発のArmアーキテクチャベースのSoC「M1」は、グラフィックス性能を測るベンチマークでGeForce GTX 1050 TiやRadeon RX 560を上回ったと報告されるなど、非常に高性能であることが話題になっています。そんなM1を搭載したMacBook Airで単精度浮動小数点数演算性能を計測した結果を、開発者のBram Wasti氏が自身のウェブサイトで報告しています。
jott - m1_webgpu_perf
https://jott.live/markdown/m1_webgpu_perf
MacBookに標準搭載されているSafariは、WebGPUというGPU APIを実験的にサポートしています。そこで、行列の乗算を最適化するためのチューナーをコーディングしました。
チューナーのコードは以下で公開されています。
jott - webgpu_mm.js
https://jott.live/code/webgpu_mm.js
Wasti氏は「基本的な考え方は、メモリアクセスを並列化してベクトル化し、積和演算命令を使って、スレッド化とディスパッチパラメータをチューニングするということです」と解説しています。その結果、コンピュートカーネルは以下のようになるとのこと。
[numthreads(2, 8, 1)]
compute void main(constant float4[] A : register(u0),
constant float4[] B : register(u1),
device float4[] C : register(u2),
float3 threadID : SV_DispatchThreadID) {
uint m = uint(threadID.x);
uint n = uint(threadID.y);
float4 result_0_0 = float4(0.0, 0.0, 0.0, 0.0);
float4 result_1_0 = float4(0.0, 0.0, 0.0, 0.0);
float4 result_2_0 = float4(0.0, 0.0, 0.0, 0.0);
float4 result_3_0 = float4(0.0, 0.0, 0.0, 0.0);
for (uint k = 0; k < 256; k++) {
float4 a_0_0 = A[(m * 4 + 0) * 256 + (k * 1 + 0)];
float4 a_1_0 = A[(m * 4 + 1) * 256 + (k * 1 + 0)];
float4 a_2_0 = A[(m * 4 + 2) * 256 + (k * 1 + 0)];
float4 a_3_0 = A[(m * 4 + 3) * 256 + (k * 1 + 0)];
float4 b_0_0 = B[(k * 4 + 0) * 256 + (n * 1 + 0)];
float4 b_0_1 = B[(k * 4 + 1) * 256 + (n * 1 + 0)];
float4 b_0_2 = B[(k * 4 + 2) * 256 + (n * 1 + 0)];
float4 b_0_3 = B[(k * 4 + 3) * 256 + (n * 1 + 0)];
result_0_0 += mul(a_0_0.x, b_0_0);
result_1_0 += mul(a_1_0.x, b_0_0);
result_2_0 += mul(a_2_0.x, b_0_0);
result_3_0 += mul(a_3_0.x, b_0_0);
result_0_0 += mul(a_0_0.y, b_0_1);
result_1_0 += mul(a_1_0.y, b_0_1);
result_2_0 += mul(a_2_0.y, b_0_1);
result_3_0 += mul(a_3_0.y, b_0_1);
result_0_0 += mul(a_0_0.z, b_0_2);
result_1_0 += mul(a_1_0.z, b_0_2);
result_2_0 += mul(a_2_0.z, b_0_2);
result_3_0 += mul(a_3_0.z, b_0_2);
result_0_0 += mul(a_0_0.w, b_0_3);
result_1_0 += mul(a_1_0.w, b_0_3);
result_2_0 += mul(a_2_0.w, b_0_3);
result_3_0 += mul(a_3_0.w, b_0_3);
}
C[(m * 4 + 0) * 256 + (n * 1 + 0)] = result_0_0;
C[(m * 4 + 1) * 256 + (n * 1 + 0)] = result_1_0;
C[(m * 4 + 2) * 256 + (n * 1 + 0)] = result_2_0;
C[(m * 4 + 3) * 256 + (n * 1 + 0)] = result_3_0;
}
dispatch params: 128,32,1
Wasti氏は、自身が所有するM1チップ搭載のMacBook Airで数秒のチューニングを行い、実際に計算を行ったところ、900GFLOPSを達成したと報告しています。Wasti氏が所有するIntel Core i9搭載のMacBook Pro(2019年モデル・16インチ)は、同じ処理で100GFLOPSしか達成しなかったとのこと。
また、M1チップ搭載のMacBook AirでエッジAI「MobileNet v3 Large(x1.0)」のパフォーマンスを計測すると最大219MFLOPSで、これは1秒当たり4500回の推論を実行できるレベルだとのこと。また、Googleの自然言語処理モデルであるBERT(12層構成)では11.2GFLOPSで、1秒間に90回の推論を実行できるパフォーマンスに相当するそうです。
Wasti氏は「調整すればもっと多くのことができますが、結果には非常に満足しています。ブラウザでほぼ1TFLOPSをたたき出したという事実は非常に心強く、M1のようなテクノロジーが利用可能になることにワクワクさせられます」と述べました。