[Proposal] tgpu.macro

[iwoplaza]

Note

This proposal depends on the "std.isKnownAtComptime" proposal #3027

Calling shader functions generates a call in WGSL, calling a tgpu.comptime functions requires all arguments to be known at comptime, pregenerates the result and injects it into the WGSL. We currently have no way of extracting mixed comptime/runtime behavior out of a shader function in a reusable way. tgpu.macro would be an API to allow for that.

tgpu.macro would accept a 'use gpu' callback that needs to meet a specific requirement:

After all comptime pruning, there needs to be exactly one return statement.

This lets us essentially inline the expression into the callee, and act as if we wrote it there. The arguments passed into the macro need to be side-effect free, as the arguments could be used multiple times inside of the macro (e.g. tgpu.macro((a) => { 'use gpu'; return a * a; }))

Unlocked capabilities

• Using macros & std.isKnownAtComptime, we can create functions that precompute if all arguments are known at comptime, and generate a function call otherwise.
• Cleaner utilities

Overly simplified example

const double = tgpu.macro((a: number) => {
  'use gpu';
  return a * 2;
});

function foo() {
  'use gpu';
  return double(d.f32(4));
}
// fn foo() -> f32 {
//   return 8f;
// }

const time = root.createUniform(d.f32);

function bar() {
  'use gpu';
  return double(time.$);
}
// var<uniform> time: f32;
// 
// fn bar() -> f32 {
//   return time * 2f;
// }

tryUnroll utility

/**
 * Unrolls the iterable if its length is known, falls back to a rolled loop otherwise.
 */
const tryUnroll = tgpu.macro(<T extends unknown[]>(iterable: T): T => {
  'use gpu';
  return std.isKnownAtComptime(iterable.length) ? tgpu.unroll(iterable) : iterable;
});

function main() {
  'use gpu';
  for (const boid of tryUnroll(boids.$)) {
    // ...
  }
}

withTryComptime utility

Using macros, we can auto-magically run a function at comp-time if the argument is known at comptime, and at runtime otherwise

function _hexToRgb(hex: number) {
  'use gpu';
  // ...
}

const hexToRgb = tgpu.macro((hex: number) => {
  'use gpu';
  if (std.isKnownAtComptime(hex)) {
    return tgpu.comptime(_hexToRgb)(hex);
  }
  return _hexToRgb(hex);
});

When we support the ...args rest parameter, we can turn this into a general utility:

const withTryComptime = (fn) => {
  return tgpu.macro((...args) => {
    'use gpu';
    if (std.isKnownAtComptime(args)) {
      return tgpu.comptime(fn)(...args);
    }
    return fn(...args);
  });
}

const hexToRgb = withTryComptime(() => {
  'use gpu';
  // ...
});