Babylon.jsでWebGPUのコンピュートシェーダを使い、リアルタイムピクセルソートをやってみた

※この記事は人間が書く→AIで清書→人間がチェックして品質担保しています。なんか間違ってるとこあったら言ってね。

はじめに

Babylon.jsとWebGPUのコンピュートシェーダを使って、動画に対してリアルタイムでピクセルソートエフェクトを適用するWebアプリケーションを作成しました。 Babylon.jsもWebGPUも初心者でしたが、なんとか実装できたので、簡略化した実装周りのコツを共有します。

筆者の経験レベル

• Three.jsは慣れている
• GLSLは分かる
• WebGPUは全く触ったことがない（WGSLってなに...？）
• Babylon.jsも全く触ったことがない（Three.jsと何が違うん...？）

完成したものの見た目

まず、完成したものの見た目をご覧ください。

動画に対してリアルタイムでエフェクトが適用されているのが分かると思います。 あと120fps出てます。

そもそもピクセルソートとは

キム・アセンドルフ（Kim Asendorf）による画像処理・演出手法で、多くのアート作品で使われています。

参考：

今回、彼の実装に倣い、画像の各行や各列のピクセルを輝度値に基づいてソートする画像処理エフェクトを実装しようと考えました。

フラグメントシェーダで組むの辛くない？

最初はGLSLのフラグメントシェーダで実装しようと考えましたが、1ピクセルごとに処理するフラグメントシェーダでピクセルソートを実装するのは非効率的で、かなりトリッキーな実装を求められると感じました。 また、リアルタイムに変化するピクセルソートを作りたいと考えていたので、処理が重くなる実装になりそうな要因は避けたいと思いました。

じゃあ何を使うの？

そこで、GPUを使ってリアルタイム性を担保し、高速な計算処理ができるGPGPUを思い出しました。 （このセッションを見ておいて良かったです：Frontend Conf Hokkaido 2025）

ただし、WebGL2のGPGPUではCompute Shaderが使えないため、WebGPUを使うことにしました。 （その他にも、WebGPUの方がGPGPUをより効率的に使える仕組みや機能が揃えられている点も決め手でした。）

しかし、WebGPUは全く触ったことがなく、WebGLの生実装と同様のつらさを感じたため（本当はちゃんと学んだ方が良いですが）、Babylon.jsでラップされたWebGPUを使うことにしました。

次項で実装の全体像を説明します。

実装の全体像

実装は大きく以下の3つのステップで構成されています：

1. 元の動画フレーム（texture_2d<f32>）をストレージテクスチャ（texture_storage_2d<rgba8unorm, write>）にコピー
2. 動画の各行（または各列）を横幅いっぱいのセグメントとして生成（CPU側の処理）
3. 各セグメントに対してBitonic Sortを実行

ステップ1と3はWebGPUのコンピュートシェーダで並列処理され、ステップ2はCPU側で処理されます。これによりリアルタイムでエフェクトを適用できます。

この全体像を実現するために、どのようなコード構成になっているか説明します。

プロジェクト構成

プロジェクトはざっくり以下のような構成になっています：

src/
├── main.ts                    # メインアプリケーションロジック
├── core/                      # コア機能
│   ├── initialization.ts      # エンジン・シーンの初期化
│   ├── textures.ts            # テクスチャ作成
│   ├── plane.ts               # プレーン作成
│   └── uniformBuffers.ts     # ユニフォームバッファ管理
├── resources/                 # リソース管理
│   ├── pixelSortResources.ts # ピクセルソートリソース
│   └── segmentResources.ts   # セグメントリソース
├── shaders/                   # WebGPUコンピュートシェーダ
│   ├── pixelSort.compute.ts   # ピクセルソート処理
│   ├── pixelSort.compute.wgsl # ピクセルソートシェーダー
│   └── copyFrame.compute.ts   # フレームコピー
|   └── copyFrame.compute.wgsl # フレームコピーシェーダー
├── compute/                   # コンピュートシェーダーまとめて作成
│   └── shaders.ts
└── ui/                        # UI関連
    └── controls.ts             # コントロールパネル

この構成の中で、特に重要なのがコンピュートシェーダの実装です。このプロジェクトでは、2つのコンピュートシェーダを使用しています。次項で各シェーダがどのように動作するか説明します。

コンピュートシェーダの実装

このプロジェクトでは、以下の2つのコンピュートシェーダを使用しています：

1. フレームコピーシェーダ (copyFrame.compute.wgsl): 動画テクスチャ（texture_2d<f32>）をストレージテクスチャ（texture_storage_2d<rgba8unorm, write>）にコピー
2. ピクセルソートシェーダ (pixelSort.compute.wgsl): セグメントごとにピクセルをソート

コンピュートシェーダの実装に入る前に、理解しておくべき基本的な概念を説明します。これらを理解しておくと、後で見るWGSLコードがより理解しやすくなります。

コンピュートシェーダの基本概念

ワークグループとスレッド

コンピュートシェーダは、ワークグループという単位で並列実行されます。各ワークグループは複数のスレッド（処理単位）で構成され、これらのスレッドが同時に実行されます。

例えば、@workgroup_size(16, 16, 1)と指定すると、1つのワークグループに16×16=256個のスレッドが含まれます。各スレッドは独立して処理を実行し、同じコードを異なるデータに対して実行します。

ワークグループのサイズの限界は、論理デバイスのlimitsから取得できます。 これを知るためには下記のようなコードを実行してみましょう。

const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();

const limits = device.limits;

console.log(limits);

グローバルIDとローカルID

各スレッドは以下のIDを持っています：

• グローバルID (global_invocation_id): 全体のスレッドの中での位置
• ローカルID (local_invocation_id): ワークグループ内での位置
• ワークグループID (workgroup_id): ワークグループの位置

これらのIDを使って、どのスレッドがどのデータを処理するかを決定します。

ストレージテクスチャとバッファ

コンピュートシェーダでは、以下のリソースを使用します：

• テクスチャ (texture_2d<f32>): 読み取り専用の画像データ
• ストレージテクスチャ (texture_storage_2d<rgba8unorm, write>): 書き込み可能な画像データ
• ストレージバッファ (var<storage, read>): 読み書き可能な構造化データ

ピクセルソートでは、入力テクスチャから読み取り、ストレージテクスチャに書き込みます。

ワークグループ共有メモリ

var<workgroup>で宣言された変数は、同じワークグループ内のスレッド間で共有されます。これは、グローバルメモリよりも高速にアクセスできるため、Bitonic Sortのような並列アルゴリズムで重要な役割を果たします。

WGSLの基本構文

WGSL（WebGPU Shading Language）は、WebGPU用のシェーダー言語です。GLSLと似ていますが、いくつか違いがあります：

• 型システムが厳格（u32, f32, vec3<f32>など）
• @で始まるアノテーション（@compute, @bindingなど）
• letとvarの使い分け（letは不変、varは可変）

これらの概念を理解した上で、実際の実装を見ていきましょう。

1. フレームコピー（copyFrame.compute.wgsl）

ピクセルソート処理の前に、まず元の動画フレームを出力テクスチャにコピーします。後続のピクセルソート処理のベース画像となります：

@group(0) @binding(0) var inputTexture : texture_2d<f32>;
@group(0) @binding(1) var outputTexture : texture_storage_2d<rgba8unorm, write>;

@compute @workgroup_size(16, 16, 1)
fn main(@builtin(global_invocation_id) globalId : vec3<u32>) {
    let dims = textureDimensions(inputTexture);
    if (globalId.x >= dims.x || globalId.y >= dims.y) {
        return;
    }

    let coords = vec2<i32>(i32(globalId.x), i32(globalId.y));
    let color = textureLoad(inputTexture, coords, 0);
    textureStore(outputTexture, coords, color);
}

2. セグメントの生成（CPU側の処理）

次に、ピクセルソートを適用する範囲を定義するため、動画の各行（または各列）を横幅いっぱいのセグメントとして生成します。これはコンピュートシェーダではなく、CPU側（TypeScript）で処理されます。セグメントは以下の構造で表現されます：

struct Segment {
    row: u32;      // 行番号（または列番号）
    start: u32;    // セグメントの開始位置（常に0）
    length: u32;   // セグメントの長さ（横幅または縦幅）
    reserved: u32; // 予約フィールド
}

横方向のソートの場合、各行が1つのセグメントとなり、セグメント数は動画の高さと等しくなります。

export function populateFullWidthSegments(
  resources: SegmentExtractionResources,
  width: number,
  height: number,
  orientation: number,
): number {
  const fullSegmentCount = orientation === 0 ? height : width;
  const segmentLength = orientation === 0 ? width : height;
  const segmentCount = Math.min(fullSegmentCount, resources.maxSegments);
  
  const buffer = ensureFullWidthSegmentsBuffer(resources, segmentCount);
  for (let line = 0; line < segmentCount; line++) {
    const base = line * SEGMENT_COMPONENTS;
    buffer[base] = line;        // row
    buffer[base + 1] = 0;       // start
    buffer[base + 2] = segmentLength; // length
    buffer[base + 3] = 0;       // reserved
  }
  
  resources.segmentsBuffer.update(view);
  return segmentCount;
}

3. ピクセルソート（Bitonic Sort）

最後に、生成したセグメントごとにBitonic Sortを実行してピクセルをソートします。Bitonic Sortは並列ソートアルゴリズムの一種で、今回の用途に（多分）向いています。

処理の流れは以下の通りです：

1. ピクセルの読み込み: セグメント内のピクセルをワークグループ共有メモリに読み込む
2. 輝度の計算: 各ピクセルの輝度値を計算し、閾値以上のピクセルのみをソート対象とする
3. Bitonic Sort: ワークグループ共有メモリ上でBitonic Sortを実行
4. 結果の書き込み: ソート済みのピクセルを出力テクスチャに書き込む

struct Params {
    width : u32,
    height : u32,
    maxSegmentLength : u32,
    sortDirection : u32,
    threshold : f32,
    segmentCount : u32,
    sortOrientation : u32,
    padding0 : f32,
};

struct Segment {
    row : u32,
    start : u32,
    length : u32,
    reserved : u32,
};

@group(0) @binding(0) var inputTexture : texture_2d<f32>;
@group(0) @binding(1) var outputTexture : texture_storage_2d<rgba8unorm, write>;
@group(0) @binding(2) var<uniform> params : Params;
@group(0) @binding(3) var<storage, read> segments : array<Segment>;

struct SharedPixel {
    colorPacked : u32,
    brightness : f32,
};

var<workgroup> sharedPixels : array<SharedPixel, 2048>;

fn luminance(color : vec3<f32>) -> f32 {
    return dot(color, vec3<f32>(0.299, 0.587, 0.114));
}

fn sentinelValue(sortDirection : u32) -> f32 {
    if (sortDirection == 0u) {
        return -1e6;
    }
    return 1e6;
}

@compute @workgroup_size(128, 1, 1)
fn main(
    @builtin(global_invocation_id) globalId : vec3<u32>,
    @builtin(local_invocation_id) localId : vec3<u32>,
    @builtin(workgroup_id) groupId : vec3<u32>
) {
    let segmentIndex = groupId.x;
    if (segmentIndex >= params.segmentCount) {
        return;
    }

    let segment = segments[segmentIndex];
    let segmentStart = segment.start;
    let orientation = params.sortOrientation;
    var axisLimit = params.width;
    var crossLimit = params.height;

    // 横方向・縦方向の切り替え
    if (orientation == 0u) {
        axisLimit = params.width;
        crossLimit = params.height;
        if (segment.row >= crossLimit) {
            return;
        }
    } else {
        axisLimit = params.height;
        crossLimit = params.width;
        if (segment.row >= crossLimit) {
            return;
        }
    }

    // セグメント長の計算（2のべき乗に調整）
    let segmentLengthComputed = min(segment.length, params.maxSegmentLength);
    let availableLength = axisLimit - segmentStart;
    let segmentLengthActual = min(segmentLengthComputed, availableLength);

    if (segmentLengthActual == 0u) {
        return;
    }

    let line = segment.row;

    // Bitonic Sortは2のべき乗の長さで効率的に動作するため、最大長を2のべき乗に調整
    var maxLength = 1u;
    loop {
        if (maxLength >= segmentLengthActual || maxLength >= params.maxSegmentLength) {
            break;
        }
        maxLength = maxLength << 1u;
    }
    maxLength = clamp(maxLength, 1u, params.maxSegmentLength);
    if (maxLength < segmentLengthActual) {
        maxLength = params.maxSegmentLength;
    }

    let sentinel = sentinelValue(params.sortDirection);

    // ステップ1: ピクセルの読み込みと輝度の計算
    var loadIndex = localId.x;
    loop {
        if (loadIndex >= maxLength) {
            break;
        }

        let withinSegment = loadIndex < segmentLengthActual;
        var color = vec4<f32>(0.0);
        var brightness = sentinel;

        if (withinSegment) {
            // 横方向または縦方向に応じて座標を計算
            if (orientation == 0u) {
                let x = segmentStart + loadIndex;
                if (x < params.width) {
                    let source = textureLoad(inputTexture, vec2<i32>(i32(x), i32(line)), 0);
                    color = source;
                    let value = luminance(color.rgb);
                    // 閾値以上のピクセルのみをソート対象とする
                    if (value >= params.threshold) {
                        brightness = value;
                    }
                }
            } else {
                let y = segmentStart + loadIndex;
                if (y < params.height) {
                    let source = textureLoad(inputTexture, vec2<i32>(i32(line), i32(y)), 0);
                    color = source;
                    let value = luminance(color.rgb);
                    if (value >= params.threshold) {
                        brightness = value;
                    }
                }
            }
        } else {
            // 範囲外のピクセルをマークするため、アルファを0にする
            color.a = 0.0;
        }

        sharedPixels[loadIndex].colorPacked = pack4x8unorm(color);
        sharedPixels[loadIndex].brightness = brightness;

        loadIndex = loadIndex + 128u;
    }

    // すべてのスレッドが読み込みを完了するまで待機
    workgroupBarrier();

    // ステップ2: Bitonic Sortの実行
    var k = 2u;
    loop {
        if (k > maxLength) {
            break;
        }

        var j = k / 2u;
        loop {
            var index = localId.x;
            loop {
                if (index >= maxLength) {
                    break;
                }

                let ix = index;
                let ixj = ix ^ j;  // XOR演算で比較対象のインデックスを取得

                if (ixj > ix && ixj < maxLength) {
                    let ascending = (ix & k) == 0u;  // 昇順か降順かを判定
                    let brightnessIx = sharedPixels[ix].brightness;
                    let brightnessIxj = sharedPixels[ixj].brightness;
                    var shouldSwap = false;

                    // ソート方向に応じてスワップ条件を決定
                    if (params.sortDirection == 0u) {
                        // 降順ソート
                        if (ascending) {
                            shouldSwap = brightnessIx < brightnessIxj;
                        } else {
                            shouldSwap = brightnessIx > brightnessIxj;
                        }
                    } else {
                        // 昇順ソート
                        if (ascending) {
                            shouldSwap = brightnessIx > brightnessIxj;
                        } else {
                            shouldSwap = brightnessIx < brightnessIxj;
                        }
                    }

                    // スワップ処理
                    if (shouldSwap) {
                        let tmpColor = sharedPixels[ix].colorPacked;
                        let tmpBrightness = brightnessIx;
                        sharedPixels[ix].colorPacked = sharedPixels[ixj].colorPacked;
                        sharedPixels[ix].brightness = brightnessIxj;
                        sharedPixels[ixj].colorPacked = tmpColor;
                        sharedPixels[ixj].brightness = tmpBrightness;
                    }
                }

                index = index + 128u;
            }

            workgroupBarrier();

            if (j <= 1u) {
                break;
            }
            j = j / 2u;
        }

        k = k * 2u;
    }

    // ステップ3: ソート済みのピクセルを出力テクスチャに書き込む
    var storeIndex = localId.x;
    loop {
        if (storeIndex >= segmentLengthActual) {
            break;
        }

        if (orientation == 0u) {
            let x = segmentStart + storeIndex;
            if (x < params.width) {
                let sortedColor = unpack4x8unorm(sharedPixels[storeIndex].colorPacked);
                // 範囲外のピクセル（アルファが0）は元のテクスチャから読み込む
                let isOutOfRange = sortedColor.a == 0.0;
                if (isOutOfRange) {
                    let originalColor = textureLoad(inputTexture, vec2<i32>(i32(x), i32(line)), 0);
                    textureStore(outputTexture, vec2<i32>(i32(x), i32(line)), originalColor);
                } else {
                    textureStore(outputTexture, vec2<i32>(i32(x), i32(line)), sortedColor);
                }
            }
        } else {
            let y = segmentStart + storeIndex;
            if (y < params.height) {
                let sortedColor = unpack4x8unorm(sharedPixels[storeIndex].colorPacked);
                let isOutOfRange = sortedColor.a == 0.0;
                if (isOutOfRange) {
                    let originalColor = textureLoad(inputTexture, vec2<i32>(i32(line), i32(y)), 0);
                    textureStore(outputTexture, vec2<i32>(i32(line), i32(y)), originalColor);
                } else {
                    textureStore(outputTexture, vec2<i32>(i32(line), i32(y)), sortedColor);
                }
            }
        }

        storeIndex = storeIndex + 128u;
    }
}

Bitonic Sortアルゴリズム

Bitonic Sortは、以下の手順でソートを実行します：

1. 段階的な比較: kを2から始めて、2倍ずつ増やしながら比較範囲を広げます
2. ペアの比較: ix ^ jで比較対象のペアを決定し、昇順/降順に応じてスワップします
3. 同期: 各比較段階の後にworkgroupBarrier()で同期を取り、すべてのスレッドが完了するまで待機します

Bitonic Sortの特徴は、比較とスワップが完全に並列化できるため、GPU上で非常に効率的に実行できることです。

結果の書き込み

• 範囲外の処理: アルファが0のピクセル（範囲外）は、元のテクスチャから読み込んでそのまま書き込みます
• ソート済みピクセルの書き込み: ソート済みのピクセルを出力テクスチャに書き込みます

これまで、2つのコンピュートシェーダ（フレームコピーとピクセルソート）と、CPU側でのセグメント生成処理を見てきました。ただし、シェーダーコードを書いただけでは動きません。Babylon.jsを使って、これらのシェーダーを実際に実行するための準備が必要です。

Babylon.jsでの実装

Babylon.jsでは、WGSLシェーダーを実行するために、エンジンの初期化、コンピュートシェーダの作成、テクスチャやバッファの準備など、いくつかのステップが必要です。順番に見ていきましょう。

エンジンとシーンの初期化

まず、WebGPUエンジンとシーンを初期化します：

export async function createEngine(canvas: HTMLCanvasElement): Promise<WebGPUEngine> {
  const engine = new WebGPUEngine(canvas, {
    adaptToDeviceRatio: true,
    antialias: true,
  });

  await engine.initAsync();
  return engine;
}

export function createScene(engine: WebGPUEngine): Scene {
  const scene = new Scene(engine);
  
  new ArcRotateCamera(
    'camera',
    Math.PI / 2,
    Math.PI / 2,
    1.5,
    Vector3.Zero(),
    scene,
    true,
  );

  return scene;
}

Babylon.jsもThree.jsも似たような記述が多いからわかりやすいですね。 WebGPUEngineは、Babylon.jsがWebGPUをラップしたエンジンです。initAsync()で非同期に初期化します。

エンジンとシーンが準備できたら、次はコンピュートシェーダを作成します。

コンピュートシェーダの作成

Babylon.jsでは、ComputeShaderクラスを使って、先ほど書いたWGSLシェーダーコードを実行可能なコンピュートシェーダに変換します：

import { ComputeShader } from '@babylonjs/core/Compute/computeShader';
import { WebGPUEngine } from '@babylonjs/core/Engines/webgpuEngine';
import { pixelSortComputeShader } from '../shaders/pixelSort.compute';

export function createPixelSortComputeShader(engine: WebGPUEngine): ComputeShader {
  return new ComputeShader(
    'pixelSort',
    engine,
    { computeSource: pixelSortComputeShader },
    {
      bindingsMapping: {
        inputTexture: { group: 0, binding: 0 },
        outputTexture: { group: 0, binding: 1 },
        params: { group: 0, binding: 2 },
        segments: { group: 0, binding: 3 },
      },
    },
  );
}

bindingsMappingで、WGSLシェーダー内の変数名とバインディング位置を対応付けます。これにより、シェーダー内の@binding(0)などが、どの変数に対応するかが決まります。

コンピュートシェーダが作成できたら、次は処理対象となるテクスチャを準備します。

テクスチャの作成

コンピュートシェーダで処理する動画テクスチャと、結果を書き込むストレージテクスチャを作成します：

export async function createVideoTexture(scene: Scene): Promise<VideoTexture> {
  const video = document.createElement('video');
  video.src = VIDEO_PATH;
  video.loop = true;
  video.muted = true;
  video.autoplay = true;
  
  await video.play();
  await readyPromise(video);

  const texture = new VideoTexture(
    'videoTexture',
    video,
    scene,
    false,
    false,
    Texture.BILINEAR_SAMPLINGMODE,
    {
      autoPlay: true,
      loop: true,
      autoUpdateTexture: true,
    },
  );

  return texture;
}

export function createStorageTexture(engine: WebGPUEngine, width: number, height: number): RawTexture {
  const texture = RawTexture.CreateRGBAStorageTexture(
    null,
    width,
    height,
    engine,
    false,
    false,
    Texture.NEAREST_SAMPLINGMODE,
  );
  return texture;
}

VideoTextureは動画をテクスチャとして扱うためのクラスで、RawTexture.CreateRGBAStorageTextureは書き込み可能なストレージテクスチャを作成します。

テクスチャが準備できたら、コンピュートシェーダにこれらのリソースをバインドする必要があります。

コンピュートシェーダへのリソースのバインド

作成したテクスチャやバッファを、コンピュートシェーダにバインドします。これにより、シェーダー内でこれらのリソースにアクセスできるようになります：

// テクスチャのバインド
pixelSortResources.sortCompute.setTexture('inputTexture', videoTexture, false);
pixelSortResources.sortCompute.setStorageTexture('outputTexture', outputTexture);

// ユニフォームバッファのバインド
pixelSortResources.sortCompute.setUniformBuffer('params', uniforms);

// ストレージバッファのバインド
pixelSortResources.sortCompute.setStorageBuffer('segments', segmentResources.segmentsBuffer);

setTextureは読み取り専用テクスチャ、setStorageTextureは書き込み可能テクスチャ、setUniformBufferはユニフォームバッファ、setStorageBufferはストレージバッファをバインドします。

リソースのバインドが完了したら、いよいよコンピュートシェーダを実行できます。

コンピュートシェーダの実行

dispatchメソッドでコンピュートシェーダを実行します。実行するワークグループの数を指定することで、どの範囲のデータを処理するかを決定します：

// ワークグループ数を計算
const copyGroupsX = Math.ceil(size.width / COPY_WORKGROUP_SIZE_X);
const copyGroupsY = Math.ceil(size.height / COPY_WORKGROUP_SIZE_Y);

// コンピュートシェーダの実行
pixelSortResources.copyCompute.dispatch(copyGroupsX, copyGroupsY, 1);
pixelSortResources.sortCompute.dispatch(effectiveSegmentCount, 1, 1);

dispatchの引数は、実行するワークグループの数です。例えば、dispatch(10, 5, 1)とすると、X方向に10個、Y方向に5個、Z方向に1個のワークグループが実行されます。

コンピュートシェーダを実行する際、シェーダーにパラメータを渡すためにユニフォームバッファを使用します。

ユニフォームバッファの更新

ユニフォームバッファは、シェーダーにパラメータを渡すために使用します。閾値やソート方向などの設定値を、このバッファを通じてシェーダーに渡します：

export function createPixelSortUniformBuffer(engine: WebGPUEngine): UniformBuffer {
  const buffer = new UniformBuffer(engine);
  buffer.addUniform('width', 1);
  buffer.addUniform('height', 1);
  buffer.addUniform('maxSegmentLength', 1);
  buffer.addUniform('sortDirection', 1);
  buffer.addUniform('threshold', 1);
  buffer.addUniform('segmentCount', 1);
  buffer.addUniform('sortOrientation', 1);
  buffer.addUniform('padding0', 1);
  buffer.create();
  return buffer;
}

export function updatePixelSortUniformBuffer(
  uniforms: UniformBuffer,
  width: number,
  height: number,
  threshold: number,
  direction: number,
  orientation: number,
  segmentCount: number,
  maxSegmentLength: number,
) {
  uniforms.updateUInt('width', width);
  uniforms.updateUInt('height', height);
  uniforms.updateFloat('threshold', threshold);
  // ... 他のパラメータも更新
  uniforms.update();
}

UniformBufferは、シェーダーにパラメータを渡すためのバッファです。addUniformでフィールドを追加し、updateメソッドで値を更新します。

これで、Babylon.jsを使ったコンピュートシェーダの実装の準備が整いました。先ほど説明した3つのWGSLシェーダー（フレームコピー、セグメント生成、ピクセルソート）を、これらのBabylon.jsを使って実行できるようになります。

次に、これらがメインループでどのように呼び出され、連携しているかを見てみましょう。

メインループの処理

ここまでで、WGSLシェーダーの実装とBabylon.jsでの実行準備が整いました。最後に、これらを組み合わせて、毎フレームごとにどのように処理を実行するかを見てみましょう。

メインループでは、毎フレームごとに以下の処理を順番に実行します：

scene.onBeforeRenderObservable.add(() => {
  videoTexture.update();
  
  // 1. セグメントリソースの初期化・更新
  segmentResources = ensureSegmentResources(engine, segmentResources, size.width, size.height);
  
  // 2. 横幅いっぱいのセグメントを生成
  const fullWidthCount = populateFullWidthSegments(
    segmentResources,
    size.width,
    size.height,
    orientation,
  );
  
  // 3. 動画フレームをコピー
  pixelSortResources.copyCompute.dispatch(copyGroupsX, copyGroupsY, 1);
  
  // 4. ユニフォームバッファの更新
  updatePixelSortUniformBuffer(
    pixelSortResources.uniforms,
    size.width,
    size.height,
    threshold,
    direction,
    orientation,
    effectiveSegmentCount,
    dynamicMaxSegmentLength,
  );
  
  // 5. ピクセルソートの実行
  pixelSortResources.sortCompute.setStorageBuffer('segments', segmentResources.segmentsBuffer);
  pixelSortResources.sortCompute.dispatch(effectiveSegmentCount, 1, 1);
  
  // 6. 結果を表示
  material.diffuseTexture = pixelSortResources.outputTexture;
  material.emissiveTexture = pixelSortResources.outputTexture;
});

大変だったところ

• wgsl、自分でメモリ管理する必要があって面倒(メモリ管理ミスって途中で何度も画面が真っ暗になったり、てんかん発作起こしそうなぐらい点滅したりした)
• ピクセルソートのロジック結構複雑めだった＆そこにwgslの癖が入り込んで理解にちょっと手間取った

まとめ

Babylon.jsとWebGPUのコンピュートシェーダを使って、リアルタイムでピクセルソートエフェクトを実装しました。コンピュートシェーダを使うことで、フラグメントシェーダでは難しい並列ソート処理を効率的に実行できました。

特に、Bitonic Sortアルゴリズムとコンピュートシェーダの組み合わせにより、GPU上で高速なソート処理を実現できたのが良かったです。

WebGPUはまだ新しい技術ですが、Babylon.jsを使うことで比較的簡単にコンピュートシェーダを扱えることが分かりました。案件とかで使う機会があればよいですね。

（そのためにも早くいろんなブラウザでちゃんと動くようになってほしいものです...）

参考資料

• Babylon.js Documentation
• WebGPU Specification
• WGSL Specification
• Bitonic Sort - Wikipedia
• 面白法人カヤック - 【WebGPU】Compute Shader で Curl Noise を計算してパーティクルを動かす

元記事: Qiita