Unrolling for loops when looping over comptime-known iterables

[iwoplaza]

Idea: if the user does a for of loop on a compile-time known iterable, and wraps it with tgpu.unroll, we could unroll the loop.

Constraints:

• Must be a for of loop, since a counter loop is harder to track (the index could be mutated inside the loop.
• If the loop could be done at runtime, and the unrolled code gets too big, we warn the user in dev-mode.

Unrolling an iterated array

const foo = () => {
  'use gpu';
  let result = d.f32(0);
  for (const item of tgpu.unroll([1, 2, 3])) {
    result += item;
  }
  return result;
};

Generated WGSL:

fn foo() -> f32 {
  var result = 0f;

  { result += 1; }
  { result += 2; }
  { result += 3; }

  return result;
}

Conditional unroll

const collection = [1, 1, 2, 3, 5, 8];
const shouldUnroll = true;

const foo = () => {
  'use gpu';
  for (const x of shouldUnroll ? tgpu.unroll(collection) : collection) {
    // ...
  }
};

Iterating over struct fields

const Weights = d.struct(Object.fromEntries(
  Object.keys(variants).map((name) => [`weight_${name}`, d.f32]),
));

const compute = (weights: d.Infer<typeof Weights>) => {
  'use gpu';

  let p = d.f32(0);
  for (const key of tgpu.unroll(Object.keys(variants))) {
    p += weights[key] * variants[key](p);
  }
  return p;
};

Generated WGSL:

struct Weights {
  weight_foo: f32,
  weight_bar: f32,
  weight_baz: f32,
}

// ...
// variant functions
// ...

fn compute(weights: Weights) -> f32 {
  var p = 0f;
  
  { p += weights.weight_foo * compute_foo(p); }
  { p += weights.weight_bar * compute_bar(p); }
  { p += weights.weight_baz * compute_baz(p); }
  
  return p;
}

[deluksic]

This sounds like an awesome change! I think I would prefer a build failure if you mark something unroll, but it doesn't due to code size. Otherwise we risk it feeling like inline (just a suggestion) from C.

[iwoplaza]

We're workshopping a way to mark unrolling explicitly, where you'd at most get a warning in dev-mode if the generated code gets too big, but would get unrolled anyway. The other way would work as well, where we would never unroll a loop for you unless you explicitly mark it for unrolling.

[aleksanderkatan]

I like this idea, if somebody wants to unroll a regular for, they can just transform it into a for ... of ....

What happens when somebody passes in an array but does not unroll it? Will it get copied in every spot it is accessed in, or will it be assigned to a variable?

[iwoplaza]

Good question. Let's consider the shown example without tgpu.unroll(...). I think we should error out and request the dev to either wrap it in tgpu.unroll, or store it in a const / variable.

const foo = () => {
  'use gpu';
  let result = d.f32(0);
  const arr = [1, 2, 3];
  for (const item of arr) {
    result += d.f32(item);
  }
  return result;
};

Generated WGSL:

fn foo() -> f32 {
  var result = 0f;
  var arr = array<i32, 3>(1, 2, 3);
  for (var i = 0u; i < 3; i++) {
    result += f32(arr[i]);
  }
  return result;
}

[iwoplaza]

We're doing it #1965