feat: loop unrolling support

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,

apps/typegpu-docs/src/examples/simple/liquid-glass/index.ts

@@ -95,7 +95,7 @@ const sampleWithChromaticAberration = (
 ) => {
   'use gpu';
   const samples = d.arrayOf(d.vec3f, 3)();
-  for (let i = 0; i < 3; i++) {
+  for (const i of tgpu.unroll([0, 1, 2])) {
     const channelOffset = dir.mul((d.f32(i) - 1.0) * offset);
     samples[i] =
       std.textureSampleBias(tex, sampler, uv.sub(channelOffset), blur).rgb;

apps/typegpu-docs/src/examples/simple/ripple-cube/pbr.ts

@@ -86,13 +86,14 @@ export const evaluateLight = (
     .mul(ndotl);
 };
 
+const lightCountIterations = Array.from({ length: LIGHT_COUNT }, (_, i) => i);
 export const shade = (p: d.v3f, n: d.v3f, v: d.v3f): d.v3f => {
   'use gpu';
   const material = materialAccess.$;
   const f0 = std.mix(d.vec3f(0.04), material.albedo, material.metallic);
 
   let lo = d.vec3f(0);
-  for (let i = 0; i < LIGHT_COUNT; i++) {
+  for (const i of tgpu.unroll(lightCountIterations)) {
     lo = lo.add(evaluateLight(p, n, v, lightsAccess.$[i], material, f0));
   }
 

apps/typegpu-docs/src/examples/simple/ripple-cube/post-processing.ts

@@ -81,8 +81,8 @@ export function createPostProcessingPipelines(
     let minColor = d.vec3f(9999);
     let maxColor = d.vec3f(-9999);
 
-    for (let ox = -1; ox <= 1; ox++) {
-      for (let oy = -1; oy <= 1; oy++) {
+    for (const ox of tgpu.unroll([-1, 0, 1])) {
+      for (const oy of tgpu.unroll([-1, 0, 1])) {
         const sampleCoord = coord.add(d.vec2i(ox, oy));
         const clampedCoord = std.clamp(
           sampleCoord,

apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts

@@ -160,21 +160,19 @@ const computeVelocity = (x: number, y: number): d.v2f => {
   ];
   let dirChoiceCount = 1;
 
-  for (const offset of neighborOffsets) {
+  for (const offset of tgpu.unroll(neighborOffsets)) {
     const neighborDensity = getCell(x + offset.x, y + offset.y);
     const cost = neighborDensity.z + d.f32(offset.y) * gravityCost;
 
-    if (!isValidFlowOut(x + offset.x, y + offset.y)) {
-      continue;
-    }
-
-    if (cost === leastCost) {
-      dirChoices[dirChoiceCount] = d.vec2f(d.f32(offset.x), d.f32(offset.y));
-      dirChoiceCount++;
-    } else if (cost < leastCost) {
-      leastCost = cost;
-      dirChoices[0] = d.vec2f(d.f32(offset.x), d.f32(offset.y));
-      dirChoiceCount = 1;
+    if (isValidFlowOut(x + offset.x, y + offset.y)) {
+      if (cost === leastCost) {
+        dirChoices[dirChoiceCount] = d.vec2f(d.f32(offset.x), d.f32(offset.y));
+        dirChoiceCount++;
+      } else if (cost < leastCost) {
+        leastCost = cost;
+        dirChoices[0] = d.vec2f(d.f32(offset.x), d.f32(offset.y));
+        dirChoiceCount = 1;
+      }
     }
   }
 

apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts

@@ -176,6 +176,8 @@ const getPerpendicular = (dir: d.v3f) => {
   return std.normalize(std.cross(dir, axis));
 };
 
+const numSamples = 8;
+const samplesIterations = Array.from({ length: numSamples }, (_, i) => i);
 const sense3D = (pos: d.v3f, direction: d.v3f) => {
   'use gpu';
   const dims = std.textureDimensions(computeLayout.$.oldState);
@@ -187,8 +189,7 @@ const sense3D = (pos: d.v3f, direction: d.v3f) => {
   const perp1 = getPerpendicular(direction);
   const perp2 = std.cross(direction, perp1);
 
-  const numSamples = 8;
-  for (let i = 0; i < numSamples; i++) {
+  for (const i of tgpu.unroll(samplesIterations)) {
     const theta = (i / numSamples) * 2 * Math.PI;
 
     const coneOffset = perp1.mul(std.cos(theta)).add(perp2.mul(std.sin(theta)));

apps/typegpu-docs/src/examples/simulation/slime-mold/index.ts

@@ -168,8 +168,8 @@ const blur = tgpu.computeFn({
   let count = d.f32();
 
   // 3x3 blur kernel
-  for (let offsetY = -1; offsetY <= 1; offsetY++) {
-    for (let offsetX = -1; offsetX <= 1; offsetX++) {
+  for (const offsetY of tgpu.unroll([-1, 0, 1])) {
+    for (const offsetX of tgpu.unroll([-1, 0, 1])) {
       const samplePos = d.vec2i(gid.xy).add(d.vec2i(offsetX, offsetY));
       const dimsi = d.vec2i(dims);
 

apps/typegpu-docs/src/examples/simulation/stable-fluid/simulation.ts

@@ -9,7 +9,7 @@ const getNeighbors = tgpu.fn([d.vec2i, d.vec2i], d.arrayOf(d.vec2i, 4))(
       d.vec2i(1, 0),
       d.vec2i(0, 1),
     ];
-    for (let i = 0; i < 4; i++) {
+    for (const i of tgpu.unroll([0, 1, 2, 3])) {
       adjacentOffsets[i] = std.clamp(
         std.add(coords, adjacentOffsets[i]),
         d.vec2i(),

packages/typegpu/src/core/unroll/tgpuUnroll.ts

@@ -0,0 +1,54 @@
+import { stitch } from '../resolve/stitch.ts';
+import {
+  $gpuCallable,
+  $internal,
+  $resolve,
+} from '../../../src/shared/symbols.ts';
+import { setName } from '../../../src/shared/meta.ts';
+import type { DualFn } from '../../../src/types.ts';
+import type { AnyData } from '../../../src/data/dataTypes.ts';
+import {
+  type ResolvedSnippet,
+  snip,
+  type Snippet,
+} from '../../../src/data/snippet.ts';
+import type { ResolutionCtx, SelfResolvable } from '../../../src/types.ts';
+
+/**
+ * The result of calling `tgpu.unroll(...)`. The code responsible for
+ * generating shader code can check if the value of a snippet is
+ * an instance of `UnrollableIterable`, and act accordingly.
+ */
+export class UnrollableIterable implements SelfResolvable {
+  readonly [$internal] = true;
+
+  constructor(public readonly snippet: Snippet) {}
+
+  [$resolve](_ctx: ResolutionCtx): ResolvedSnippet {
+    return snip(
+      stitch`${this.snippet}`,
+      this.snippet.dataType as AnyData,
+      this.snippet.origin,
+    );
+  }
+}
+
+/**
+ * Marks an iterable to be unrolled by the wgslGenerator.
+ */
+export const unroll = (() => {
+  const impl = (<T extends Iterable<unknown>>(value: T) => value) as unknown as
+    & DualFn<(<T extends Iterable<unknown>>(value: T) => T)>
+    & { [$internal]: true };
+
+  setName(impl, 'unroll');
+  impl.toString = () => 'unroll';
+  impl[$internal] = true;
+  impl[$gpuCallable] = {
+    call(_ctx, [value]) {
+      return snip(new UnrollableIterable(value), value.dataType, value.origin);
+    },
+  };
+
+  return impl;
+})();

packages/typegpu/src/tgpu.ts

@@ -14,5 +14,6 @@ export { bindGroupLayout } from './tgpuBindGroupLayout.ts';
 export { computeFn } from './core/function/tgpuComputeFn.ts';
 export { fragmentFn } from './core/function/tgpuFragmentFn.ts';
 export { vertexFn } from './core/function/tgpuVertexFn.ts';
+export { unroll } from './core/unroll/tgpuUnroll.ts';
 
 export * as '~unstable' from './tgpuUnstable.ts';