Give compute shaders access to workgroup info

src/core/p5.Renderer3D.js

@@ -2388,7 +2388,7 @@ function renderer3D(p5, fn) {
       );
       return;
     }
-    return this.baseComputeShader().modify(cb, context, { hook: 'iteration' });
+    return this.baseComputeShader().modify(cb, context, { hook: 'computeIteration' });
   };
 
   /**

src/webgpu/p5.RendererWebGPU.js

@@ -181,7 +181,7 @@ function rendererWebGPU(p5, fn) {
      * Copies data from the GPU to the CPU using a temporary buffer,
      * so it must be awaited. Returns a `Float32Array` for number
      * buffers, or an array of plain objects for struct buffers.
-     * 
+     *
      * Note: This is a GPU -> CPU read, so calling it often (like every frame)
      * can be slow.
      *
@@ -2313,7 +2313,7 @@ function rendererWebGPU(p5, fn) {
 
       // Extract storage buffers
       const storageBuffers = {};
-      const storageRegex = /@group\((\d+)\)\s*@binding\((\d+)\)\s*var<storage,\s*(read|read_write)>\s+(\w+)\s*:\s*array<\w+>/g;
+      const storageRegex = /@group\((\d+)\)\s*@binding\((\d+)\)\s*var<storage,\s*(read|read_write)>\s+(\w+)\s*:\s*array<(\w+|atomic<\w+>)>/g;
 
       // Track which bindings are taken by the struct properties we've parsed
       // (the rest should be textures/samplers)
@@ -4027,7 +4027,7 @@ ${hookUniformFields}}
           baseComputeShader,
           {
             compute: {
-              'void iteration': '(index: vec3<i32>) {}',
+              'void computeIteration': '(inputs: ComputeInputs) {}',
             },
           }
         );

src/webgpu/shaders/compute.js

@@ -5,6 +5,14 @@ struct ComputeUniforms {
 }
 @group(0) @binding(0) var<uniform> uniforms: ComputeUniforms;
 
+struct ComputeInputs {
+  index: vec3<i32>,
+  globalID: vec3<i32>,
+  localIndex: i32,
+  localID: vec3<i32>,
+  workgroupID: vec3<i32>,
+}
+
 @compute @workgroup_size(8, 8, 1)
 fn main(
   @builtin(global_invocation_id) globalId: vec3<u32>,
@@ -19,12 +27,20 @@ fn main(
     return;
   }
 
+  var inputs: ComputeInputs;
+
   var index = vec3<i32>(0);
   index.x = i32(physicalId % u32(uniforms.uTotalCount.x));
   let remainingY = physicalId / u32(uniforms.uTotalCount.x);
   index.y = i32(remainingY % u32(uniforms.uTotalCount.y));
   index.z = i32(remainingY / u32(uniforms.uTotalCount.y));
 
-  HOOK_iteration(index);
+  inputs.index = index;
+  inputs.localID = vec3<i32>(localId);
+  inputs.workgroupID = vec3<i32>(workgroupId);
+  inputs.localIndex = i32(localIndex);
+  inputs.globalID = vec3<i32>(globalId);
+
+  HOOK_computeIteration(inputs);
 }
 `;