「M1搭載のMacBook Airはブラウザから900GFLOPSの計算が可能」という検証結果

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のようなテクノロジーが利用可能になることにワクワクさせられます」と述べました。