feat: loop unrolling support (#2178)

Author: cieplypolar

apps/typegpu-docs/src/content/docs/fundamentals/utils.mdx

@@ -267,3 +267,59 @@ Otherwise, for example when using `tgpu.resolve` on a WGSL template, logs are ig
 - `console.log` only works in fragment and compute shaders.
 This is due to a [WebGPU limitation](https://www.w3.org/TR/WGSL/#address-space) that does not allow modifying buffers during the vertex shader stage.
 - `console.log` currently does not support template literals (but you can use [string substitutions](https://developer.mozilla.org/en-US/docs/Web/API/console#using_string_substitutions), or just pass multiple arguments instead).
+
+## *for...of...* loops
+
+TypeGPU supports `for...of...` loops in shader functions. The only constraints are that the loop variable must be declared with `const` and the iterable must be stored in a variable.
+
+```ts twoslash
+import tgpu, { d } from 'typegpu';
+
+const processNeighbor = (cell: d.v2i) => {};
+
+// ---cut---
+const processNeighbors = (cell: d.v2i) => {
+  'use gpu';
+
+  const offsets = [
+    d.vec2i(0, 1),
+    d.vec2i(0, -1),
+    d.vec2i(1, 0),
+    d.vec2i(-1, 0),
+  ];
+
+  for (const offset of offsets) {
+    processNeighbor(cell.add(offset));
+  }
+};
+```
+
+## *tgpu.unroll*
+
+For code with small, fixed iteration counts, you can use `tgpu.unroll` to unroll loops at compile time. This eliminates branch prediction overhead and can significantly improve performance.
+
+### Usage
+
+Wrap your iterable with `tgpu.unroll()`:
+
+```ts twoslash
+import tgpu, { d } from 'typegpu';
+
+const processNeighbor = (cell: d.v2i) => {};
+
+// ---cut---
+const processNeighbors = (cell: d.v2i) => {
+  'use gpu';
+
+  for (const dy of tgpu.unroll([-1, 0, 1])) {
+    for (const dx of tgpu.unroll([-1, 0, 1])) {
+      processNeighbor(cell.add(d.vec2i(dx, dy)));
+    }
+  }
+};
+```
+
+:::note
+- There are no constraints on how large a loop can be for unrolling. We will always try to unroll it, and if we can't, you'll receive an error.
+- You cannot use `continue` or `break` inside loop that you intend to unroll later.
+:::

apps/typegpu-docs/src/examples/algorithms/jump-flood-distance/index.ts

@@ -186,8 +186,8 @@ const jumpFlood = root.createGuardedComputePipeline((x, y) => {
   let bestInsideDist = 1e20;
   let bestOutsideDist = 1e20;
 
-  for (let dy = -1; dy <= 1; dy++) {
-    for (let dx = -1; dx <= 1; dx++) {
+  for (const dx of tgpu.unroll([-1, 0, 1])) {
+    for (const dy of tgpu.unroll([-1, 0, 1])) {
       const sample = sampleWithOffset(
         pingPongLayout.$.readView,
         d.vec2i(x, y),

apps/typegpu-docs/src/examples/algorithms/jump-flood-voronoi/index.ts

@@ -165,22 +165,23 @@ const jumpFlood = root.createGuardedComputePipeline((x, y) => {
   let minDist = 1e20;
   let bestSample = SampleResult({ color: d.vec4f(), coord: d.vec2f(-1) });
 
-  for (let dy = -1; dy <= 1; dy++) {
-    for (let dx = -1; dx <= 1; dx++) {
+  for (const dy of tgpu.unroll([-1, 0, 1])) {
+    for (const dx of tgpu.unroll([-1, 0, 1])) {
       const sample = sampleWithOffset(
         pingPongLayout.$.readView,
         d.vec2i(x, y),
         d.vec2i(dx * offset, dy * offset),
       );
 
-      if (sample.coord.x < 0) {
-        continue;
-      }
-
-      const dist = std.distance(d.vec2f(x, y), sample.coord.mul(d.vec2f(size)));
-      if (dist < minDist) {
-        minDist = dist;
-        bestSample = SampleResult(sample);
+      if (sample.coord.x >= 0) {
+        const dist = std.distance(
+          d.vec2f(x, y),
+          sample.coord.mul(d.vec2f(size)),
+        );
+        if (dist < minDist) {
+          minDist = dist;
+          bestSample = SampleResult(sample);
+        }
       }
     }
   }

apps/typegpu-docs/src/examples/image-processing/background-segmentation/shaders.ts

@@ -56,8 +56,8 @@ export const computeFn = tgpu.computeFn({
   ).sub(d.vec2i(filterOffset, 0));
 
   // Load a tile of pixels into shared memory
-  for (let r = 0; r < 4; r++) {
-    for (let c = 0; c < 4; c++) {
+  for (const r of tgpu.unroll([0, 1, 2, 3])) {
+    for (const c of tgpu.unroll([0, 1, 2, 3])) {
       let loadIndex = baseIndex.add(d.vec2i(c, r));
       if (flipAccess.$) {
         loadIndex = loadIndex.yx;
@@ -75,8 +75,8 @@ export const computeFn = tgpu.computeFn({
   std.workgroupBarrier();
 
   // Apply the horizontal blur filter and write to the output texture
-  for (let r = 0; r < 4; r++) {
-    for (let c = 0; c < 4; c++) {
+  for (const r of tgpu.unroll([0, 1, 2, 3])) {
+    for (const c of tgpu.unroll([0, 1, 2, 3])) {
       let writeIndex = baseIndex.add(d.vec2i(c, r));
       if (flipAccess.$) {
         writeIndex = writeIndex.yx;

apps/typegpu-docs/src/examples/image-processing/blur/index.ts

@@ -73,8 +73,8 @@ const computeFn = tgpu.computeFn({
   ).sub(d.vec2i(filterOffset, 0));
 
   // Load a tile of pixels into shared memory
-  for (let r = 0; r < 4; r++) {
-    for (let c = 0; c < 4; c++) {
+  for (const r of tgpu.unroll([0, 1, 2, 3])) {
+    for (const c of tgpu.unroll([0, 1, 2, 3])) {
       let loadIndex = baseIndex.add(d.vec2i(c, r));
       if (ioLayout.$.flip !== 0) {
         loadIndex = loadIndex.yx;
@@ -92,8 +92,8 @@ const computeFn = tgpu.computeFn({
   std.workgroupBarrier();
 
   // Apply the horizontal blur filter and write to the output texture
-  for (let r = 0; r < 4; r++) {
-    for (let c = 0; c < 4; c++) {
+  for (const r of tgpu.unroll([0, 1, 2, 3])) {
+    for (const c of tgpu.unroll([0, 1, 2, 3])) {
       let writeIndex = baseIndex.add(d.vec2i(c, r));
       if (ioLayout.$.flip !== 0) {
         writeIndex = writeIndex.yx;

apps/typegpu-docs/src/examples/rendering/3d-fish/compute.ts

@@ -1,4 +1,4 @@
-import { d, std } from 'typegpu';
+import tgpu, { d, std } from 'typegpu';
 import * as p from './params.ts';
 import { computeBindGroupLayout as layout } from './schemas.ts';
 import { projectPointOnLine } from './tgsl-helpers.ts';
@@ -39,7 +39,7 @@ export const simulate = (fishIndex: number) => {
   if (cohesionCount > 0) {
     cohesion = cohesion / cohesionCount - fishData.position;
   }
-  for (let i = 0; i < 3; i += 1) {
+  for (const i of tgpu.unroll([0, 1, 2])) {
     const repulsion = d.vec3f();
     repulsion[i] = 1;
 

apps/typegpu-docs/src/examples/rendering/clouds/utils.ts

@@ -72,12 +72,13 @@ export const raymarch = tgpu.fn([d.vec3f, d.vec3f, d.vec3f], d.vec4f)(
   },
 );
 
+const iterations = Array.from({ length: FBM_OCTAVES }, (_, i) => i);
 const fbm = tgpu.fn([d.vec3f], d.f32)((pos) => {
   let sum = d.f32();
   let amp = d.f32(CLOUD_AMPLITUDE);
   let freq = d.f32(CLOUD_FREQUENCY);
 
-  for (let i = 0; i < FBM_OCTAVES; i++) {
+  for (const _i of tgpu.unroll(iterations)) {
     sum += noise3d(std.mul(pos, freq)) * amp;
     amp *= FBM_PERSISTENCE;
     freq *= FBM_LACUNARITY;

apps/typegpu-docs/src/examples/rendering/cubemap-reflection/icosphere.ts

@@ -171,7 +171,7 @@ export class IcosphereGenerator {
       ]);
 
       const baseIndexNext = triangleIndex * 12;
-      for (let i = d.u32(0); i < 12; i++) {
+      for (const i of tgpu.unroll([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])) {
         const reprojectedVertex = newVertices[i];
 
         const triBase = i - (i % 3);

apps/typegpu-docs/src/examples/rendering/jelly-slider/taa.ts

@@ -25,8 +25,8 @@ export const taaResolveFn = tgpu.computeFn({
 
   const dimensions = std.textureDimensions(taaResolveLayout.$.currentTexture);
 
-  for (let x = -1; x <= 1; x++) {
-    for (let y = -1; y <= 1; y++) {
+  for (const x of tgpu.unroll([-1, 0, 1])) {
+    for (const y of tgpu.unroll([-1, 0, 1])) {
       const sampleCoord = d.vec2i(gid.xy).add(d.vec2i(x, y));
       const clampedCoord = std.clamp(
         sampleCoord,

apps/typegpu-docs/src/examples/rendering/jelly-switch/taa.ts

@@ -25,8 +25,8 @@ export const taaResolveFn = tgpu.computeFn({
 
   const dimensions = std.textureDimensions(taaResolveLayout.$.currentTexture);
 
-  for (let x = -1; x <= 1; x++) {
-    for (let y = -1; y <= 1; y++) {
+  for (const x of tgpu.unroll([-1, 0, 1])) {
+    for (const y of tgpu.unroll([-1, 0, 1])) {
       const sampleCoord = d.vec2i(gid.xy).add(d.vec2i(x, y));
       const clampedCoord = std.clamp(
         sampleCoord,