Bump allocation for Uniform Buffers on WebGPU (#5438)

* Bump allocation for Uniform Buffers on WebGPU

* lint

* types

---------

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>

package.json

@@ -36,6 +36,7 @@
       "glslang": "readonly",
       "GPUBufferUsage": "readonly",
       "GPUColorWrite": "readonly",
+      "GPUMapMode": "readonly",
       "GPUShaderStage": "readonly",
       "GPUTextureUsage": "readonly",
       "opentype": "readonly",

src/platform/graphics/bind-group.js

@@ -20,6 +20,17 @@ class BindGroup {
      */
     renderVersionUpdated = -1;
 
+    /** @type {import('./uniform-buffer.js').UniformBuffer[]} */
+    uniformBuffers;
+
+    /**
+     * An array of offsets for each uniform buffer in the bind group. This is the offset in the
+     * buffer where the uniform buffer data starts.
+     *
+     * @type {number[]}
+     */
+    uniformBufferOffsets = [];
+
     /**
      * Create a new Bind Group.
      *
@@ -108,6 +119,20 @@ class BindGroup {
             this.setTexture(textureFormat.name, value);
         }
 
+        // update uniform buffer offsets
+        this.uniformBufferOffsets.length = this.uniformBuffers.length;
+        for (let i = 0; i < this.uniformBuffers.length; i++) {
+            const uniformBuffer = this.uniformBuffers[i];
+
+            // offset
+            this.uniformBufferOffsets[i] = uniformBuffer.offset;
+
+            // test if any of the uniform buffers have changed (not their content, but the buffer container itself)
+            if (this.renderVersionUpdated < uniformBuffer.renderVersionDirty) {
+                this.dirty = true;
+            }
+        }
+
         if (this.dirty) {
             this.dirty = false;
             this.renderVersionUpdated = this.device.renderVersion;

src/platform/graphics/dynamic-buffers.js

@@ -0,0 +1,218 @@
+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} */
+    gpuBuffer;
+
+    /** @type {DynamicBuffer} */
+    stagingBuffer;
+
+    /**
+     * The beginning position of the used area that needs to be copied from staging to to the GPU
+     * buffer.
+     *
+     * @type {number}
+     */
+    offset;
+
+    /**
+     * Used byte size of the buffer, from the offset.
+     *
+     * @type {number}
+     */
+    size;
+}
+
+/**
+ * A container for storing the return values of an allocation function.
+ *
+ * @ignore
+ */
+class DynamicBufferAllocation {
+    /**
+     * The storage access to the allocated data in the staging buffer.
+     *
+     * @type {Int32Array}
+     */
+    storage;
+
+    /**
+     * The gpu buffer this allocation will be copied to.
+     *
+     * @type {DynamicBuffer}
+     */
+    gpuBuffer;
+
+    /**
+     * Offset in the gpuBuffer where the data will be copied to.
+     *
+     * @type {number}
+     */
+    offset;
+}
+
+/**
+ * The DynamicBuffers class provides a dynamic memory allocation system for uniform buffer data,
+ * particularly for non-persistent uniform buffers. This class utilizes a bump allocator to
+ * efficiently allocate aligned memory space from a set of large buffers managed internally. To
+ * utilize this system, the user writes data to CPU-accessible staging buffers. When submitting
+ * command buffers that require these buffers, the system automatically uploads the data to the GPU
+ * buffers. This approach ensures efficient memory management and smooth data transfer between the
+ * CPU and GPU.
+ *
+ * @ignore
+ */
+class DynamicBuffers {
+    /**
+     * Allocation size of the underlying buffers.
+     *
+     * @type {number}
+     */
+    bufferSize;
+
+    /**
+     * Internally allocated gpu buffers.
+     *
+     * @type {DynamicBuffer[]}
+     */
+    gpuBuffers = [];
+
+    /**
+     * Internally allocated staging buffers (CPU writable)
+     *
+     * @type {DynamicBuffer[]}
+     */
+    stagingBuffers = [];
+
+    /**
+     * @type {UsedBuffer[]}
+     */
+    usedBuffers = [];
+
+    /**
+     * @type {UsedBuffer}
+     */
+    activeBuffer = null;
+
+    /**
+     * Create the system of dynamic buffers.
+     *
+     * @param {import('./graphics-device.js').GraphicsDevice} device - The graphics device.
+     * @param {number} bufferSize - The size of the underlying large buffers.
+     * @param {number} bufferAlignment - Alignment of each allocation.
+     */
+    constructor(device, bufferSize, bufferAlignment) {
+        this.device = device;
+        this.bufferSize = bufferSize;
+        this.bufferAlignment = bufferAlignment;
+    }
+
+    /**
+     * Destroy the system of dynamic buffers.
+     */
+    destroy() {
+
+        this.gpuBuffers.forEach((gpuBuffer) => {
+            gpuBuffer.destroy(this.device);
+        });
+        this.gpuBuffers = null;
+
+        this.stagingBuffers.forEach((stagingBuffer) => {
+            stagingBuffer.destroy(this.device);
+        });
+        this.stagingBuffers = null;
+
+        this.usedBuffers = null;
+        this.activeBuffer = null;
+    }
+
+    /**
+     * Allocate an aligned space of the given size from a dynamic buffer.
+     *
+     * @param {DynamicBufferAllocation} allocation - The allocation info to fill.
+     * @param {number} size - The size of the allocation.
+     */
+    alloc(allocation, size) {
+
+        // if we have active buffer without enough space
+        if (this.activeBuffer) {
+            const alignedStart = math.roundUp(this.activeBuffer.size, this.bufferAlignment);
+            const space = this.bufferSize - alignedStart;
+            if (space < size) {
+
+                // we're done with this buffer, schedule it for submit
+                this.scheduleSubmit();
+            }
+        }
+
+        // if we don't have an active buffer, allocate new one
+        if (!this.activeBuffer) {
+
+            // gpu buffer
+            let gpuBuffer = this.gpuBuffers.pop();
+            if (!gpuBuffer) {
+                gpuBuffer = this.createBuffer(this.device, this.bufferSize, false);
+            }
+
+            // staging buffer
+            let stagingBuffer = this.stagingBuffers.pop();
+            if (!stagingBuffer) {
+                stagingBuffer = this.createBuffer(this.device, this.bufferSize, true);
+            }
+
+            this.activeBuffer = new UsedBuffer();
+            this.activeBuffer.stagingBuffer = stagingBuffer;
+            this.activeBuffer.gpuBuffer = gpuBuffer;
+            this.activeBuffer.offset = 0;
+            this.activeBuffer.size = 0;
+        }
+
+        // allocate from active buffer
+        const activeBuffer = this.activeBuffer;
+        const alignedStart = math.roundUp(activeBuffer.size, this.bufferAlignment);
+        Debug.assert(alignedStart + size <= this.bufferSize, `The allocation size of ${size} is larger than the buffer size of ${this.bufferSize}`);
+
+        allocation.gpuBuffer = activeBuffer.gpuBuffer;
+        allocation.offset = alignedStart;
+        allocation.storage = activeBuffer.stagingBuffer.alloc(alignedStart, size);
+
+        // take the allocation from the buffer
+        activeBuffer.size = alignedStart + size;
+    }
+
+    scheduleSubmit() {
+
+        if (this.activeBuffer) {
+            this.usedBuffers.push(this.activeBuffer);
+            this.activeBuffer = null;
+        }
+    }
+
+    submit() {
+
+        // schedule currently active buffer for submit
+        this.scheduleSubmit();
+    }
+}
+
+export { DynamicBuffer, DynamicBuffers, DynamicBufferAllocation };

src/platform/graphics/graphics-device.js

@@ -244,6 +244,14 @@ class GraphicsDevice extends EventHandler {
      */
     stencilBack = new StencilParameters();
 
+    /**
+     * The dynamic buffer manager.
+     *
+     * @type {import('./dynamic-buffers.js').DynamicBuffers}
+     * @ignore
+     */
+    dynamicBuffers;
+
     defaultClearOptions = {
         color: [0, 0, 0, 1],
         depth: 1,
@@ -355,6 +363,9 @@ class GraphicsDevice extends EventHandler {
 
         this.quadVertexBuffer?.destroy();
         this.quadVertexBuffer = null;
+
+        this.dynamicBuffers?.destroy();
+        this.dynamicBuffers = null;
     }
 
     onDestroyShader(shader) {

src/platform/graphics/texture.js

@@ -47,6 +47,13 @@ class Texture {
     /** @protected */
     _lockedLevel = -1;
 
+    /**
+     * A render version used to track the last time the texture properties requiring bind group
+     * to be updated were changed.
+     *
+     * @type {number}
+     * @ignore
+     */
     renderVersionDirty = 0;
 
     /**