Initial refactor to how bind groups can be used with non-persistent U…

…Bs (#6341) * Initial refactor to how bind groups can be used with non-persistent UBs * types fix * reverted a small thing --------- Co-authored-by: Martin Valigursky <[email protected]>
playcanvas · May 7, 2024 · 95ddb6c · 95ddb6c
1 parent 9b9285f
commit 95ddb6c
Show file tree

Hide file tree

Showing 8 changed files with 127 additions and 66 deletions.
diff --git a/src/platform/graphics/bind-group.js b/src/platform/graphics/bind-group.js
@@ -5,6 +5,18 @@ import { DebugGraphics } from './debug-graphics.js';
 
 let id = 0;
 
+/**
+ * Data structure to hold a bind group and its offsets. This is used by {@link UniformBuffer#update}
+ * to return a dynamic bind group and offset for the uniform buffer.
+ *
+ * @ignore
+ */
+class DynamicBindGroup {
+    bindGroup;
+
+    offsets = [];
+}
+
 /**
  * A bind group represents a collection of {@link UniformBuffer}, {@link Texture} and
  * {@link StorageBuffer} instanced, which can be bind on a GPU for rendering.
@@ -201,4 +213,4 @@ class BindGroup {
     }
 }
 
-export { BindGroup };
+export { BindGroup, DynamicBindGroup };
diff --git a/src/platform/graphics/dynamic-buffer.js b/src/platform/graphics/dynamic-buffer.js
@@ -0,0 +1,50 @@
+import { DebugHelper } from '../../core/debug.js';
+import { BindGroupFormat, BindUniformBufferFormat } from './bind-group-format.js';
+import { BindGroup } from './bind-group.js';
+import { SHADERSTAGE_FRAGMENT, SHADERSTAGE_VERTEX, UNIFORM_BUFFER_DEFAULT_SLOT_NAME } from './constants.js';
+
+/**
+ * A base class representing a single per platform buffer.
+ *
+ * @ignore
+ */
+class DynamicBuffer {
+    /** @type {import('./graphics-device.js').GraphicsDevice} */
+    device;
+
+    /**
+     * A cache of bind groups for each uniform buffer size, which is used to avoid creating a new
+     * bind group for each uniform buffer.
+     *
+     * @type {Map<number, BindGroup>}
+     */
+    bindGroupCache = new Map();
+
+    constructor(device) {
+        this.device = device;
+
+        // format of the bind group
+        this.bindGroupFormat = new BindGroupFormat(this.device, [
+            new BindUniformBufferFormat(UNIFORM_BUFFER_DEFAULT_SLOT_NAME, SHADERSTAGE_VERTEX | SHADERSTAGE_FRAGMENT)
+        ]);
+    }
+
+    getBindGroup(ub) {
+        const ubSize = ub.format.byteSize;
+        let bindGroup = this.bindGroupCache.get(ubSize);
+        if (!bindGroup) {
+
+            // bind group
+            // we pass ub to it, but internally only its size is used
+            bindGroup = new BindGroup(this.device, this.bindGroupFormat, ub);
+            DebugHelper.setName(bindGroup, `DynamicBuffer-BindGroup_${bindGroup.id}-${ubSize}`);
+            bindGroup.update();
+
+            this.bindGroupCache.set(ubSize, bindGroup);
+        }
+
+        return bindGroup;
+    }
+}
+
+export { DynamicBuffer };
diff --git a/src/platform/graphics/dynamic-buffers.js b/src/platform/graphics/dynamic-buffers.js
@@ -1,30 +1,16 @@
 import { Debug } from '../../core/debug.js';
 import { math } from '../../core/math/math.js';
 
-/**
- * A base class representing a single per platform buffer.
- *
- * @ignore
- */
-class DynamicBuffer {
-    /** @type {import('./graphics-device.js').GraphicsDevice} */
-    device;
-
-    constructor(device) {
-        this.device = device;
-    }
-}
-
 /**
  * A container for storing the used areas of a pair of staging and gpu buffers.
  *
  * @ignore
  */
 class UsedBuffer {
-    /** @type {DynamicBuffer} */
+    /** @type {import('./dynamic-buffer.js').DynamicBuffer} */
     gpuBuffer;
 
-    /** @type {DynamicBuffer} */
+    /** @type {import('./dynamic-buffer.js').DynamicBuffer} */
     stagingBuffer;
 
     /**
@@ -59,7 +45,7 @@ class DynamicBufferAllocation {
     /**
      * The gpu buffer this allocation will be copied to.
      *
-     * @type {DynamicBuffer}
+     * @type {import('./dynamic-buffer.js').DynamicBuffer}
      */
     gpuBuffer;
 
@@ -93,14 +79,14 @@ class DynamicBuffers {
     /**
      * Internally allocated gpu buffers.
      *
-     * @type {DynamicBuffer[]}
+     * @type {import('./dynamic-buffer.js').DynamicBuffer[]}
      */
     gpuBuffers = [];
 
     /**
      * Internally allocated staging buffers (CPU writable)
      *
-     * @type {DynamicBuffer[]}
+     * @type {import('./dynamic-buffer.js').DynamicBuffer[]}
      */
     stagingBuffers = [];
 
@@ -215,4 +201,4 @@ class DynamicBuffers {
     }
 }
 
-export { DynamicBuffer, DynamicBuffers, DynamicBufferAllocation };
+export { DynamicBuffers, DynamicBufferAllocation };
diff --git a/src/platform/graphics/uniform-buffer.js b/src/platform/graphics/uniform-buffer.js
@@ -304,11 +304,11 @@ class UniformBuffer {
      *
      * @param {import('./uniform-buffer-format.js').UniformFormat} uniformFormat - The format of
      * the uniform.
+     * @param {any} value - The value to assign to the uniform.
      */
-    setUniform(uniformFormat) {
+    setUniform(uniformFormat, value) {
         Debug.assert(uniformFormat);
         const offset = uniformFormat.offset;
-        const value = uniformFormat.scopeId.value;
 
         if (value !== null && value !== undefined) {
 
@@ -328,46 +328,70 @@ class UniformBuffer {
      * Assign a value to the uniform specified by name.
      *
      * @param {string} name - The name of the uniform.
+     * @param {any} value - The value to assign to the uniform.
      */
-    set(name) {
+    set(name, value) {
         const uniformFormat = this.format.map.get(name);
         Debug.assert(uniformFormat, `Uniform name [${name}] is not part of the Uniform buffer.`);
         if (uniformFormat) {
-            this.setUniform(uniformFormat);
+            this.setUniform(uniformFormat, value);
         }
     }
 
-    update() {
+    startUpdate(dynamicBindGroup) {
 
-        const persistent = this.persistent;
-        if (!persistent) {
+        if (!this.persistent) {
 
             // allocate memory from dynamic buffer for this frame
             const allocation = this.allocation;
             const oldGpuBuffer = allocation.gpuBuffer;
             this.device.dynamicBuffers.alloc(allocation, this.format.byteSize);
             this.assignStorage(allocation.storage);
 
+            // get info about bind group we can use for this non-persistent UB for this frame
+            if (dynamicBindGroup) {
+                dynamicBindGroup.bindGroup = allocation.gpuBuffer.getBindGroup(this);
+                dynamicBindGroup.offsets[0] = allocation.offset;
+            }
+
             // buffer has changed, update the render version to force bind group to be updated
             if (oldGpuBuffer !== allocation.gpuBuffer) {
                 this.renderVersionDirty = this.device.renderVersion;
             }
         }
+    }
 
-        // set new values
-        const uniforms = this.format.uniforms;
-        for (let i = 0; i < uniforms.length; i++) {
-            this.setUniform(uniforms[i]);
-        }
+    endUpdate() {
 
-        if (persistent) {
+        if (this.persistent) {
             // Upload the new data
             this.impl.unlock(this);
         } else {
             this.storageFloat32 = null;
             this.storageInt32 = null;
         }
     }
+
+    /**
+     * @param {import('./bind-group.js').DynamicBindGroup} [dynamicBindGroup] - The function fills
+     * in the info about the dynamic bind group for this frame, which uses this uniform buffer. Only
+     * used if the uniform buffer is non-persistent. This allows the uniform buffer to be used
+     * without having to create a bind group for it. Note that the bind group can only contains
+     * this single uniform buffer, and no other resources.
+     */
+    update(dynamicBindGroup) {
+
+        this.startUpdate(dynamicBindGroup);
+
+        // set new values
+        const uniforms = this.format.uniforms;
+        for (let i = 0; i < uniforms.length; i++) {
+            const value = uniforms[i].scopeId.value;
+            this.setUniform(uniforms[i], value);
+        }
+
+        this.endUpdate();
+    }
 }
 
 export { UniformBuffer };
diff --git a/src/platform/graphics/webgpu/webgpu-bind-group.js b/src/platform/graphics/webgpu/webgpu-bind-group.js
@@ -34,7 +34,6 @@ class WebgpuBindGroup {
     }
 
     destroy() {
-        // this.bindGroup?.destroy();
         this.bindGroup = null;
     }
 
@@ -44,8 +43,8 @@ class WebgpuBindGroup {
      * @param {import('./webgpu-graphics-device.js').WebgpuGraphicsDevice} device - Graphics device.
      * @param {import('../bind-group.js').BindGroup} bindGroup - Bind group to create the
      * descriptor for.
-     * @returns {object} - Returns the generated descriptor of type
-     * GPUBindGroupDescriptor, which can be used to create a GPUBindGroup
+     * @returns {object} - Returns the generated descriptor of type GPUBindGroupDescriptor, which
+     * can be used to create a GPUBindGroup
      */
     createDescriptor(device, bindGroup) {
 

diff --git a/src/platform/graphics/webgpu/webgpu-clear-renderer.js b/src/platform/graphics/webgpu/webgpu-clear-renderer.js
@@ -1,14 +1,13 @@
-import { Debug, DebugHelper } from "../../../core/debug.js";
-import { BindUniformBufferFormat, BindGroupFormat } from "../bind-group-format.js";
+import { Debug } from "../../../core/debug.js";
 import { UniformBufferFormat, UniformFormat } from "../uniform-buffer-format.js";
 import { BlendState } from "../blend-state.js";
 import {
     CULLFACE_NONE,
-    PRIMITIVE_TRISTRIP, SHADERLANGUAGE_WGSL, SHADERSTAGE_FRAGMENT, SHADERSTAGE_VERTEX,
-    UNIFORMTYPE_FLOAT, UNIFORMTYPE_VEC4, UNIFORM_BUFFER_DEFAULT_SLOT_NAME, BINDGROUP_MESH, CLEARFLAG_COLOR, CLEARFLAG_DEPTH, CLEARFLAG_STENCIL
+    PRIMITIVE_TRISTRIP, SHADERLANGUAGE_WGSL,
+    UNIFORMTYPE_FLOAT, UNIFORMTYPE_VEC4, BINDGROUP_MESH, CLEARFLAG_COLOR, CLEARFLAG_DEPTH, CLEARFLAG_STENCIL
 } from "../constants.js";
 import { Shader } from "../shader.js";
-import { BindGroup } from "../bind-group.js";
+import { DynamicBindGroup } from "../bind-group.js";
 import { UniformBuffer } from "../uniform-buffer.js";
 import { DebugGraphics } from "../debug-graphics.js";
 import { DepthState } from "../depth-state.js";
@@ -77,19 +76,10 @@ class WebgpuClearRenderer {
             new UniformFormat('depth', UNIFORMTYPE_FLOAT)
         ]), false);
 
-        // format of the bind group
-        const bindGroupFormat = new BindGroupFormat(device, [
-            new BindUniformBufferFormat(UNIFORM_BUFFER_DEFAULT_SLOT_NAME, SHADERSTAGE_VERTEX | SHADERSTAGE_FRAGMENT)
-        ]);
-
-        // bind group
-        this.bindGroup = new BindGroup(device, bindGroupFormat, this.uniformBuffer);
-        DebugHelper.setName(this.bindGroup, `ClearRenderer-BindGroup_${this.bindGroup.id}`);
+        this.dynamicBindGroup = new DynamicBindGroup();
 
         // uniform data
         this.colorData = new Float32Array(4);
-        this.colorId = device.scope.resolve('color');
-        this.depthId = device.scope.resolve('depth');
     }
 
     destroy() {
@@ -98,9 +88,6 @@ class WebgpuClearRenderer {
 
         this.uniformBuffer.destroy();
         this.uniformBuffer = null;
-
-        this.bindGroup.destroy();
-        this.bindGroup = null;
     }
 
     clear(device, renderTarget, options, defaultOptions) {
@@ -111,6 +98,11 @@ class WebgpuClearRenderer {
 
             DebugGraphics.pushGpuMarker(device, 'CLEAR-RENDERER');
 
+            // dynamic bind group for this UB
+            const { uniformBuffer, dynamicBindGroup } = this;
+            uniformBuffer.startUpdate(dynamicBindGroup);
+            device.setBindGroup(BINDGROUP_MESH, dynamicBindGroup.bindGroup, dynamicBindGroup.offsets);
+
             // setup clear color
             if ((flags & CLEARFLAG_COLOR) && (renderTarget.colorBuffer || renderTarget.impl.assignedColorTexture)) {
                 const color = options.color ?? defaultOptions.color;
@@ -120,16 +112,16 @@ class WebgpuClearRenderer {
             } else {
                 device.setBlendState(BlendState.NOWRITE);
             }
-            this.colorId.setValue(this.colorData);
+            uniformBuffer.set('color', this.colorData);
 
             // setup depth clear
             if ((flags & CLEARFLAG_DEPTH) && renderTarget.depth) {
                 const depth = options.depth ?? defaultOptions.depth;
-                this.depthId.setValue(depth);
+                uniformBuffer.set('depth', depth);
                 device.setDepthState(DepthState.WRITEDEPTH);
 
             } else {
-                this.depthId.setValue(1);
+                uniformBuffer.set('depth', 1);
                 device.setDepthState(DepthState.NODEPTH);
             }
 
@@ -138,16 +130,13 @@ class WebgpuClearRenderer {
                 Debug.warnOnce("ClearRenderer does not support stencil clear at the moment");
             }
 
+            uniformBuffer.endUpdate();
+
             device.setCullMode(CULLFACE_NONE);
 
             // render 4 vertices without vertex buffer
             device.setShader(this.shader);
 
-            const bindGroup = this.bindGroup;
-            bindGroup.defaultUniformBuffer.update();
-            bindGroup.update();
-            device.setBindGroup(BINDGROUP_MESH, bindGroup);
-
             device.draw(primitive);
 
             DebugGraphics.popGpuMarker(device);

diff --git a/src/platform/graphics/webgpu/webgpu-dynamic-buffer.js b/src/platform/graphics/webgpu/webgpu-dynamic-buffer.js
@@ -1,5 +1,5 @@
 import { DebugHelper } from "../../../core/debug.js";
-import { DynamicBuffer } from "../dynamic-buffers.js";
+import { DynamicBuffer } from "../dynamic-buffer.js";
 
 /**
  * @ignore

diff --git a/src/platform/graphics/webgpu/webgpu-graphics-device.js b/src/platform/graphics/webgpu/webgpu-graphics-device.js
@@ -419,14 +419,15 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
     /**
      * @param {number} index - Index of the bind group slot
      * @param {import('../bind-group.js').BindGroup} bindGroup - Bind group to attach
+     * @param {number[]} [offsets] - Byte offsets for all uniform buffers in the bind group.
      */
-    setBindGroup(index, bindGroup) {
+    setBindGroup(index, bindGroup, offsets) {
 
         // TODO: this condition should be removed, it's here to handle fake grab pass, which should be refactored instead
         if (this.passEncoder) {
 
             // set it on the device
-            this.passEncoder.setBindGroup(index, bindGroup.impl.bindGroup, bindGroup.uniformBufferOffsets);
+            this.passEncoder.setBindGroup(index, bindGroup.impl.bindGroup, offsets ?? bindGroup.uniformBufferOffsets);
 
             // store the active formats, used by the pipeline creation
             this.bindGroupFormats[index] = bindGroup.format.impl;