From 92e9118d9c365619222ad0f02dc769e6613714e9 Mon Sep 17 00:00:00 2001
From: Martin Valigursky <mvaligursky@snapchat.com>
Date: Tue, 23 Apr 2024 12:03:03 +0100
Subject: [PATCH] Fix to WebGPU vertex buffer setup

---
 .../graphics/mesh-morph-many/config.mjs       |  2 +-
 .../graphics/webgpu/webgpu-graphics-device.js | 36 +++++++++++++++++--
 .../webgpu/webgpu-vertex-buffer-layout.js     |  6 ++++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/examples/src/examples/graphics/mesh-morph-many/config.mjs b/examples/src/examples/graphics/mesh-morph-many/config.mjs
index a88f9ed7363..b7fc3987320 100644
--- a/examples/src/examples/graphics/mesh-morph-many/config.mjs
+++ b/examples/src/examples/graphics/mesh-morph-many/config.mjs
@@ -2,5 +2,5 @@
  * @type {import('../../../../types.mjs').ExampleConfig}
  */
 export default {
-    WEBGPU_ENABLED: false
+    WEBGPU_ENABLED: true
 };
diff --git a/src/platform/graphics/webgpu/webgpu-graphics-device.js b/src/platform/graphics/webgpu/webgpu-graphics-device.js
index 7ac173c7ef3..ce24cb14f93 100644
--- a/src/platform/graphics/webgpu/webgpu-graphics-device.js
+++ b/src/platform/graphics/webgpu/webgpu-graphics-device.js
@@ -4,7 +4,7 @@ import { path } from '../../../core/path.js';
 
 import {
     PIXELFORMAT_RGBA32F, PIXELFORMAT_RGBA8, PIXELFORMAT_BGRA8, DEVICETYPE_WEBGPU,
-    BUFFERUSAGE_READ, BUFFERUSAGE_COPY_DST
+    BUFFERUSAGE_READ, BUFFERUSAGE_COPY_DST, semanticToLocation
 } from '../constants.js';
 import { GraphicsDevice } from '../graphics-device.js';
 import { DebugGraphics } from '../debug-graphics.js';
@@ -30,6 +30,8 @@ import { WebgpuResolver } from './webgpu-resolver.js';
 import { WebgpuCompute } from './webgpu-compute.js';
 import { WebgpuBuffer } from './webgpu-buffer.js';
 
+const _uniqueLocations = new Map();
+
 class WebgpuGraphicsDevice extends GraphicsDevice {
     /**
      * Object responsible for caching and creation of render pipelines.
@@ -452,9 +454,18 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
 
     submitVertexBuffer(vertexBuffer, slot) {
 
-        const elements = vertexBuffer.format.elements;
+        const format = vertexBuffer.format;
+        const { interleaved, elements } = format;
         const elementCount = elements.length;
         const vbBuffer = vertexBuffer.impl.buffer;
+
+        if (interleaved) {
+            // for interleaved buffers, we use a single vertex buffer, and attributes are specified using the layout
+            this.passEncoder.setVertexBuffer(slot, vbBuffer);
+            return 1;
+        }
+
+        // non-interleaved - vertex buffer per attribute
         for (let i = 0; i < elementCount; i++) {
             this.passEncoder.setVertexBuffer(slot + i, vbBuffer, elements[i].offset);
         }
@@ -462,6 +473,26 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
         return elementCount;
     }
 
+    validateVBLocations(vb0, vb1) {
+
+        // in case of multiple VBs, validate all elements use unique locations
+        const validateVB = (vb) => {
+            const { elements } = vb.format;
+            for (let i = 0; i < elements.length; i++) {
+                const name = elements[i].name;
+                const location = semanticToLocation[name];
+                if (_uniqueLocations.has(location)) {
+                    Debug.errorOnce(`Vertex buffer element location ${location} used by [${name}] is already used by element [${_uniqueLocations.get(location)}], while rendering [${DebugGraphics.toString()}]`);
+                }
+                _uniqueLocations.set(location, name);
+            }
+        };
+
+        validateVB(vb0);
+        validateVB(vb1);
+        _uniqueLocations.clear();
+    }
+
     draw(primitive, numInstances = 1, keepBuffers) {
 
         if (this.shader.ready && !this.shader.failed) {
@@ -479,6 +510,7 @@ class WebgpuGraphicsDevice extends GraphicsDevice {
             if (vb0) {
                 const vbSlot = this.submitVertexBuffer(vb0, 0);
                 if (vb1) {
+                    Debug.call(() => this.validateVBLocations(vb0, vb1));
                     this.submitVertexBuffer(vb1, vbSlot);
                 }
             }
diff --git a/src/platform/graphics/webgpu/webgpu-vertex-buffer-layout.js b/src/platform/graphics/webgpu/webgpu-vertex-buffer-layout.js
index 2d442187fb4..6be0eda5a35 100644
--- a/src/platform/graphics/webgpu/webgpu-vertex-buffer-layout.js
+++ b/src/platform/graphics/webgpu/webgpu-vertex-buffer-layout.js
@@ -66,6 +66,12 @@ class WebgpuVertexBufferLayout {
         // type  {GPUVertexBufferLayout[]}
         const layout = [];
 
+        // Note: If the VertexFormat is interleaved, we use a single vertex buffer with multiple
+        // attributes. This uses a smaller number of vertex buffers (1), which has performance
+        // benefits when setting it up on the device.
+        // If the VertexFormat is not interleaved, we use multiple vertex buffers, one per
+        // attribute. This is less efficient, but is required as there is a pretty small
+        // limit on the attribute offsets in the vertex buffer layout.
         const addFormat = (format) => {
             const interleaved = format.interleaved;
             const stepMode = format.instancing ? 'instance' : 'vertex';