Utilities

tgpu.rawCodeSnippet

When working on top of some existing shader code, sometimes you may know for certain that some variable will be already defined and should be accessible in the code. In such scenario you can use tgpu['~unstable'].rawCodeSnippet — an advanced API that creates a typed shader expression which can be injected into the final shader bundle upon use.

// `EXISTING_GLOBAL` is an identifier that we know will be in the
// final shader bundle, but we cannot
// refer to it in any other way.
const existingGlobal = tgpu['~unstable']
  .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');

const foo = tgpu.fn([], d.f32)(() => {
  'use gpu';
  return existingGlobal.$ * 22;
});

const wgsl = tgpu.resolve([foo]);
// fn foo() -> f32 {
//   return (EXISTING_GLOBAL * 22f);
// }

Note

There currently is no way to create a rawCodeSnippet that refers to an existing function.

Which origin to choose?

The optional third parameter origin lets TypeGPU transpiler know how to optimize the code snippet, as well as allows for some transpilation-time validity checks.

Usually 'runtime' (the default) is a safe bet, but if you’re sure that the expression or computation is constant (either a reference to a constant, a numeric literal, or an operation on constants), then pass 'constant' as it might lead to better optimizations.

If what the expression is a direct reference to an existing value (e.g. a uniform, a storage binding, …), then choose from 'uniform', 'mutable', 'readonly', 'workgroup', 'private' or 'handle' depending on the address space of the referred value.

Tip

TgpuGuardedComputePipeline provides getters for the underlying pipeline and the size buffer. Those might be useful for tgpu.resolve, since you cannot resolve a guarded pipeline directly.

const innerPipeline = doubleUpPipeline.with(bindGroup1).pipeline;
tgpu.resolve([innerPipeline]);

tgpu.comptime

tgpu.comptime(func) creates a version of func that instead of being transpiled to WGSL, will be called during the WGSL code generation. This can be used to precompute and inject a value into the final shader code.

const color = tgpu.comptime((int: number) => {
  const r = (int >> 16) & 0xff;
  const g = (int >> 8) & 0xff;
  const b = int & 0xff;
  return d.vec3f(r / 255, g / 255, b / 255);
});

const material = tgpu.fn([d.vec3f], d.vec3f)((diffuse) => {
  'use gpu';
  const albedo = color(0xff00ff);
  return albedo.mul(diffuse);
});

const shader = tgpu.resolve([material]);
// fn material(diffuse: vec3f) -> vec3f {
//   var albedo = vec3f(1, 0, 1);
//   return (albedo * diffuse);
// }

Note how the function passed into comptime doesn’t have to be marked with 'use gpu' and can use Math. That’s because the function doesn’t execute on the GPU, it gets executed before the shader code gets sent to the GPU.

Branch pruning

If a condition is known at resolution time (comptime), then typegpu prunes the unvisited block. Comptime-known conditions include:

• referenced js values and operations on such, like std.pow(userSelection, 2) < THRESHOLD (note that these values will be concretized during shader resolution, if userSelection may change over time, use buffers to provide it),
• values provided via slots,
• values returned by comptime functions.

const counterEnabledSlot = tgpu.slot<boolean>(false);
const counter = root.createMutable(d.u32);

const myFunction = tgpu.fn([])(() => {
  if (counterEnabledSlot.$) counter.$++;
});
//  fn myFunction() {
//
//  }

const myFunctionWithCounter = myFunction.with(counterEnabledSlot, true);
//  "@group(0) @binding(0) var<storage, read_write> counter: u32;
//
//  fn myFunction() {
//    {
//      counter++;
//    }
//  }"

Note

Conditions involving tgpu.const won’t be automatically pruned, since we treat tgpu.const as a way to opt-out of inlining.

Branch pruning also works for ternary operators.

const myFunction = tgpu.fn([])(() => {
  counterEnabledSlot.$ ? counter.$++ : undefined;
});

Caution

Ternary operator’s condition must be a comptime-known value. This restriction is due to WGSL having no ternary operator equivalent.

tgpu.rawCodeSnippet

When working on top of some existing shader code, sometimes you may know for certain that some variable will be already defined and should be accessible in the code. In such scenario you can use tgpu['~unstable'].rawCodeSnippet — an advanced API that creates a typed shader expression which can be injected into the final shader bundle upon use.

// `EXISTING_GLOBAL` is an identifier that we know will be in the
// final shader bundle, but we cannot
// refer to it in any other way.
const existingGlobal = tgpu['~unstable']
  .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');

const foo = tgpu.fn([], d.f32)(() => {
  'use gpu';
  return existingGlobal.$ * 22;
});

const wgsl = tgpu.resolve([foo]);
// fn foo() -> f32 {
//   return (EXISTING_GLOBAL * 22f);
// }

Note

There currently is no way to create a rawCodeSnippet that refers to an existing function.

Which origin to choose?

The optional third parameter origin lets TypeGPU transpiler know how to optimize the code snippet, as well as allows for some transpilation-time validity checks.

Usually 'runtime' (the default) is a safe bet, but if you’re sure that the expression or computation is constant (either a reference to a constant, a numeric literal, or an operation on constants), then pass 'constant' as it might lead to better optimizations.

If what the expression is a direct reference to an existing value (e.g. a uniform, a storage binding, …), then choose from 'uniform', 'mutable', 'readonly', 'workgroup', 'private' or 'handle' depending on the address space of the referred value.

console.log

Yes, you read that correctly, TypeGPU implements logging to the console on the GPU! Just call console.log like you would in plain JavaScript, and open the console to see the results.

const callCountMutable = root.createMutable(d.u32, 0);
const compute = root.createGuardedComputePipeline(() => {
  'use gpu';
  callCountMutable.$ += 1;
  console.log('Call number', callCountMutable.$);
});

compute.dispatchThreads();
compute.dispatchThreads();

// Eventually...
// "[GPU] Call number 1"
// "[GPU] Call number 2"

Currently supported data types for logging include scalars, vectors, matrices, structs, and fixed-size arrays.

Under the hood, TypeGPU translates console.log to a series of serializing functions that write the logged arguments to a buffer that is read and deserialized after every draw/dispatch call.

The buffer is of fixed size, which may limit the total amount of information that can be logged; if the buffer overflows, additional logs are dropped. If that’s an issue, you may specify the size manually when creating the root object.

const root = await tgpu.init({
  unstable_logOptions: {
    logCountLimit: 32,
    logSizeLimit: 8, // in bytes, enough to fit 2*u32
  },
});

/* vertex shader */

const mainFragment = tgpu.fragmentFn({
  in: { pos: d.builtin.position },
  out: d.vec4f,
})(({ pos }) => {
  // this log fits in 8 bytes
  // static strings do not count towards the serialized log size
  console.log('X:', d.u32(pos.x), 'Y:', d.u32(pos.y));
  return d.vec4f(0, 1, 1, 1);
});

/* pipeline creation and draw call */

Note

The logs are written to console only after the dispatch finishes and the buffer is read. This may happen with a noticeable delay.

Caution

When using console.log, atomic operations are injected into the WGSL code to safely synchronize logging from multiple threads. This synchronization can introduce overhead and significantly impact shader performance.

Other supported console functionalities include console.debug, console.info, console.warn, console.error and console.clear.

There are some limitations (some of which we intend to alleviate in the future):

• console.log only works when used in TypeGPU functions that are transitively called in a TypeGPU pipeline. Otherwise, for example when using tgpu.resolve on a WGSL template, logs are ignored.
• console.log only works in fragment and compute shaders. This is due to a WebGPU limitation 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, 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.

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():

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.