Compute shaders support Uniform Buffer (WebGPU) (#6186)

* Compute shaders support Uniform Buffer (WebGPU)

* example handles no compute support

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

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

---------

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

examples/assets/models/icosahedron.txt

@@ -0,0 +1,8 @@
+Model Information:
+* title:	UXR Icosahedron
+* source:	https://sketchfab.com/3d-models/uxr-icosahedron-66c69bd0538a455197aebe81ae3a4961
+* author:	enealefons 
+
+Model License:
+* license type:	CC-BY-4.0 (http://creativecommons.org/licenses/by/4.0/)
+* requirements:	Author must be credited. Commercial use is allowed.

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

@@ -6,12 +6,18 @@ export default {
     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 outputTexture: texture_storage_2d<rgba8unorm, write>;
 
-            // color used to tint the source texture
-            const tintColor: vec4<f32> = vec4<f32>(1.0, 0.7, 0.7, 1.0);
+            struct ub_compute {
+                tint : vec4<f32>,
+                offset: f32,
+                frequency: f32
+            }
+
+            @group(0) @binding(0) var<uniform> ubCompute : ub_compute;
+
+            @group(0) @binding(1) var inputTexture: texture_2d<f32>;
+            // @group(0) @binding(2) is a sampler of the inputTexture, but we don't need it in the shader.
+            @group(0) @binding(3) var outputTexture: texture_storage_2d<rgba8unorm, write>;
 
             @compute @workgroup_size(1, 1, 1)
             fn main(@builtin(global_invocation_id) global_id : vec3u) {
@@ -22,8 +28,14 @@ export default {
                 var texColor = textureLoad(inputTexture, uv, 0);
 
                 // tint it
+                var tintColor: vec4<f32> = ubCompute.tint;
                 texColor *= tintColor;
 
+                // scroll a darkness effect over the texture
+                let uvFloat = vec2<f32>(f32(uv.x), f32(uv.y));
+                var darkness: f32 = sin(ubCompute.offset + ubCompute.frequency * length(uvFloat)) * 0.2 + 0.8;
+                texColor *= darkness;
+
                 // write it to the output texture
                 textureStore(outputTexture, vec2<i32>(global_id.xy), texColor);
             }

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

@@ -9,6 +9,7 @@ if (!(canvas instanceof HTMLCanvasElement)) {
 
 const assets = {
     texture: new pc.Asset('color', 'texture', { url: rootPath + '/static/assets/textures/seaside-rocks01-color.jpg' }),
+    solid: new pc.Asset('solid', 'container', { url: rootPath + '/static/assets/models/icosahedron.glb' }),
     helipad: new pc.Asset(
         'helipad-env-atlas',
         'texture',
@@ -27,8 +28,16 @@ const device = await pc.createGraphicsDevice(canvas, gfxOptions);
 const createOptions = new pc.AppOptions();
 createOptions.graphicsDevice = device;
 
-createOptions.componentSystems = [pc.RenderComponentSystem, pc.CameraComponentSystem, pc.LightComponentSystem];
-createOptions.resourceHandlers = [pc.TextureHandler, pc.ContainerHandler];
+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);
@@ -53,7 +62,7 @@ assetListLoader.load(() => {
 
     // setup skydome
     app.scene.skyboxMip = 1;
-    app.scene.skyboxIntensity = 1.7;
+    app.scene.skyboxIntensity = 3;
     app.scene.envAtlas = assets.helipad.resource;
 
     // create camera entity
@@ -62,80 +71,129 @@ assetListLoader.load(() => {
         clearColor: new pc.Color(0.5, 0.6, 0.9)
     });
     app.root.addChild(camera);
-    camera.setPosition(0, 0, 3);
+    camera.setPosition(0.6, 0, 5);
 
     // create directional light entity
     const light = new pc.Entity('light');
     light.addComponent('light');
     app.root.addChild(light);
     light.setEulerAngles(45, 0, 0);
 
-    // create a storage texture, that the compute shader will write to. Make it the same dimensions
-    // as the loaded input texture
-    const storageTexture = new pc.Texture(app.graphicsDevice, {
-        name: 'outputTexture',
-        width: assets.texture.resource.width,
-        height: assets.texture.resource.height,
-        format: pc.PIXELFORMAT_RGBA8,
-        mipmaps: false,
-        minFilter: pc.FILTER_LINEAR,
-        magFilter: pc.FILTER_LINEAR,
-        addressU: pc.ADDRESS_CLAMP_TO_EDGE,
-        addressV: pc.ADDRESS_CLAMP_TO_EDGE,
-
-        // this is a storage texture, allowing compute shader to write to it
-        storage: true
-    });
-
-    // bind group for the compute shader - this needs to match the bindings in the shader
-    const buffers = [];
-    const textures = [
-        new pc.BindTextureFormat('inTexture', pc.SHADERSTAGE_COMPUTE)
-    ];
-    const storageTextures = [
-        new pc.BindStorageTextureFormat('outTexture', pc.PIXELFORMAT_RGBA8, pc.TEXTUREDIMENSION_2D)
-    ];
+    // a helper script that rotates the entity
+    const Rotator = pc.createScript('rotator');
+    Rotator.prototype.update = function (/** @type {number} */ dt) {
+        this.entity.rotate(10 * dt, 20 * dt, 30 * dt);
+    };
 
     // a compute shader that will tint the input texture and write the result to the storage texture
-    const shader = new pc.Shader(device, {
+    const shader = device.supportsCompute ? new pc.Shader(device, {
         name: 'ComputeShader',
         shaderLanguage: pc.SHADERLANGUAGE_WGSL,
         cshader: files['compute-shader.wgsl'],
 
-        computeBindGroupFormat: new pc.BindGroupFormat(device, buffers, textures, storageTextures, {
+        // format of a uniform buffer used by the compute shader
+        computeUniformBufferFormat: new pc.UniformBufferFormat(device, [
+            new pc.UniformFormat('tint', pc.UNIFORMTYPE_VEC4),
+            new pc.UniformFormat('offset', pc.UNIFORMTYPE_FLOAT),
+            new pc.UniformFormat('frequency', pc.UNIFORMTYPE_FLOAT)
+        ]),
+
+        // format of a bind group, providing resources for the compute shader
+        computeBindGroupFormat: new pc.BindGroupFormat(device, [
+            // a uniform buffer we provided format for
+            new pc.BindBufferFormat(pc.UNIFORM_BUFFER_DEFAULT_SLOT_NAME, pc.SHADERSTAGE_COMPUTE)
+        ], [
+            // input textures
+            new pc.BindTextureFormat('inTexture', pc.SHADERSTAGE_COMPUTE)
+        ], [
+            // output storage textures
+            new pc.BindStorageTextureFormat('outTexture', pc.PIXELFORMAT_RGBA8, pc.TEXTUREDIMENSION_2D)
+        ], {
             compute: true
         })
-    });
+    }) : null;
+
+    // helper function, which creates a cube entity, and an instance of the compute shader that will
+    // update its texture each frame
+    const createCubeInstance = (/** @type {pc.Vec3} */position) => {
+
+        if (!device.supportsCompute)
+            return null;
+
+        // create a storage texture, that the compute shader will write to. Make it the same dimensions
+        // as the loaded input texture
+        const storageTexture = new pc.Texture(app.graphicsDevice, {
+            name: 'outputTexture',
+            width: assets.texture.resource.width,
+            height: assets.texture.resource.height,
+            format: pc.PIXELFORMAT_RGBA8,
+            mipmaps: false,
+            minFilter: pc.FILTER_LINEAR,
+            magFilter: pc.FILTER_LINEAR,
+            addressU: pc.ADDRESS_CLAMP_TO_EDGE,
+            addressV: pc.ADDRESS_CLAMP_TO_EDGE,
+
+            // this is a storage texture, allowing compute shader to write to it
+            storage: true
+        });
+
+        // create an instance of the compute shader, and set the input and output textures
+        const compute = new pc.Compute(device, shader);
+        compute.setParameter('inTexture', assets.texture.resource);
+        compute.setParameter('outTexture', storageTexture);
+
+        // add a box in the scene, using the storage texture as a material
+        const material = new pc.StandardMaterial();
+        material.diffuseMap = storageTexture;
+        material.gloss = 0.6;
+        material.metalness = 0.4;
+        material.useMetalness = true;
+        material.update();
+
+        const solid = assets.solid.resource.instantiateRenderEntity();
+        solid.findByName('Object_3').render.meshInstances[0].material = material;
+
+        // add the script to rotate the object
+        solid.addComponent('script');
+        solid.script.create('rotator');
+
+        // place it in the world
+        solid.setLocalPosition(position);
+        solid.setLocalScale(0.25, 0.25, 0.25);
+        app.root.addChild(solid);
+
+        return compute;
+    };
+
+    // create two instances of cube / compute shader
+    const compute1 = createCubeInstance(new pc.Vec3(0, 1, 0));
+    const compute2 = createCubeInstance(new pc.Vec3(0, -1, 0));
+
+    let time = 0;
+    const srcTexture = assets.texture.resource;
+    app.on('update', function (/** @type {number} */ dt) {
+        time += dt;
 
-    // create an instance of the compute shader, and set the input and output textures
-    const compute = new pc.Compute(device, shader);
-    compute.setParameter('outTexture', storageTexture);
-    compute.setParameter('inTexture', assets.texture.resource);
-
-    // add a box in the scene, using the storage texture as a material
-    const material = new pc.StandardMaterial();
-    material.diffuseMap = storageTexture;
-    material.update();
-
-    // create box entity
-    const box = new pc.Entity('cube');
-    box.addComponent('render', {
-        type: 'box',
-        material: material
-    });
-    app.root.addChild(box);
+        if (device.supportsCompute) {
+            compute1.setParameter('offset', 20 * time);
+            compute1.setParameter('frequency', 0.1);
+            compute1.setParameter('tint', [Math.sin(time) * 0.5 + 0.5, 1, 0, 1]);
 
-    app.on('update', function (/** @type {number} */ dt) {
+            // run the compute shader each frame to update the texture
+            compute1.dispatch(srcTexture.width, srcTexture.height);
 
-        box.rotate(10 * dt, 20 * dt, 30 * dt);
+            compute2.setParameter('offset', 10 * time);
+            compute2.setParameter('frequency', 0.03);
+            compute2.setParameter('tint', [1, 0, Math.sin(time) * 0.5 + 0.5, 1]);
 
-        // run the compute shader each frame (even though it generates the same output)
-        compute.dispatch(storageTexture.width, storageTexture.height);
+            // run the compute shader each frame to update the texture
+            compute2.dispatch(srcTexture.width, srcTexture.height);
 
-        // debug render the generated texture
-        app.drawTexture(0.6, -0.7, 0.6, 0.3, storageTexture);
+            // debug render the generated textures
+            app.drawTexture(0.6, 0.5, 0.6, 0.3, compute1.getParameter('outTexture'));
+            app.drawTexture(0.6, -0.5, 0.6, 0.3, compute2.getParameter('outTexture'));
+        }
     });
-
 });
 
 export { app };

src/platform/graphics/compute.js

@@ -61,6 +61,17 @@ class Compute {
         param.value = value;
     }
 
+    /**
+     * Returns the value of a shader parameter from the compute instance.
+     *
+     * @param {string} name - The name of the parameter to get.
+     * @returns {number|number[]|Float32Array|import('./texture.js').Texture|undefined} The value of the
+     * specified parameter.
+     */
+    getParameter(name) {
+        return this.parameters.get(name)?.value;
+    }
+
     /**
      * Deletes a shader parameter from the compute instance.
      *

src/platform/graphics/webgpu/webgpu-compute.js

@@ -1,5 +1,6 @@
 import { Debug, DebugHelper } from "../../../core/debug.js";
 import { BindGroup } from "../bind-group.js";
+import { UniformBuffer } from "../uniform-buffer.js";
 
 /**
  * A WebGPU implementation of the Compute.
@@ -13,9 +14,15 @@ class WebgpuCompute {
         const { device, shader } = compute;
 
         // create bind group
-        const { computeBindGroupFormat } = shader.impl;
+        const { computeBindGroupFormat, computeUniformBufferFormat } = shader.impl;
         Debug.assert(computeBindGroupFormat, 'Compute shader does not have computeBindGroupFormat specified', shader);
-        this.bindGroup = new BindGroup(device, computeBindGroupFormat);
+
+        if (computeUniformBufferFormat) {
+            // TODO: investigate implications of using a non-persistent uniform buffer
+            this.uniformBuffer = new UniformBuffer(device, computeUniformBufferFormat, true);
+        }
+
+        this.bindGroup = new BindGroup(device, computeBindGroupFormat, this.uniformBuffer);
         DebugHelper.setName(this.bindGroup, `Compute-BindGroup_${this.bindGroup.id}`);
 
         // pipeline
@@ -31,6 +38,7 @@ class WebgpuCompute {
 
         // bind group data
         const { bindGroup } = this;
+        bindGroup.defaultUniformBuffer?.update();
         bindGroup.update();
         device.setBindGroup(0, bindGroup);
 

src/platform/graphics/webgpu/webgpu-shader.js

@@ -62,9 +62,13 @@ class WebgpuShader {
             this._vertexCode = definition.vshader ?? null;
             this._fragmentCode = definition.fshader ?? null;
             this._computeCode = definition.cshader ?? null;
+
             this.meshUniformBufferFormat = definition.meshUniformBufferFormat;
             this.meshBindGroupFormat = definition.meshBindGroupFormat;
+
+            this.computeUniformBufferFormat = definition.computeUniformBufferFormat;
             this.computeBindGroupFormat = definition.computeBindGroupFormat;
+
             this.vertexEntryPoint = 'vertexMain';
             this.fragmentEntryPoint = 'fragmentMain';
             shader.ready = true;

src/platform/graphics/webgpu/webgpu-uniform-buffer.js

@@ -21,7 +21,7 @@ class WebgpuUniformBuffer extends WebgpuBuffer {
     unlock(uniformBuffer) {
 
         const device = uniformBuffer.device;
-        super.unlock(device, undefined, GPUBufferUsage.UNIFORM, uniformBuffer.storage);
+        super.unlock(device, undefined, GPUBufferUsage.UNIFORM, uniformBuffer.storageInt32.buffer);
     }
 }