Add support for instanceID() in WebGPU mode

Author: davepagurek

preview/index.html

@@ -24,12 +24,13 @@
 
     const sketch = function (p) {
       let fbo;
-      let sh;
+      let sh, sh2;
       let ssh;
       let tex;
       let font;
       let redFilter;
       let env;
+      let instance;
 
       p.setup = async function () {
         await p.createCanvas(400, 400, p.WEBGPU);
@@ -39,6 +40,8 @@
         );
         fbo = p.createFramebuffer();
 
+        instance = p.buildGeometry(() => p.sphere(5));
+
         redFilter = p.baseFilterShader().modify(() => {
           p.getColor((inputs, canvasContent) => {
             let col = p.getTexture(canvasContent, inputs.texCoord);
@@ -84,6 +87,12 @@
             return inputs;
           });
         }, { p })
+        sh2 = p.baseMaterialShader().modify(() => {
+          p.getWorldInputs((inputs) => {
+            inputs.position.x += 20 * p.instanceID();
+            return inputs;
+          });
+        }, { p })
 	/*ssh = p.baseStrokeShader().modify({
           uniforms: {
             'f32 time': () => p.millis(),
@@ -151,6 +160,13 @@
         }
         p.pop();
 
+        p.push();
+        p.shader(sh2);
+        p.noStroke();
+        p.fill('red');
+        p.model(instance, 10);
+        p.pop();
+
         // Test beginShape/endShape with immediate mode shapes
         p.push();
         p.translate(0, 100, 0);

src/strands/strands_api.js

@@ -50,7 +50,7 @@ export function initGlobalStrandsAPI(p5, fn, strandsContext) {
   };
   p5.break = fn.break;
   fn.instanceID = function() {
-    const node = build.variableNode(strandsContext, { baseType: BaseType.INT, dimension: 1 }, 'gl_InstanceID');
+    const node = build.variableNode(strandsContext, { baseType: BaseType.INT, dimension: 1 }, strandsContext.backend.instanceIdReference());
     return createStrandsNode(node.id, node.dimension, strandsContext);
   }
   // Internal methods use p5 static methods; user-facing methods use fn.

src/webgl/strands_glslBackend.js

@@ -379,5 +379,9 @@ export const glslBackend = {
       }]
     });
     return { id, dimension };
-  }
+  },
+
+  instanceIdReference() {
+    return 'gl_InstanceID';
+  },
 }

src/webgpu/p5.RendererWebGPU.js

@@ -289,7 +289,7 @@ function rendererWebGPU(p5, fn) {
       // Check if we already have a buffer for this data
       let existingBuffer = buffers[dst];
       const needsNewBuffer = !existingBuffer;
-      
+
       // Only create new buffer and write data if buffer doesn't exist or data is dirty
       if (needsNewBuffer || geometry.dirtyFlags[src] !== false) {
         const raw = map ? map(srcData) : srcData;
@@ -1349,7 +1349,7 @@ function rendererWebGPU(p5, fn) {
     }
 
     _getShaderAttributes(shader) {
-      const mainMatch = /fn main\(.+:\s*(\S+)\s*\)/.exec(shader._vertSrc);
+      const mainMatch = /fn main\(.+:\s*([^\s\)]+)/.exec(shader._vertSrc);
       if (!mainMatch) throw new Error("Can't find `fn main` in vertex shader source");
       const inputType = mainMatch[1];
 
@@ -1740,7 +1740,12 @@ function rendererWebGPU(p5, fn) {
         }
       );
 
-      let [preMain, main, postMain] = src.split(/((?:@(?:vertex|fragment)\s*)?fn main)/);
+      let [preMain, main, postMain] = src.split(/((?:@(?:vertex|fragment)\s*)?fn main[^{]+\{)/);
+      if (shaderType !== 'fragment') {
+        if (!main.match(/\@builtin\s*\(\s*instance_index\s*\)/)) {
+          main = main.replace(/\)\s*(->|\{)/, ', @builtin(instance_index) instanceID: u32) $1');
+        }
+      }
 
       let uniforms = '';
       for (const key in shader.hooks.uniforms) {
@@ -1850,14 +1855,57 @@ function rendererWebGPU(p5, fn) {
           shader.hooks.modified[shaderType][hookDef] ? 'true' : 'false'
         };\n`;
 
-        const [_, params, body] = /^(\([^\)]*\))((?:.|\n)*)$/.exec(shader.hooks[shaderType][hookDef]);
+        let [_, params, body] = /^(\([^\)]*\))((?:.|\n)*)$/.exec(shader.hooks[shaderType][hookDef]);
+
+        if (shaderType !== 'fragment') {
+          // Splice the instance ID in as a final parameter to every WGSL hook function
+          let hasParams = !!params.match(/^\(\s*\S+.*\)$/);
+          params = params.slice(0, -1) + (hasParams ? ', ' : '') + 'instanceID: u32)';
+        }
+
         if (hookType === 'void') {
           hooks += `fn HOOK_${hookName}${params}${body}\n`;
         } else {
           hooks += `fn HOOK_${hookName}${params} -> ${hookType}${body}\n`;
         }
       }
 
+      // Add the instance ID as a final parameter to each hook call
+      if (shaderType !== 'fragment') {
+        const addInstanceIDParam = (src) => {
+          let result = src;
+          let idx = 0;
+          let match;
+          do {
+            match = /HOOK_\w+\(/.exec(result.slice(idx));
+            if (match) {
+              idx += match.index + match[0].length - 1;
+              let nesting = 0;
+              let hasParams = false;
+              while (idx < result.length) {
+                if (result[idx] === '(') {
+                  nesting++;
+                } else if (result[idx] === ')') {
+                  nesting--;
+                } else if (result[idx].match(/\S/)) {
+                  hasParams = true;
+                }
+                idx++;
+                if (nesting === 0) {
+                  break;
+                }
+              }
+              const insertion = (hasParams ? ', ' : '') + 'instanceID';
+              result = result.slice(0, idx-1) + insertion + result.slice(idx-1);
+              idx += insertion.length;
+            }
+          } while (match);
+          return result;
+        };
+        preMain = addInstanceIDParam(preMain);
+        postMain = addInstanceIDParam(postMain);
+      }
+
       return preMain + '\n' + defines + hooks + main + postMain;
     }
 

src/webgpu/shaders/color.js

@@ -46,7 +46,7 @@ fn main(input: VertexInput) -> VertexOutput {
   HOOK_beforeVertex();
   var output: VertexOutput;
 
-  let useVertexColor = (uniforms.uUseVertexColor != 0);
+  let useVertexColor = (uniforms.uUseVertexColor != 0 && input.aVertexColor.x >= 0.0);
   var inputs = Vertex(
     input.aPosition,
     input.aNormal,

src/webgpu/shaders/material.js

@@ -87,7 +87,7 @@ fn main(input: VertexInput) -> VertexOutput {
   HOOK_beforeVertex();
   var output: VertexOutput;
 
-  let useVertexColor = (uniforms.uUseVertexColor != 0);
+  let useVertexColor = (uniforms.uUseVertexColor != 0 && input.aVertexColor.x >= 0.0);
   var inputs = Vertex(
     input.aPosition,
     input.aNormal,

src/webgpu/strands_wgslBackend.js

@@ -503,5 +503,9 @@ export const wgslBackend = {
       }]
     });
     return { id, dimension };
-  }
+  },
+
+  instanceIdReference() {
+    return 'instanceID';
+  },
 }

test/unit/visual/cases/webgpu.js

@@ -114,6 +114,22 @@ visualSuite("WebGPU", function () {
         await screenshot();
       },
     );
+
+    visualTest('Instanced rendering', async function(p5, screenshot) {
+      await p5.createCanvas(50, 50, p5.WEBGPU);
+      const model = p5.buildGeometry(() => p5.sphere(5));
+      const shader = p5.baseMaterialShader().modify(() => {
+        p5.getWorldInputs((inputs) => {
+          inputs.position += (p5.instanceID() - 1) * 15
+          return inputs;
+        });
+      }, { p5 });
+      p5.noStroke();
+      p5.fill(0);
+      p5.shader(shader);
+      p5.model(model, 3);
+      await screenshot();
+    });
   });
 
   visualSuite('filters', function() {

test/unit/visual/screenshots/WebGPU/Shaders/Instanced rendering/000.png

@@ -0,0 +1,3 @@
+{
+  "numScreenshots": 1
+}