Implementation of StorageBuffer on WebGPU (#6201)

* Implementation of StorageBuffer on WebGPU

* comment

* allow storage buffer data to be returned to a preallocated typed buffer

* comment

* Update src/platform/graphics/storage-buffer.js

Co-authored-by: Will Eastcott <will@playcanvas.com>

* Update src/platform/graphics/webgpu/webgpu-buffer.js

Co-authored-by: Will Eastcott <will@playcanvas.com>

* Update src/platform/graphics/webgpu/webgpu-graphics-device.js

Co-authored-by: Will Eastcott <will@playcanvas.com>

* changing name to ‘Unnamed’

---------

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
Co-authored-by: Will Eastcott <will@playcanvas.com>

examples/src/examples/compute/histogram/config.mjs

@@ -0,0 +1,30 @@
+/**
+ * @type {import('../../../../types.mjs').ExampleConfig}
+ */
+export default {
+    HIDDEN: true,
+    WEBGPU_REQUIRED: true,
+    FILES: {
+        'compute-shader.wgsl': /* wgsl */ `
+
+            @group(0) @binding(0) var inputTexture: texture_2d<f32>;
+            // @group(0) @binding(1) is a sampler of the inputTexture, but we don't need it in the shader.
+            @group(0) @binding(2) var<storage, read_write> bins: array<atomic<u32>>;
+    
+            fn luminance(color: vec3f) -> f32 {
+                return saturate(dot(color, vec3f(0.2126, 0.7152, 0.0722)));
+            }
+            
+            @compute @workgroup_size(1, 1, 1)
+            fn main(@builtin(global_invocation_id) global_invocation_id: vec3u) {
+                let numBins = f32(arrayLength(&bins));
+                let lastBinIndex = u32(numBins - 1);
+                let position = global_invocation_id.xy;
+                let color = textureLoad(inputTexture, position, 0);
+                let v = luminance(color.rgb);
+                let bin = min(u32(v * numBins), lastBinIndex);
+                atomicAdd(&bins[bin], 1u);
+            }
+        `
+    }
+};

examples/src/examples/compute/histogram/example.mjs

@@ -0,0 +1,178 @@
+import * as pc from 'playcanvas';
+import { deviceType, rootPath } from '@examples/utils';
+import files from '@examples/files';
+
+// Note: the example is based on this article:
+// https://webgpufundamentals.org/webgpu/lessons/webgpu-compute-shaders-histogram.html
+// A simpler but less performant version of the compute shader is used for simplicity.
+
+const canvas = document.getElementById('application-canvas');
+if (!(canvas instanceof HTMLCanvasElement)) {
+    throw new Error('No canvas found');
+}
+
+const assets = {
+    solid: new pc.Asset('solid', 'container', { url: rootPath + '/static/assets/models/icosahedron.glb' }),
+    helipad: new pc.Asset(
+        'helipad-env-atlas',
+        'texture',
+        { url: rootPath + '/static/assets/cubemaps/helipad-env-atlas.png' },
+        { type: pc.TEXTURETYPE_RGBP, mipmaps: false }
+    )
+};
+
+const gfxOptions = {
+    deviceTypes: [deviceType],
+    glslangUrl: rootPath + '/static/lib/glslang/glslang.js',
+    twgslUrl: rootPath + '/static/lib/twgsl/twgsl.js'
+};
+
+const device = await pc.createGraphicsDevice(canvas, gfxOptions);
+const createOptions = new pc.AppOptions();
+createOptions.graphicsDevice = device;
+
+createOptions.componentSystems = [
+    pc.RenderComponentSystem,
+    pc.CameraComponentSystem,
+    pc.LightComponentSystem,
+    pc.ScriptComponentSystem
+];
+createOptions.resourceHandlers = [
+    pc.TextureHandler,
+    pc.ContainerHandler
+];
+
+const app = new pc.AppBase(canvas);
+app.init(createOptions);
+app.start();
+
+// Set the canvas to fill the window and automatically change resolution to be the same as the canvas size
+app.setCanvasFillMode(pc.FILLMODE_FILL_WINDOW);
+app.setCanvasResolution(pc.RESOLUTION_AUTO);
+
+// Ensure canvas is resized when window changes size
+const resize = () => app.resizeCanvas();
+window.addEventListener('resize', resize);
+app.on('destroy', () => {
+    window.removeEventListener('resize', resize);
+});
+
+const assetListLoader = new pc.AssetListLoader(Object.values(assets), app.assets);
+assetListLoader.load(() => {
+
+    // set up some general scene rendering properties
+    app.scene.toneMapping = pc.TONEMAP_ACES;
+
+    // setup skydome
+    app.scene.skyboxMip = 2;
+    app.scene.skyboxIntensity = 0.3;
+    app.scene.envAtlas = assets.helipad.resource;
+
+    // create camera entity
+    const camera = new pc.Entity('camera');
+    camera.addComponent('camera');
+    app.root.addChild(camera);
+    camera.setPosition(0, 0, 5);
+
+    // Enable the camera to render the scene's color map, available as uSceneColorMap in the shaders.
+    // This allows us to use the rendered scene as an input for the histogram compute shader.
+    camera.camera.requestSceneColorMap(true);
+
+    // create directional light entity
+    const light = new pc.Entity('light');
+    light.addComponent('light', {
+        type: 'directional',
+        color: new pc.Color(1, 1, 1),
+        intensity: 15
+    });
+    app.root.addChild(light);
+    light.setEulerAngles(45, 0, 40);
+
+    // a helper script that rotates the entity
+    const Rotator = pc.createScript('rotator');
+    Rotator.prototype.update = function (/** @type {number} */ dt) {
+        this.entity.rotate(5 * dt, 10 * dt, -15 * dt);
+    };
+
+    // a compute shader that will compute the histogram of the input texture and write the result to the storage buffer
+    const shader = device.supportsCompute ? new pc.Shader(device, {
+        name: 'ComputeShader',
+        shaderLanguage: pc.SHADERLANGUAGE_WGSL,
+        cshader: files['compute-shader.wgsl'],
+
+        // format of a bind group, providing resources for the compute shader
+        computeBindGroupFormat: new pc.BindGroupFormat(device, [
+            // no uniform buffer
+        ], [
+            // input texture - the scene color map
+            new pc.BindTextureFormat('uSceneColorMap', pc.SHADERSTAGE_COMPUTE)
+        ], [
+            // no storage textures
+        ], [
+            // output storage buffer
+            new pc.BindStorageBufferFormat('outBuffer', pc.SHADERSTAGE_COMPUTE)
+        ])
+    }) : null;
+
+    // Create a storage buffer to which the compute shader will write the histogram values.
+    const numBins = 256;
+    const histogramStorageBuffer = new pc.StorageBuffer(
+        device, numBins * 4,     // 4 bytes per value, storing unsigned int
+        pc.BUFFERUSAGE_COPY_SRC | // needed for reading back the data to CPU
+        pc.BUFFERUSAGE_COPY_DST   // needed for clearing the buffer
+    );
+
+    // Create an instance of the compute shader, and set the input and output data. Note that we do
+    // not provide a value for `uSceneColorMap` as this is done by the engine internally.
+    const compute = new pc.Compute(device, shader, 'ComputeHistogram');
+    compute.setParameter('outBuffer', histogramStorageBuffer);
+
+    // instantiate the spinning mesh
+    const solid = assets.solid.resource.instantiateRenderEntity();
+    solid.addComponent('script');
+    solid.script.create('rotator');
+    solid.setLocalPosition(0, 0.4, 0);
+    solid.setLocalScale(0.35, 0.35, 0.35);
+    app.root.addChild(solid);
+
+    let firstFrame = true;
+    app.on('update', function (/** @type {number} */ dt) {
+
+        // The update function runs every frame before the frame gets rendered. On the first time it
+        // runs, the scene color map has not been rendered yet, so we skip the first frame.
+        if (firstFrame) {
+            firstFrame = false;
+            return;
+        }
+
+        if (device.supportsCompute) {
+
+            // clear the storage buffer, to avoid the accumulation buildup
+            histogramStorageBuffer.clear();
+
+            // dispatch the compute shader
+            compute.setupDispatch(app.graphicsDevice.width, app.graphicsDevice.height);
+            device.computeDispatch([compute]);
+
+            // Read back the histogram data from the storage buffer. None that the returned promise
+            // will be resolved later, when the GPU is done running it, and so the histogram on the
+            // screen will be up to few frames behind.
+            const histogramData = new Uint32Array(numBins);
+            histogramStorageBuffer.read(0, undefined, histogramData).then(
+                (data) => {
+                    // render the histogram using lines
+                    const scale = 1 / 50000;
+                    const positions = [];
+                    for (let x = 0; x < data.length; x++) {
+                        const value = pc.math.clamp(data[x] * scale, 0, 0.2);
+                        positions.push(x * 0.001, -0.35, 4);
+                        positions.push(x * 0.001, value - 0.35, 4);
+                    }
+                    app.drawLineArrays(positions, pc.Color.YELLOW);
+                }
+            );
+        }
+    });
+});
+
+export { app };

examples/src/examples/compute/texture-gen/example.mjs

@@ -136,7 +136,7 @@ assetListLoader.load(() => {
         });
 
         // create an instance of the compute shader, and set the input and output textures
-        const compute = new pc.Compute(device, shader);
+        const compute = new pc.Compute(device, shader, 'ComputeModifyTexture');
         compute.setParameter('inTexture', assets.texture.resource);
         compute.setParameter('outTexture', storageTexture);
 

src/core/constants.js

@@ -95,6 +95,14 @@ export const TRACEID_VRAM_VB = 'VRAM.Vb';
  */
 export const TRACEID_VRAM_IB = 'VRAM.Ib';
 
+/**
+ * Logs the vram use by the storage buffers.
+ *
+ * @type {string}
+ * @category Debug
+ */
+export const TRACEID_VRAM_SB = 'VRAM.Sb';
+
 /**
  * Logs the creation of bind groups.
  *

src/index.js

@@ -72,7 +72,7 @@ export * from './platform/audio/constants.js';
 // PLATFORM / GRAPHICS
 export * from './platform/graphics/constants.js';
 export { createGraphicsDevice } from './platform/graphics/graphics-device-create.js';
-export { BindGroupFormat, BindBufferFormat, BindTextureFormat, BindStorageTextureFormat } from './platform/graphics/bind-group-format.js';
+export { BindGroupFormat, BindBufferFormat, BindTextureFormat, BindStorageTextureFormat, BindStorageBufferFormat } from './platform/graphics/bind-group-format.js';
 export { BlendState } from './platform/graphics/blend-state.js';
 export { Compute } from './platform/graphics/compute.js';
 export { DepthState } from './platform/graphics/depth-state.js';
@@ -85,6 +85,7 @@ export { ScopeSpace } from './platform/graphics/scope-space.js';
 export { Shader } from './platform/graphics/shader.js';
 export { ShaderProcessorOptions } from './platform/graphics/shader-processor-options.js';   // used by splats in extras
 export { ShaderUtils } from './platform/graphics/shader-utils.js';  // used by splats in extras
+export { StorageBuffer } from './platform/graphics/storage-buffer.js';
 export { Texture } from './platform/graphics/texture.js';
 export { TextureUtils } from './platform/graphics/texture-utils.js';
 export { TransformFeedback } from './platform/graphics/transform-feedback.js';

src/platform/graphics/bind-group-format.js

@@ -28,6 +28,22 @@ class BindBufferFormat {
     }
 }
 
+/**
+ * @ignore
+ */
+class BindStorageBufferFormat {
+    /** @type {import('./scope-id.js').ScopeId} */
+    scopeId;
+
+    constructor(name, visibility) {
+        /** @type {string} */
+        this.name = name;
+
+        // SHADERSTAGE_VERTEX, SHADERSTAGE_FRAGMENT, SHADERSTAGE_COMPUTE
+        this.visibility = visibility;
+    }
+}
+
 /**
  * @ignore
  */
@@ -73,8 +89,6 @@ class BindStorageTextureFormat {
  * @ignore
  */
 class BindGroupFormat {
-    compute = false;
-
     /**
      * @param {import('./graphics-device.js').GraphicsDevice} graphicsDevice - The graphics device
      * used to manage this vertex format.
@@ -84,8 +98,10 @@ class BindGroupFormat {
      * Defaults to an empty array.
      * @param {BindStorageTextureFormat[]} [storageTextureFormats] - An array of bind storage texture
      * formats (storage textures), used by the compute shader. Defaults to an empty array.
+     * @param {BindStorageBufferFormat[]} [storageBufferFormats] - An array of bind storage buffer
+     * formats. Defaults to an empty array.
      */
-    constructor(graphicsDevice, bufferFormats = [], textureFormats = [], storageTextureFormats = []) {
+    constructor(graphicsDevice, bufferFormats = [], textureFormats = [], storageTextureFormats = [], storageBufferFormats = []) {
         this.id = id++;
         DebugHelper.setName(this, `BindGroupFormat_${this.id}`);
 
@@ -127,6 +143,19 @@ class BindGroupFormat {
             tf.scopeId = scope.resolve(tf.name);
         });
 
+        /** @type {BindStorageBufferFormat[]} */
+        this.storageBufferFormats = storageBufferFormats;
+
+        // maps a storage buffer format name to a slot index
+        /** @type {Map<string, number>} */
+        this.storageBufferFormatsMap = new Map();
+        storageBufferFormats.forEach((bf, i) => {
+            this.storageBufferFormatsMap.set(bf.name, i);
+
+            // resolve scope id
+            bf.scopeId = scope.resolve(bf.name);
+        });
+
         this.impl = graphicsDevice.createBindGroupFormatImpl(this);
 
         Debug.trace(TRACEID_BINDGROUPFORMAT_ALLOC, `Alloc: Id ${this.id}`, this);
@@ -205,4 +234,4 @@ class BindGroupFormat {
     }
 }
 
-export { BindBufferFormat, BindTextureFormat, BindGroupFormat, BindStorageTextureFormat };
+export { BindBufferFormat, BindTextureFormat, BindGroupFormat, BindStorageTextureFormat, BindStorageBufferFormat };

src/platform/graphics/bind-group.js

@@ -6,8 +6,8 @@ import { DebugGraphics } from './debug-graphics.js';
 let id = 0;
 
 /**
- * A bind group represents an collection of {@link UniformBuffer} and {@link Texture} instance,
- * which can be bind on a GPU for rendering.
+ * A bind group represents a collection of {@link UniformBuffer}, {@link Texture} and
+ * {@link StorageBuffer} instanced, which can be bind on a GPU for rendering.
  *
  * @ignore
  */
@@ -50,6 +50,7 @@ class BindGroup {
 
         this.textures = [];
         this.storageTextures = [];
+        this.storageBuffers = [];
         this.uniformBuffers = [];
 
         /** @type {import('./uniform-buffer.js').UniformBuffer} */
@@ -87,6 +88,22 @@ class BindGroup {
         }
     }
 
+    /**
+     * Assign a storage buffer to a slot.
+     *
+     * @param {string} name - The name of the storage buffer slot.
+     * @param {import('./storage-buffer.js').StorageBuffer} storageBuffer - The storage buffer to
+     * assign to the slot.
+     */
+    setStorageBuffer(name, storageBuffer) {
+        const index = this.format.storageBufferFormatsMap.get(name);
+        Debug.assert(index !== undefined, `Setting a storage buffer [${name}] on a bind group with id: ${this.id} which does not contain in, while rendering [${DebugGraphics.toString()}]`, this);
+        if (this.storageBuffers[index] !== storageBuffer) {
+            this.storageBuffers[index] = storageBuffer;
+            this.dirty = true;
+        }
+    }
+
     /**
      * Assign a texture to a named slot.
      *
@@ -129,7 +146,7 @@ class BindGroup {
     update() {
 
         // TODO: implement faster version of this, which does not call SetTexture, which does a map lookup
-        const { textureFormats, storageTextureFormats } = this.format;
+        const { textureFormats, storageTextureFormats, storageBufferFormats } = this.format;
 
         for (let i = 0; i < textureFormats.length; i++) {
             const textureFormat = textureFormats[i];
@@ -145,6 +162,13 @@ class BindGroup {
             this.setStorageTexture(storageTextureFormat.name, value);
         }
 
+        for (let i = 0; i < storageBufferFormats.length; i++) {
+            const storageBufferFormat = storageBufferFormats[i];
+            const value = storageBufferFormat.scopeId.value;
+            Debug.assert(value, `Value was not set when assigning storage buffer slot [${storageBufferFormat.name}] to a bind group, while rendering [${DebugGraphics.toString()}]`, this);
+            this.setStorageBuffer(storageBufferFormat.name, value);
+        }
+
         // update uniform buffer offsets
         this.uniformBufferOffsets.length = this.uniformBuffers.length;
         for (let i = 0; i < this.uniformBuffers.length; i++) {