1 files changed, 5 insertions, 7 deletions
diff --git a/cli/tests/webgpu_computepass_shader.wgsl b/cli/tests/webgpu_computepass_shader.wgsl
index 2433f3243..7d4748e2a 100644
--- a/cli/tests/webgpu_computepass_shader.wgsl
+++ b/cli/tests/webgpu_computepass_shader.wgsl
@@ -1,14 +1,9 @@
-[[builtin(global_invocation_id)]]
-var global_id: vec3<u32>;
-
 [[block]]
 struct PrimeIndices {
     data: [[stride(4)]] array<u32>;
 }; // this is used as both input and output for convenience
-
 [[group(0), binding(0)]]
 var<storage> v_indices: [[access(read_write)]] PrimeIndices;
-
 // The Collatz Conjecture states that for any integer n:
 // If n is even, n = n/2
 // If n is odd, n = 3n+1
@@ -26,14 +21,17 @@ fn collatz_iterations(n_base: u32) -> u32{
             n = n / 2u;
         }
         else {
+            // Overflow? (i.e. 3*n + 1 > 0xffffffffu?)
+            if (n >= 1431655765u) {   // 0x55555555u
+                return 4294967295u;   // 0xffffffffu
+            }
             n = 3u * n + 1u;
         }
         i = i + 1u;
     }
     return i;
 }
-
 [[stage(compute), workgroup_size(1)]]
-fn main() {
+fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
     v_indices.data[global_id.x] = collatz_iterations(v_indices.data[global_id.x]);
 }