WebGPU support for int and uint vertex attributes (#5050)

* WebGPU support for int and uint vertex attributes

* unused imports

* moved declaration

---------

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

src/platform/graphics/shader-processor-options.js

@@ -12,19 +12,26 @@ class ShaderProcessorOptions {
     /** @type {import('./bind-group-format.js').BindGroupFormat[]} */
     bindGroupFormats = [];
 
+    /** @type {import('./vertex-format.js').VertexFormat[]} */
+    vertexFormat;
+
     /**
      * Constructs shader processing options, used to process the shader for uniform buffer support.
      *
      * @param {import('./uniform-buffer-format.js').UniformBufferFormat} [viewUniformFormat] - Format
      * of the uniform buffer.
      * @param {import('./bind-group-format.js').BindGroupFormat} [viewBindGroupFormat] - Format of
      * the bind group.
+     * @param {import('./vertex-format.js').VertexFormat} [vertexFormat] - Format of the vertex
+     * buffer.
      */
-    constructor(viewUniformFormat, viewBindGroupFormat) {
+    constructor(viewUniformFormat, viewBindGroupFormat, vertexFormat) {
 
         // construct a sparse array
         this.uniformFormats[BINDGROUP_VIEW] = viewUniformFormat;
         this.bindGroupFormats[BINDGROUP_VIEW] = viewBindGroupFormat;
+
+        this.vertexFormat = vertexFormat;
     }
 
     /**
@@ -63,6 +70,10 @@ class ShaderProcessorOptions {
         return false;
     }
 
+    getVertexElement(semantic) {
+        return this.vertexFormat?.elements.find(element => element.name === semantic);
+    }
+
     /**
      * Generate unique key represending the processing options.
      *
@@ -71,7 +82,9 @@ class ShaderProcessorOptions {
     generateKey() {
         // TODO: Optimize. Uniform and BindGroup formats should have their keys evaluated in their
         // constructors, and here we should simply concatenate those.
-        return JSON.stringify(this);
+        return JSON.stringify(this.uniformFormats) +
+        JSON.stringify(this.bindGroupFormats) +
+        this.vertexFormat?.renderingingHashString;
     }
 }
 

src/platform/graphics/shader-processor.js

@@ -4,7 +4,8 @@ import {
     SHADERSTAGE_VERTEX, SHADERSTAGE_FRAGMENT,
     UNIFORM_BUFFER_DEFAULT_SLOT_NAME,
     SAMPLETYPE_FLOAT, SAMPLETYPE_DEPTH, SAMPLETYPE_UNFILTERABLE_FLOAT,
-    TEXTUREDIMENSION_2D, TEXTUREDIMENSION_2D_ARRAY, TEXTUREDIMENSION_CUBE, TEXTUREDIMENSION_3D
+    TEXTUREDIMENSION_2D, TEXTUREDIMENSION_2D_ARRAY, TEXTUREDIMENSION_CUBE, TEXTUREDIMENSION_3D,
+    TYPE_FLOAT32, TYPE_INT8, TYPE_INT16, TYPE_INT32
 } from './constants.js';
 import { UniformFormat, UniformBufferFormat } from './uniform-buffer-format.js';
 import { BindGroupFormat, BindBufferFormat, BindTextureFormat } from './bind-group-format.js';
@@ -74,7 +75,7 @@ class UniformLine {
 
             // simple uniform
             this.name = words.shift();
-            this.arraySize = 1;
+            this.arraySize = 0;
         }
 
         this.isSampler = this.type.indexOf('sampler') !== -1;
@@ -106,7 +107,7 @@ class ShaderProcessor {
         const fragmentExtracted = ShaderProcessor.extract(shaderDefinition.fshader);
 
         // VS - convert a list of attributes to a shader block with fixed locations
-        const attributesBlock = ShaderProcessor.processAttributes(vertexExtracted.attributes, shaderDefinition.attributes);
+        const attributesBlock = ShaderProcessor.processAttributes(vertexExtracted.attributes, shaderDefinition.attributes, shaderDefinition.processingOptions);
 
         // VS - convert a list of varyings to a shader block
         const vertexVaryingsBlock = ShaderProcessor.processVaryings(vertexExtracted.varyings, varyingMap, true);
@@ -353,13 +354,20 @@ class ShaderProcessor {
         return block;
     }
 
-    static processAttributes(attributeLines, shaderDefinitionAttributes) {
+    // extract count from type ('vec3' => 3, 'float' => 1)
+    static getTypeCount(type) {
+        const lastChar = type.substring(type.length - 1);
+        const num = parseInt(lastChar, 10);
+        return isNaN(num) ? 1 : num;
+    }
+
+    static processAttributes(attributeLines, shaderDefinitionAttributes, processingOptions) {
         let block = '';
         const usedLocations = {};
         attributeLines.forEach((line) => {
             const words = ShaderProcessor.splitToWords(line);
-            const type = words[0];
-            const name = words[1];
+            let type = words[0];
+            let name = words[1];
 
             if (shaderDefinitionAttributes.hasOwnProperty(name)) {
                 const semantic = shaderDefinitionAttributes[name];
@@ -369,8 +377,41 @@ class ShaderProcessor {
                              `WARNING: Two vertex attributes are mapped to the same location in a shader: ${usedLocations[location]} and ${semantic}`);
                 usedLocations[location] = semantic;
 
+                // if vertex format for this attribute is not of a float type, we need to adjust the attribute format, for example we convert
+                //      attribute vec4 vertex_position;
+                // to
+                //      attribute ivec4 _private_vertex_position;
+                //      vec4 vertex_position = vec4(_private_vertex_position);
+                let copyCode;
+                const element = processingOptions.getVertexElement(semantic);
+                if (element) {
+                    const dataType = element.dataType;
+                    if (dataType !== TYPE_FLOAT32) {
+
+                        const attribNumElements = ShaderProcessor.getTypeCount(type);
+                        const newName = `_private_${name}`;
+
+                        // second line of new code, copy private (u)int type into vec type
+                        copyCode = `vec${attribNumElements} ${name} = vec${attribNumElements}(${newName});\n`;
+
+                        name = newName;
+
+                        // new attribute type, based on the vertex format element type, example: vec3 -> ivec3
+                        const isSignedType = dataType === TYPE_INT8 || dataType === TYPE_INT16 || dataType === TYPE_INT32;
+                        if (attribNumElements === 1) {
+                            type = isSignedType ? 'int' : 'uint';
+                        } else {
+                            type = isSignedType ? `ivec${attribNumElements}` : `uvec${attribNumElements}`;
+                        }
+                    }
+                }
+
                 // generates: 'layout(location = 0) in vec4 position;'
                 block += `layout(location = ${location}) in ${type} ${name};\n`;
+
+                if (copyCode) {
+                    block += copyCode;
+                }
             }
         });
         return block;

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

@@ -54,24 +54,24 @@ class UniformFormat {
     scopeId;
 
     /**
-     * Count of elements for arrays, otherwise 1.
+     * Count of elements for arrays, otherwise 0.
      *
      * @type {number}
      */
     count;
 
-    constructor(name, type, count = 1) {
+    constructor(name, type, count = 0) {
 
         // just a name
         this.shortName = name;
 
         // name with [0] if this is an array
-        this.name = count > 1 ? `${name}[0]` : name;
+        this.name = count ? `${name}[0]` : name;
 
         this.type = type;
 
         this.updateType = type;
-        if (count > 1) {
+        if (count) {
 
             switch (type) {
                 case UNIFORMTYPE_FLOAT: this.updateType = UNIFORMTYPE_FLOATARRAY; break;
@@ -100,10 +100,13 @@ class UniformFormat {
         Debug.assert(elementSize, `Unhandled uniform format ${type} used for ${name}`);
 
         // element size for arrays is aligned up to vec4
-        if (count > 1)
+        if (count)
             elementSize = math.roundUp(elementSize, 4);
 
-        this.byteSize = count * elementSize * 4;
+        this.byteSize = elementSize * 4;
+        if (count)
+            this.byteSize *= count;
+
         Debug.assert(this.byteSize, `Unknown byte size for uniform format ${type} used for ${name}`);
     }
 
@@ -115,7 +118,7 @@ class UniformFormat {
         let alignment = this.byteSize <= 8 ? this.byteSize : 16;
 
         // arrays have vec4 alignments
-        if (this.count > 1)
+        if (this.count)
             alignment = 16;
 
         // align the start offset
@@ -183,7 +186,7 @@ class UniformBufferFormat {
         this.uniforms.forEach((uniform) => {
             const typeString = uniformTypeToName[uniform.type];
             Debug.assert(typeString.length > 0, `Uniform type ${uniform.type} is not handled.`);
-            code += `    ${typeString} ${uniform.shortName}${uniform.count !== 1 ? `[${uniform.count}]` : ''};\n`;
+            code += `    ${typeString} ${uniform.shortName}${uniform.count ? `[${uniform.count}]` : ''};\n`;
         });
 
         return code + '};\n';

src/platform/graphics/webgpu/webgpu-vertex-buffer-layout.js

@@ -1,10 +1,7 @@
-import { Debug } from "../../../core/debug.js";
-
 import {
     semanticToLocation,
-    TYPE_INT8, TYPE_UINT8, TYPE_INT16, TYPE_UINT16, TYPE_INT32, TYPE_UINT32, TYPE_FLOAT32, vertexTypesNames
+    TYPE_INT8, TYPE_UINT8, TYPE_INT16, TYPE_UINT16, TYPE_INT32, TYPE_UINT32, TYPE_FLOAT32
 } from '../constants.js';
-import { DebugGraphics } from "../debug-graphics.js";
 
 // map of TYPE_*** to GPUVertexFormat
 const gpuVertexFormats = [];
@@ -67,9 +64,6 @@ class WebgpuVertexBufferLayout {
                 const element = format.elements[i];
                 const location = semanticToLocation[element.name];
 
-                // A WGL shader needs attributes to have matching types, but glslang translator we use does not allow us to set those
-                Debug.assert(element.dataType === TYPE_FLOAT32, `Only float vertex attributes are supported, ${vertexTypesNames[element.dataType]} is not supported, semantic: ${element.name}, stack: ${DebugGraphics.toString()}`, element);
-
                 attributes.push({
                     shaderLocation: location,
                     offset: interleaved ? element.offset : 0,

src/scene/materials/basic-material.js

@@ -82,7 +82,7 @@ class BasicMaterial extends Material {
         }
     }
 
-    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat) {
+    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat, vertexFormat) {
 
         // Note: this is deprecated function Editor and possibly other projects use: they define
         // updateShader callback on their BasicMaterial, so we handle it here.
@@ -106,7 +106,7 @@ class BasicMaterial extends Material {
             pass: pass
         };
 
-        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat);
+        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat, vertexFormat);
 
         const library = getProgramLibrary(device);
         library.register('basic', basic);

src/scene/materials/material.js

@@ -428,10 +428,10 @@ class Material {
     updateUniforms(device, scene) {
     }
 
-    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat) {
+    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat, vertexFormat) {
 
         // generate shader variant - its the same shader, but with different processing options
-        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat);
+        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat, vertexFormat);
         return processShader(this._shader, processingOptions);
     }
 

src/scene/materials/standard-material.js

@@ -828,7 +828,7 @@ class StandardMaterial extends Material {
         this._processParameters('_activeLightingParams');
     }
 
-    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat) {
+    getShaderVariant(device, scene, objDefs, staticLightList, pass, sortedLights, viewUniformFormat, viewBindGroupFormat, vertexFormat) {
 
         // update prefiltered lighting data
         this.updateEnvUniforms(device, scene);
@@ -847,7 +847,7 @@ class StandardMaterial extends Material {
             options = this.onUpdateShader(options);
         }
 
-        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat);
+        const processingOptions = new ShaderProcessorOptions(viewUniformFormat, viewBindGroupFormat, vertexFormat);
 
         const library = getProgramLibrary(device);
         library.register('standard', standard);

src/scene/mesh-instance.js

@@ -742,7 +742,7 @@ class MeshInstance {
      */
     updatePassShader(scene, pass, staticLightList, sortedLights, viewUniformFormat, viewBindGroupFormat) {
         this._shader[pass] = this.material.getShaderVariant(this.mesh.device, scene, this._shaderDefs, staticLightList, pass, sortedLights,
-                                                            viewUniformFormat, viewBindGroupFormat);
+                                                            viewUniformFormat, viewBindGroupFormat, this._mesh.vertexBuffer.format);
     }
 
     ensureMaterial(device) {

src/scene/shader-lib/chunks/common/vert/skinTex.js

@@ -1,21 +1,12 @@
 export default /* glsl */`
 
 attribute vec4 vertex_boneWeights;
-
-#ifdef WEBGPU
-    attribute uvec4 vertex_boneIndices;
-    #define BoneIndexFormat uint
-    #define BoneIndexFormat4 uvec4
-#else
-    attribute vec4 vertex_boneIndices;
-    #define BoneIndexFormat float
-    #define BoneIndexFormat4 vec4
-#endif
+attribute vec4 vertex_boneIndices;
 
 uniform highp sampler2D texture_poseMap;
 uniform vec4 texture_poseMapSize;
 
-void getBoneMatrix(const in BoneIndexFormat index, out vec4 v1, out vec4 v2, out vec4 v3) {
+void getBoneMatrix(const in float index, out vec4 v1, out vec4 v2, out vec4 v3) {
 
     float i = float(index);
     float j = i * 3.0;
@@ -32,7 +23,7 @@ void getBoneMatrix(const in BoneIndexFormat index, out vec4 v1, out vec4 v2, out
     v3 = texture2D(texture_poseMap, vec2(dx * (x + 2.5), y));
 }
 
-mat4 getSkinMatrix(const in BoneIndexFormat4 indices, const in vec4 weights) {
+mat4 getSkinMatrix(const in vec4 indices, const in vec4 weights) {
     // get 4 bone matrices
     vec4 a1, a2, a3;
     getBoneMatrix(indices.x, a1, a2, a3);