「GoとMetalシェーディング言語を通じてAppleのGPUをプログラミングする」

「AppleのGPUをプログラミングする魅力的なGoとMetalシェーディング言語」

Go、Cgo、Metalシェーディング言語、Metalパフォーマンスシェーダーの調査と行列の乗算に対する異なるアプローチのベンチマーキング

以下では、GoとネイティブC間のインターフェースとしてのcgoの使用方法、これを使用してAppleのMetalパフォーマンスシェーダーフレームワークのObjective-Cバインディングとのインターフェース方法、Metalシェーディング言語で記述されたカスタムGPUコード(シェーダー)とのインターフェース方法、そして手書きおよびOpenBLASに基づくGoベースの行列乗算演算とのベンチマークについて説明します。この記事は私のM2 MacBookで実行するために書かれました。

ソースのレイアウトは、ここでGitHubでご覧いただけます。以下のセクションに分けて説明しますが、ベンチマークへも直接移動することもできます。

GPUと浮動小数点並列処理

この時点で、ほとんどの人はGPUが特定の種類の計算タスクにおいて非常に強力であることに直感的に理解していると思います。特に、機械学習をサポートするいくつかのタスクにおいては、それがいかにCPUよりも強力であるかを自分自身で理解するまで、私自身はMetalを使って遊び始めるまで分かりませんでした。

設計上、GPUは大量の並列浮動小数点演算を効率的に行うことができ、高いメモリ帯域幅を要求します。私のMacBook M2には8つのCPUコアと8つのGPUコアがありますが、比較のために、NvidiaのRTX 4090には16384のコアが搭載されており、H100には16896のCUDAコアと数百の特殊なテンソルコアが搭載されています。通常、GPUはSIMD処理をサポートしており、同じ命令を複数のデータポイントに対して同時に実行することができます。

グラフィックス以外にも、行列の積や線形代数のタスク全般は、高度に並列化可能なアルゴリズムによってこの並列処理の恩恵を受けています。これにより、トレーニングや推論などの核となる機械学習の作業がサポートされています[1] [2]。

CUDAは、おそらく最もよく知られたGPUプログラミングプラットフォームで、Nvidiaのハードウェアに特化しています。 OpenGL用の数学フレームワークも利用可能です。TensorFlowやPyTorchなどのフレームワークは、GPUハードウェアと簡単かつ透明に統合することができます。Thisは、spaCy NLPライブラリにMetalベースのGPUフレームワークをサポートすることでパフォーマンスが向上した興味深い記事です。

メタルGPU基礎知識

直接GPU計算をプログラムすることは、オンデバイスのCPU用のコードを書くことほど簡単ではありません。AppleのMetalフレームワークを使用する場合、GPU上でコードを実行するための大まかな操作のシリーズは次のようになります：

適切なGPUデバイスを見つける
コマンドを実行するためのキューを作成する（つまり、MTLCommandQueue）
データ配列へのポインターを構造化バッファーにラップする。データが実行可能なコードである場合はパイプラインステート、さもなければ通常のバッファーになります。Apple GPUは統一メモリスペースを使用しているため、実際にデータをGPU固有の物理メモリにコピーする必要はありません
コマンドバッファをコミットして実行し、結果を待機するか、完了時にイベントハンドラを設定する
応答バッファーからバイトを抽出し、CPUのプログラムコードでローカルにフォーマットする

Raw GPUプログラミングは非同期モデルを使用します。

メタルシェーディング言語

Metalシェーディング言語は、Metal互換のGPU上で実行されるカスタムロジック（「シェーダー」と呼ばれる）を構築するために使用される、C++14の派生言語です。一般的には、可能な限りMPSフレームワーク（後で説明します）を使用した方が良いです。これは、通常、一般的なGPU対応の使用ケース（行列の乗算やニューラルネットワークなど）に最適化されています。

MSLコードのデバッグは非常に困難です。Xcodeを介してシェーダーデバッガーを使用することができますが、Xcodeなしで中間値を検査または表示する場合は、データを応答デバッグバッファに書き込み、C++またはObjective-Cのラッパーでプリミティブを解析する必要があります。

MSL関数は、kernel指定を介して公開インターフェースとして公開されます。Metalフレームワークは、現在の呼び出しスレッドのコンテキストまたはスレッドグループのためのIDを渡します。これは、重複しない書き込みを保証するために使用できます。スレッドは、三次元のIDシステムで表現することができます。このスレッドスペースの寸法は、ラッパーコードで設定されます。

以下は、単純な行列乗算アルゴリズムの実装です。いくつかのループ展開を組み合わせたもので、意外にもパフォーマンスが著しく向上しました。これは比較のためにのみです。通常は、MPSのMPSMatrixMultiplication機能の方が適しています。

kernel void matrix_multiply_naive(  device const MatrixParams *params,  constant float *A,  constant float *B,  device float *C,  // 実行されているスレッドの一意な位置を示す  // uint2型は2次元座標で、xとyのフィールドがそれぞれの軸のインデックスを表す  // このパラメータは呼び出し元のコードから直接提供されるわけではありませんが、Metalフレームワークから提供されます  uint2 gid [[thread_position_in_grid]]) {  if (gid.x >= params->a_rows || gid.y >= params->b_cols) {    return; // このスレッドは行列の次元範囲外なので何もしない  }  float sum = 0.0;  int k;  // ループの展開。パフォーマンスが大幅に改善されます  for (k = 0; k <= params->a_cols - 4; k += 4) {    sum += A[gid.x * params->a_cols + k]        * B[k * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 1]        * B[(k + 1) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 2]        * B[(k + 2) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 3]        * B[(k + 3) * params->b_cols + gid.y];  }  // 残りの要素を処理  for (; k < params->a_cols; ++k) {    sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];  }  C[gid.x * params->b_cols + gid.y] = sum;}

比較のため、MSLでnaive-transpose関数も実装しました。転置された行列が与えられた場合、このロジックへの微調整は簡単で、内部のループはBの列ではなく行に対して実行されます。

// ループの展開。パフォーマンスが大幅に改善されますfor (k = 0; k <= params->a_cols - 4; k += 4) {  sum += A[gid.x * params->a_cols + k]         * B[gid.y * params->b_cols + k]; // ここには gid.y * params->b_cols に k を加えた値が入ります  sum += A[gid.x * params->a_cols + k + 1]     * B[gid.y * params->b_cols + k + 1];  sum += A[gid.x * params->a_cols + k + 2]     * B[gid.y * params->b_cols + k + 2];  sum += A[gid.x * params->a_cols + k + 3]     * B[gid.y * params->b_cols + k + 3];}// 残りの要素を処理for (; k < params->a_cols; ++k) {  sum += A[gid.x * params->a_cols + k] * B[gid.y * params->b_cols + k];}

私は以前のブログ記事でもこのアプローチについて議論しました。これは、少なくともCPU上では単純なアルゴリズムのスカラーパフォーマンスを向上させるための非常に簡単な方法です。後で詳しく説明します。

Objective-Cバインディング

Metalフレームワークは、Metalソースコードからライブラリをコンパイルできる機能を提供します。ファイルの内容が読み込まれたら、バインディングコードではカーネル関数を名前で検索し、コンパイルされた関数コードを表す新しいMTLComputePipelineStateを初期化します。

id<MTLDevice> device = MTLCreateSystemDefaultDevice();// 提供されたソースパスにあるライブラリをコンパイルして初期化するMTLCompileOptions *compileOptions = [MTLCompileOptions new];compileOptions.languageVersion = MTLLanguageVersion3_0;// 入力ソースパスの文字列をラップNSString *ss = [NSString stringWithUTF8String:source_path];// コンパイルされたシェーダー関数を含む新しいライブラリを初期化するid<MTLLibrary> lib = [device newLibraryWithSource:ss  options:compileOptions  error:&error];// Metalライブラリ内のnaive乗算パブリックシェーダー関数を作成するid<MTLFunction> naiveFunction =    [lib newFunctionWithName:@"matrix_multiply_naive"];// 新しいコンピュートパイプラインステートを作成するid<MTLComputePipelineState> pipelineStateNaive = [device newComputePipelineStateWithFunction:naiveFunction  error:&error];

ネイティブなMetalコードを実際に呼び出すには、スレッド構成を設定する必要があります。そして、GPUバッファを初期化する必要があります。

[computeEncoder setComputePipelineState:pipelineStateNaive];
MTLSize threadsPerGrid = MTLSizeMake(params->a_cols, params->a_rows, 1);// スレッドグループのサイズを計算します。// https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes?language=objc
NSUInteger w = pipelineStateNaive.threadExecutionWidth;
NSUInteger h = pipelineStateNaive.maxTotalThreadsPerThreadgroup / w;
MTLSize threadsPerThreadgroup = MTLSizeMake(w, h, 1);// カーネル関数の入力をエンコードします。
[computeEncoder setBytes:params length:16 atIndex:0];
[computeEncoder setBuffer:bufferA offset:0 atIndex:1];
[computeEncoder setBuffer:bufferB offset:0 atIndex:2];
[computeEncoder setBuffer:bufferC offset:0 atIndex:3];// コンピュートコマンドをエンコードします。
[computeEncoder dispatchThreads:threadsPerGrid   threadsPerThreadgroup:threadsPerThreadgroup];// コンピュートパスを終了します。
[computeEncoder endEncoding];// コマンドを実行します。
[commandBuffer commit];

これはかなりの量ですので、関係性を以下に示します。

Objective-Cラッパー内のコンセプト、タイプ、ハードウェアのハイレベルレイアウト

Metal Performance Shadersフレームワーク

MPS Frameworkは、Appleが提供する高性能なライブラリであり、MetalファミリのGPUと一緒に使用するためのものです。それは画像処理からニューラルネットワークのサポートまで機能を提供しています。

APIは主にSwiftまたはObjective-Cを介して利用可能ですが、使用するためのMetal-cppライブラリもあります。

MPSMatrixMultiplication APIは比較的簡単に使用できます。上記のMSLコードと同様に、MPSコマンドはまだMTLCommandBufferにエンコードされ、非同期でコミットされて実行される必要があります。

// 行列の次元とバイトサイズを考慮した行列 "説明"を定義します
MPSMatrixDescriptor *descriptorA = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:a_cols  rowBytes:a_cols * sizeof(float)  dataType:MPSDataTypeFloat32];
MPSMatrixDescriptor *descriptorB = [MPSMatrixDescriptor matrixDescriptorWithDimensions:b_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// 出力行列
MPSMatrixDescriptor *descriptorC = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// 上記の説明と行列バッファを使用して行列表現を初期化します
MPSMatrix *matrixA = [[MPSMatrix alloc] initWithBuffer:bufferA descriptor:descriptorA];
MPSMatrix *matrixB = [[MPSMatrix alloc] initWithBuffer:bufferB descriptor:descriptorB];
MPSMatrix *matrixC = [[MPSMatrix alloc] initWithBuffer:bufferC descriptor:descriptorC];// 乗算インスタンスを作成します
MPSMatrixMultiplication *matrixMultiplication = [[MPSMatrixMultiplication alloc] initWithDevice:device  resultRows:a_rows  resultColumns:b_cols  interiorColumns:a_cols];// GPU用のコマンドバッファに乗算コマンドをエンコードします
id commandBuffer = [commandQueue commandBuffer];
[matrixMultiplication encodeToCommandBuffer:commandBuffer  leftMatrix:matrixA  rightMatrix:matrixB  resultMatrix:matrixC];

Goとcgo

私は特にObjective-Cでの作業は好きではありませんが、このプログラムのポイントは、GoプログラムからGPUで実行されるコードを実行することです。

Cgoは、ネイティブのCコードに関連するコメント内に含まれるコンパイラディレクティブをGoコンパイラが理解するためのGo言語の機能です。これは外部関数インターフェースのバージョンをサポートしています。

ディレクティブの設定は少し脆弱ですが、import "C"の直前のコメント（「前文」と呼ばれる）は、参照するCコードのコンパイル時にヘッダーインポートまたはコンパイル引数として解釈されます。例：

/*#cgo LDFLAGS: -framework Foundation -framework CoreGraphics -framework Metal -framework MetalPerformanceShaders -L/opt/homebrew/opt/openblas/lib -lopenblas#include <stdlib.h>#include "metal.h"*/import "C"

リンカーにリンクフラグをコマンドラインを介して渡す
Cコードを標準ヘッダーstdlib.hでコンパイルする
Cコードをプロジェクトのローカルヘッダーmetal.hでコンパイルする

MacOSで正しいリンカーフラグのセットを取得するために試行錯誤が必要でした。

Foundation：基本ライブラリ
CoreGraphics：GPUとのインターフェースを行うためにMacOSで必要
Metal：Metalのライブラリと言語サポート、MSLを含む
MetalPerformanceShaders：上記で説明したMPSのライブラリ

AppleはAccelerateフレームワークにBLAS実装をバンドルしているため、brewを介してOpenBLASをインストールするだけでなく、リンクする際にライブラリの場所も指定する必要があります：

-L/opt/homebrew/opt/openblas/lib -lopenblas

go:embedディレクティブを使用すると、Goプログラムはコンパイル時にファイルを含めることができます。これは、上述のように、MSLソースファイル（mm.metal）の内容をMetalフレームワークに渡してコンパイルする場合に便利です。

//go:embed mm.metalvar source string// シェーダーソースコードをコンパイルし、パイプラインを初期化します。 metalSource // パラメータには埋め込まれたMetalシェーディング言語ファイルの内容が含まれます。func Compile (metalSource string) { src := C.CString(metalSource) // 上記の文字列をコマンドキューが初期化された後に解放するdefer C.free(unsafe.Pointer(src)) C.initializePipelineAndCommandQueue(src)}

上記のCへの参照は、cgoを介してC APIと連携しています。例えば：

// Obj-CバインディングからinitializeMTLBuffersを呼び出すC.initializeMTLBuffers( a_data,                  // Aの入力不透明ポインタ b_data,                  // Bの入力不透明ポインタ C.int(4),                // 4をCの整数型に変換 C.int(a.Size()),          C.int(b.Size()),          C.int(a.Rows * b.Cols))params := MatrixParams{ a_rows: int32(a.Rows), a_cols: int32(a.Cols), b_rows: int32(b.Rows), b_cols: int32(b.Cols),}// このMatrixParams構造体へのポインタを、共有ヘッダーファイルで定義されたネイティブC表現にキャストして、不安全ポインターで返すreturn (*C.MatrixParams)(unsafe.Pointer(&params));

これにより、Cは予約語であり、変数名として使用することはできません。

Go実装ベースラインとOpenBLAS

GPUベースの行列乗算のパフォーマンスを、Gonumライブラリなどの高レベルの実装と、直感的で手書き（比較的効率が低い）実装と比較したかったです。

私はGoでいくつかの異なるアルゴリズムを実装しましたが、この並列転置ナイーブアルゴリズムは、乗算作業をN個のゴルーチンに分割するというナイーブな方法で実装されています：

func (a Matrix[T]) TransposeMultParallel(b *Matrix[T]) *Matrix[T] { if a.Cols != b.Rows {  panic("行列のサイズが間違っています") } c_data := make([]T, a.Rows*b.Cols) t := b.Transpose() var wg sync.WaitGroup for i := 0; i < a.Rows; i++ {  wg.Add(1) // 新しいゴルーチンのためにカウントをWaitGroupに追加  go func(i int) { // ゴルーチンを開始   defer wg.Done() // ゴルーチンが完了したらカウントを減らす   ptr := i * b.Cols   for j := 0; j < b.Cols; j++ {    var sum T = 0.0    for k := 0; k < a.Cols; k++ {     sum += a.At(i, k) * t.At(j, k)    }    c_data[ptr+j] = sum   }  }(i) } wg.Wait() // 全てのゴルーチンの完了を待つ return InitMatrixWithData(a.Rows, b.Cols, c_data)}

Gonum BLASは、BLASインターフェースを実装した純粋なGoライブラリです。ただし、OpenBLASなどのネイティブコードのBLAS実装にも設定できるようになっています。これは、OpenBLASをnetlibを介して使用するように設定することができます。

上記では、MacOS上のOpenBLASインストールに適切にリンクするようにcgoを設定する方法を示しました。アプリケーションコード内では、好みのBLAS実装を直接設定することもできます。ベンチマークコードから:

// プリミティブな配列をgonumの密行列型に変換gonum_a := mat.NewDense(a_rows, a_cols, a64_data)gonum_b := mat.NewDense(b_rows, b_cols, b64_data)gonum_c := mat.NewDense(a_rows, b_cols, nil)gonum_d := mat.NewDense(a_rows, b_cols, nil)// GonumがGonumのデフォルトGo実装を使用するように設定blas64.Use(gonum.Implementation{})// Gonum BLAS実装を使用した乗算を実行start = time.Now()gonum_c.Mul(gonum_a, gonum_b)bdata.TimeGonumNative(start)// Netlibを使用するようにGonumを設定。これにより、ネイティブのCコードBLAS実装(OpenBLAS)を使用するblas64.Use(netlib.Implementation{})// OpenBLAS実装を使用した乗算をGonum APIを介して実行start = time.Now()gonum_d.Mul(gonum_a, gonum_b)bdata.TimeGonumOpenBLAS(start)

結果

ベンチマークコードは、次の行列の乗算実装をいくつかの試行で実行し、次第に次元が増える正方行列の乗算にかかる平均時間を報告します:

- Naive multiplication, in Go- Transposed naive multiplication, in Go- Goroutine-parallelized transposed naive multiplication, in Go- Gonum pure Go-based BLAS multiplication- Gonum-wrapped OpenBLAS multiplication, written in C- Hand-implemented naive multiplication, in MSL, on GPU- Hand-implemented transposed naive multiplication, in MSL, on GPU- Metal Performance Shaders framework, called from Objective-C, on GPU

ベンチマークの出力は次のようになります（浮動小数点数はミリ秒を表しています）:

2023-12-01 11:12:51.644 go-mm[75818:22427382] デフォルトのデバイス Apple M2を使用していますelements naive transpose transpose_parallel metal_naive metal_transpose mps gonum openblas160000 196.00 201.00 42.00 8.00 9.67 0.33 4.67 6.00250000 381.33 387.67 80.67 11.00 11.67 0.00 8.33 21.00360000 801.00 789.33 159.33 19.00 16.33 0.00 14.33 4.67490000 1228.00 1075.00 411.00 23.67 24.33 1.00 26.67 16.33...

一部のクイックプロットはmatplotlibを通じて行われました