From 742b25968804b130df4663b643f8bf60aabd27e1 Mon Sep 17 00:00:00 2001
From: Dave Pagurek <davepagurek@gmail.com>
Date: Thu, 21 May 2026 18:40:24 -0400
Subject: [PATCH 1/2] Give compute shaders access to workgroup info

---
 src/core/p5.Renderer3D.js       |  2 +-
 src/webgpu/p5.RendererWebGPU.js |  4 ++--
 src/webgpu/shaders/compute.js   | 18 +++++++++++++++++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/core/p5.Renderer3D.js b/src/core/p5.Renderer3D.js
index 6974561b4f..3d482fae03 100644
--- a/src/core/p5.Renderer3D.js
+++ b/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' });
   };
 
   /**
diff --git a/src/webgpu/p5.RendererWebGPU.js b/src/webgpu/p5.RendererWebGPU.js
index 3b909230b7..6d06a78171 100644
--- a/src/webgpu/p5.RendererWebGPU.js
+++ b/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.
      *
@@ -4027,7 +4027,7 @@ ${hookUniformFields}}
           baseComputeShader,
           {
             compute: {
-              'void iteration': '(index: vec3<i32>) {}',
+              'void computeIteration': '(inputs: ComputeInputs) {}',
             },
           }
         );
diff --git a/src/webgpu/shaders/compute.js b/src/webgpu/shaders/compute.js
index dafe356ee6..b69d2d0a5b 100644
--- a/src/webgpu/shaders/compute.js
+++ b/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);
 }
 `;

From 87ca93451400d9e1cc3b1d942ad682df28fa09eb Mon Sep 17 00:00:00 2001
From: Dave Pagurek <davepagurek@gmail.com>
Date: Thu, 21 May 2026 19:14:29 -0400
Subject: [PATCH 2/2] Allow atomic storage in the regex

---
 src/webgpu/p5.RendererWebGPU.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/webgpu/p5.RendererWebGPU.js b/src/webgpu/p5.RendererWebGPU.js
index 6d06a78171..d5df978dc6 100644
--- a/src/webgpu/p5.RendererWebGPU.js
+++ b/src/webgpu/p5.RendererWebGPU.js
@@ -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)