[JavaScript] WebGPU ベースの NDArray (多次元配列) を作って公開した
0. TL;DR
WebGPU ベースの NDArray をJavaScriptで実装して公開しました。
1. はじめに
以前 Vulkan API を使ってGPGPUするPythonライブラリ vulkpy を作って公開して、記事を書きましたが、記事でも言及のようにコンピュート・シェーダーの管理が大変であんまり進みませんでした。
いつの間にか Chromium で WebGPU が使えるようになっていた[1]ので、同じく GPGPU として NDArray (多次元配列) を作ってみました。
2. 使い方
READMEの例より
import { createGPU } from "https://cdn.jsdelivr.net/gh/ymd-h/gpu-array-js/gpu-array.js";
// 1. Create GPU. (Throw Error, if WebGPU is not available)
const gpu = await createGPU();
// 2. Create NDArray
const a = gpu.Array({ shape: [2, 2] });
const b = gpu.Array({ shape: [2, 2] });
// 3. Set Data
a.set(1, 0, 0); // Set 1 at (0, 0)
a.set(1, 1, 1); // Set 1 at (1, 1)
b.set(2, 0, 1); // Set 2 at (0, 1)
b.set(3, 1, 1); // Set 3 at (1, 1)
// 4. Execute Calculation.
// (If data is updated, automatically send to GPU)
const c = gpu.add(a, b); // c = a + b
// Optional: You can send data manually.
// a.send();
// 5. Get Data
// (If gpu data is updated, automatically load from GPU)
console.log(await c.get(0, 0));
console.log(await c.get(0, 1));
console.log(await c.get(1, 0));
console.log(await c.get(1, 1));
// Optional: You can load data manually.
// await c.load();
// console.log(c.get_without_load(0, 0));
// console.log(c.get_without_load(0, 1));
// console.log(c.get_without_load(1, 0));
// console.log(c.get_without_load(1, 1));
3. 実装
3.1 Templateベースのコンピュート・シェーダー
vulkpy の記事で言及のように真面目にいろんな計算を実装すると、型だけ違う、オペレータだけ違うというコンピュート・シェーダーが多数必要になります。
Vulkanではコンパイル済みのSPIR-Vと呼ばれる中間バイナリをシェーダーの入力に取りますが、WebGPUではテキストベースのWebGPU Shading Language (WGSL)を入力とするため、動的にパッチングする事が簡単でした。
シェーダーファイルは MIME Type "text/wgsl" で配信する事もできます[2]が、動的にパッチングしたいので、stringを返す関数としてJavaScriptの中に実装しました。
3.2 更新トラッキング
CPU側、GPU側の更新の有無をそれぞれ NDArray の .cpu_dirty, .gpu_dirty プロパティで管理して、更新があったときのみ send() / load() メソッドが実際にデータをCPU・GPU間でデータをコピーする事で、ユーザーはデータの転送を意識しなくてもいい(はず)。
4. まとめ
WebGPU でGPGPUするための NDArray を実装して公開しました。
WebGPU APIはVulkan APIによく似ているので、 vulkpy の開発での経験が生きました(が、やっぱりリソース管理等がわかりにくいですね)。
まだまだいろんな機能が足りないので、少しずつでも開発を進めていきたいなと思っています。
（今のところ問題はないので、今度こそ頓挫させないように。。。。）
興味持ってもらえたら、レポジトリにスターつけたり、触ってみたり、コードを読んでみたりしてもらえると嬉しいです。
参考
-
どうやら2023年5月のChrome 113。参照: https://ics.media/entry/230426/ ↩︎
-
https://gpuweb.github.io/gpuweb/wgsl/#text-wgsl-media-type ↩︎
Discussion