Common LispでApple SiliconのGPUを使う
cl-metal
Lisp Advent Calendar 2023 23日目の記事です。
Apple製のOSを搭載したMacやiPhoneでは，MetalというグラフフィックAPIを用いることができます。これはCUDAやOpenCLと同様にベクトル計算機として深層学習モデルの計算などの科学計算にも応用することができます。
最近のMacはユニファイドメモリを採用していて，CPUに割り当てた配列のポインタをそのままGPUで実行して並列計算することができるというのは大きな強みです。というのは，一度Metalで計算の過程を記述してコンパイルしておけば，使いたい配列をCommon Lispで割り当ててCFFIを噛ませるだけなのでシームレスに連携ができるからなのですが，それでもMetal側のめんどくさい仕様やBufferの管理をどこかでやる必要がありました。
そういったライブラリはJuliaだとかPythonみたいなイケてる言語には当然用意されているのですが，例の如くCommon Lispでは誰もやってなかったので，最近作っています。
使い方
cl-metalを用いると，以下のようなLisp風DSLを直接Common Lispコードに埋め込むことができます。これらのコードはマクロなのでコードを静的に解析してC++のコードを吐き出します。(以下./examples/mandelbrot.lispから抜粋)
(define-mfunc (aux)
(uint8-t ((x float :in) (y float :in) (a float :in) (b float :in) (m uint8-t :in)))
(if (< m 100)
(let ((x1 (- (* x x) (* y y) a))
(y1 (- (* 2.0 x y) b)))
(if (> (+ (* x1 x1) (* y1 y1)) 4.0)
m
(aux x1 y1 a b (+ m 1))))
0))
(define-kernel (mandelbrot
:thread-position-in-grid id
:using (aux))
(void ((x* uint8-t :out)))
(let ((a (/ (- (mod id 2048) 512.0) 1024.0))
(b (/ (- (/ id 2048) 1024.0) 1024.0)))
(setf (aref x id) (aux 0.0 0.0 a b 1))))
(参考) 生成されたC++のコード
#include <metal_stdlib>
using namespace metal;
static inline uint8_t aux(const float x, const float y, const float a, const float b, const uint8_t m) {
if (m < 100) {auto x1 = (((x * x) - (y * y)) - a);
auto y1 = (((2.0 * x) * y) - b);
{
if (((x1 * x1) + (y1 * y1)) > 4.0) {return m;} else {return aux(x1, y1, a, b, (m + 1));}
}} else {return 0;}
}
kernel void mandelbrot(device uint8_t *x [[ buffer(0) ]], uint id [[ thread_position_in_grid ]]) {
auto a = (((id % 2048) - 512.0) / 1024.0);
auto b = (((id / 2048) - 1024.0) / 1024.0);
{
x[id] = aux(0.0, 0.0, a, b, 1);
}
}
生成したC++のコードは更にXcodeを介してコンパイルする必要があるのですが，cl-metalではその手順をtoplevel-formに記述しています。なので，REPL上でC-c C-cしてdefine-kernel ...を実行するとその場で↓みたいにコンパイルエラーがわかりますし，
REPL
./.cl_metal_tmpディレクトリにコンパイルされたシェーダーのキャッシュを作成するので，一度LispのファイルをLoadすればそれ以降はコンパイルによるオーバーヘッドも発生しません。こんな具合でエディターやCommon LispとMetalをうまく統合しながら快適にプログラミングができます。
最初に述べた通り，コンパイルしたMetalのシェーダーにCommon Lisp上で割り当てた配列をそのまま渡して計算することができます。試しにCPUとGPUでMandelbrot集合を描画するコードを比べてみると:
おまけ: CPUバージョン mandelbrot-cpu
#+sbcl(setf sb-ext:*inline-expansion-limit* 20)
(defun mandelbrot-cpu (xs)
(declare (type (simple-array (unsigned-byte 8) (*)) xs)
(optimize (speed 3) (safety 0)))
(dotimes (i #.(* 2048 2048))
(labels ((aux (x y a b m)
(declare (type single-float x y a b)
(type (unsigned-byte 8) m))
(if (< m 100)
(let ((x1 (- (* x x) (* y y) a))
(y1 (- (* 2.0 x y) b)))
(if (> (+ (* x1 x1) (* y1 y1)) 4.0)
m
(aux x1 y1 a b (+ m 1))))
0)))
#+sbcl(declare (inline aux))
(let ((a (/ (- (mod i 2048) 512) 1024.0))
(b (/ (- (/ i 2048.0) 1024) 1024.0)))
(setf (aref xs i) (aux 0.0 0.0 a b 1))))))
CPUとGPUを使うときの差異はmandelbrot-cpu関数を呼び出すか，mandelbrotを呼び出すかだけです。
(defun draw-mandelbrot (pathname xs)
(with-open-file (out pathname :direction :output
:if-does-not-exist :create
:if-exists :supersede)
(write-line "P2" out)
(write-line "2048 2048" out)
(write-line "255" out)
(dotimes (i (* 2048 2048))
(princ (min 255 (* 8 (aref xs i))) out)
(terpri out))))
(defun main ()
(let ((xs (make-array (* 2048 2048)
:initial-element 0
:element-type '(unsigned-byte 8))))
(time (mandelbrot-cpu xs)) ;; <- CPU
(time (mandelbrot xs)) ;; <- cl-metal
(draw-mandelbrot #P"./mandelbrot.pgm" xs)))
mainを叩くと./mandelbrot.pgmにそれっぽい画像が吐き出されます:
Result
並列化の恩恵が大きい今回の場合だとCPU比で67倍の高速化に成功します
MANDELBROT> (main)
Evaluation took:
0.402 seconds of real time
0.401794 seconds of total run time (0.400975 user, 0.000819 system)
100.00% CPU
0 bytes consed
Evaluation took:
0.006 seconds of real time
0.003584 seconds of total run time (0.000555 user, 0.003029 system)
66.67% CPU
0 bytes consed
パフォーマンスについて
まぁただいくらライブラリ側での最適化を頑張っても， 外部言語を使う以上どうしても不便な箇所はあったりします。
(当然)SBCLレベルでのインライン化とかは難しくて，Buffer作成のオーバーヘッドなどもあるので小さいスケールでしょうもない関数をたくさん呼び出したりすると，CPUより遅くなったりします。(N=256でやっと逆転する)
Benchmarking on: sin (N=100):
scale=2 | SBCL(CPU) -> (1.0e-5) | Metal -> (0.01027)
scale=4 | SBCL(CPU) -> (5.0e-6) | Metal -> (0.007887)
scale=8 | SBCL(CPU) -> (1.4e-5) | Metal -> (0.008005)
scale=16 | SBCL(CPU) -> (4.9e-5) | Metal -> (0.006916)
scale=32 | SBCL(CPU) -> (1.93e-4) | Metal -> (0.006828)
scale=64 | SBCL(CPU) -> (7.67e-4) | Metal -> (0.007558)
scale=128 | SBCL(CPU) -> (0.003059) | Metal -> (0.008098)
scale=256 | SBCL(CPU) -> (0.013527) | Metal -> (0.01089)
scale=512 | SBCL(CPU) -> (0.053579) | Metal -> (0.018506)
scale=1024 | SBCL(CPU) -> (0.220899) | Metal -> (0.048736)
scale=2048 | SBCL(CPU) -> (0.885407) | Metal -> (0.222491)
これを避けるために 1.「いくつかの命令を融合したでっかい関数を」2.「なるべくでかい行列で」3.「なるべく一回の呼び出しで」実行するなんてことをコード書く時に意識するのですが，それを手動でやるのはnot my cup of teaだと思ってるので，将来的にcl-waffe2のJITに組み込んで使うのが一番楽かなぁと考えています。余談ですが今度時間がある時にcl-waffe2でStable Diffusionの推論を実装してみようと考えていて，その時にcl-metalを使おうかななんて考えています。
cl-metalもまだ実装が雑なので実運用はやめておいた方がいいと思いますが，もしよかったらContributionしてくれると僕が喜びます。
Discussion