GPU.js、普通のJavaScript関数を最大370倍速くする仕組み

はじめに

大量のループ処理や行列演算をJavaScriptで書くと、どうしてもCPUの限界にぶつかります。Web Workerで並列化しても、コア数以上には速くなりません。ではブラウザが持っているもう1つの計算資源、GPUを使えないでしょうか。

GPU.jsは、普通に書いたJavaScript関数を自動的にシェーダー言語へトランスパイルし、GPU上で並列実行してくれるライブラリです。WebGLの知識がなくても、for文とプレーンな関数を書くだけでGPUの並列計算力を引き出せます。GPUが使えない環境では自動的にCPUへフォールバックするため、動作しない心配もありません。

とはいえ、読むより触った方が早いと思います。ブラウザ上でGPU.jsのカーネルを実行し、配列サイズと倍率を変えて結果を確認できるサンプルを用意しました。先に挙動を見たい方はこちらからどうぞ。

＞＞＞ まずはサンプルを確認 ＜＜＜

GPU.jsとは

GPU.jsは「GPU Accelerated JavaScript」を掲げるGPGPU（GPU上での汎用計算）ライブラリです。ブラウザとNode.jsの両方で動作し、WebGPU・WebGL2・WebGL・HeadlessGL・WebAssembly・CPUという複数のバックエンドを内部で使い分けます。実行環境にGPUがなければ自動的にCPU実行に切り替わるため、動作環境を選ばずに導入できるのが大きな強みです。

主な特徴

• JavaScript関数をそのままGPUで実行 - createKernel()に渡した関数がシェーダーコードへ自動変換されます。WebGLの知識は不要です
• 大幅な高速化 - 公式ベンチマークでは行列演算などでCPU比最大370倍程度の高速化が報告されています(処理内容やハードウェアに依存)
• 自動フォールバック - GPUが使えない環境ではCPUで同じ関数を実行するため、実行不能になることがありません
• グラフィック出力にも対応 - setGraphical(true)を使うと、計算結果をそのままCanvasに描画できます
• TypeScript対応 - 型定義が同梱されており、IKernelFunctionThisなどで関数内のthisも型付けできます

インストール

npmまたはyarnで導入します。

npm install gpu.js --save

yarn add gpu.js

Node.jsではheadless-gl経由でGPU処理を実行し、ブラウザではdist/gpu-browser.jsのビルドがそのまま使われます。

GPU.jsのサンプルを動かす

GPU.jsの中心となるAPIはgpu.createKernel()です。渡した関数の中ではthis.thread.xで現在処理中のインデックスにアクセスでき、.setOutput([n])で出力する配列サイズを指定します。下のサンプルでは、配列サイズと掛ける数値を入力すると、その場でGPU.jsのカーネルが再実行されて結果が更新されます。

import { GPU } from 'gpu.js'

const gpu = new GPU()

const kernel = gpu.createKernel(function (multiplier) {
  return this.thread.x * multiplier
}).setOutput([10])

const result = kernel(3)
// Float32Array [0, 3, 6, 9, 12, ...]

配列サイズを500まで増やしても瞬時に結果が返るのが分かるはずです。this.thread.xはGPU上の各スレッドが担当するインデックスを表しており、GPU.jsはこの関数をスレッド数ぶん並列に実行しています。掛ける数を負の値にすれば、結果がすべてマイナスに反転することも確認できます。

基本的な使い方

GPU.jsを使う最小構成は次の3ステップです。

const { GPU } = require('gpu.js') // ブラウザではimport { GPU } from 'gpu.js'

// 1. GPUインスタンスを作成
const gpu = new GPU()

// 2. カーネル関数を定義し、出力サイズを指定
const kernel = gpu.createKernel(function () {
  return this.thread.x
}).setOutput([100])

// 3. 実行する
kernel()
// Result: Float32Array [0, 1, 2, 3, ... 99]

createKernel()に渡す関数の中では、通常のJavaScriptと同じようにfor文やMath関数、配列アクセスが使えます。ただし関数内は最終的にシェーダーコードへ変換されるため、クロージャで外部の複雑なオブジェクトを参照することはできません。必要な値は引数として渡すのが基本です。

実践的なユースケース

GPU.jsの使い道は、単純な配列処理だけではありません。計算負荷の高い処理・グラフィック描画・UIをブロックしない非同期実行という3つの代表的なパターンを見ていきます。

行列演算をGPUで並列化する

GPU.jsが最も威力を発揮するのが、ネストしたループを伴う行列演算です。CPUだとO(n^3)の計算量になる行列の掛け算も、GPU.jsなら各出力要素の計算を別々のGPUスレッドに割り振れます。下のサンプルでは行列サイズNを変更して、GPUとCPUそれぞれの実行時間をperformance.now()で計測し比較できます。

const multiplyMatrix = gpu.createKernel(function (a, b, n) {
  let sum = 0
  for (let i = 0; i < n; i++) {
    sum += a[this.thread.y][i] * b[i][this.thread.x]
  }
  return sum
}).setOutput([n, n])

const c = multiplyMatrix(a, b, n)

Nを100や150まで大きくすると、GPU.jsのmultiplyMatrixと素のJavaScriptによる3重ループとの差が広がっていくはずです。カーネルの生成自体にコストがかかるため、この例では入力ごとの即時再計算ではなく「計算を実行」ボタンで測定するようにしています。

setGraphical(true)でGPU計算結果をそのまま画像として描画する

GPU.jsは数値配列だけでなく、Canvasへの描画にも対応しています。.setGraphical(true)を付けたカーネルの中でthis.color(r, g, b, a)を呼ぶと、各スレッドが担当するピクセルの色をGPU上で直接計算できます。下のサンプルはスライダーで周波数を変えると、GPU.jsが生成するプラズマ状のパターンがリアルタイムに変化します。

const render = gpu.createKernel(function (freq) {
  const x = this.thread.x / this.output.x
  const y = this.thread.y / this.output.y
  const v = Math.sin(x * freq) * Math.cos(y * freq)
  this.color(v, 0.2, 1 - v, 1)
}).setOutput([200, 200]).setGraphical(true)

render(10)
document.body.appendChild(render.canvas) // 計算結果がそのままCanvasに描画される

this.output.xとthis.output.yはカーネルの出力サイズ（ここでは200×200）を表す組み込み変数です。ピクセルごとの色計算はすべてGPU側の並列実行になるため、周波数を動かしてもパターンの再描画が滑らかに追従します。

非同期モードでUIをブロックしない

GPU.jsはnew GPU({ mode: 'async' })で非同期モードを有効にでき、カーネル呼び出しがPromiseを返すようになります。計算が終わるまでawaitで待つ形になるため、重い処理中でも他のUI更新を止めずに済みます。下のサンプルでは、ボタンでカーネルを非同期実行している間も、右側のtickカウンターが100msごとに動き続けることを確認できます。

const gpu = new GPU({ mode: 'async' })

const kernel = gpu.createKernel(function () {
  let sum = 0
  for (let i = 0; i < 1000; i++) sum += Math.sqrt(i + this.thread.x)
  return sum
}).setOutput([2000])

const result = await kernel() // Promiseとして結果を待つ

「非同期で実行」を押している間もtickカウンターが止まらずに増え続ける点に注目してください。同期モード（デフォルト）だとGPU.jsの計算自体はGPU側で並列実行されますが、mode: 'async'にすることで結果取得をawaitベースの非同期処理として扱え、他の非同期処理と組み合わせやすくなります。

まとめ

GPU.jsは、WebGLのシェーダー言語を意識することなく、普通のJavaScript関数のままGPUの並列計算力を引き出せるライブラリです。createKernel()とスレッドインデックス（this.thread.xなど)という最小限のAPIを覚えるだけで、行列演算のような重い計算から、setGraphical(true)によるグラフィック生成、mode: 'async'を使った非同期実行まで幅広くカバーできます。GPUが使えない環境でも自動的にCPUへフォールバックするため、「動かない」という事故が起きにくいのも安心材料です。

大量データの集計やシミュレーション、Canvasを使ったビジュアル表現など、CPUの計算がボトルネックになっている処理があれば、GPU.jsで置き換えられないか試してみる価値があります。

参考リンク

• GPU.js公式サイト
• GPU.js GitHub
• gpu.js - npm