Functions

TypeGPU allows writing shaders by composing typed functions, which are special wrappers around WGSL code. These functions can reference outside resources, like other user-defined or helper functions, buffers, bind group layouts etc.

Creating a function

Functions are constructed by first defining their shells, which specify their inputs and outputs. Then the actual WGSL implementation is passed in as an argument to a shell invocation. If the code string is a template literal, you can omit the parentheses, which may result in a more compact Biome/Prettier formatting.

The following code defines a function that accepts one argument and returns one value.

const getGradientColor = tgpu['~unstable']
  .fn(
    { ratio: d.f32 },
    d.vec4f,
  )(`{
    let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
    return color;
  }`)
  .$name('getGradientColor');

// or

const getGradientColor = tgpu['~unstable'].fn({ ratio: d.f32 }, d.vec4f)`{
  let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
  return color;
}`.$name('getGradientColor');

If you’re using Visual Studio Code, you can use an extension that brings syntax highlighting to the code fragments marked with /* wgsl */ comments.

External resources

Functions can use external resources passed inside a record via the $uses method. Externals can be any value or TypeGPU resource that can be resolved to WGSL (functions, buffer usages, slots, accessors, constants, variables, declarations, vectors, matrices, textures, samplers etc.).

const getBlue = tgpu['~unstable'].fn({}, d.vec4f)`{
  return vec4f(0.114, 0.447, 0.941, 1);
}`;

const purple = d.vec4f(0.769, 0.392, 1.0, 1);

const getGradientColor = tgpu['~unstable'].fn({ ratio: d.f32 }, d.vec4f)`{
  let color = mix(purple, getBlue(), ratio);
  return color;
}`
  .$uses({ purple, getBlue })
  .$name('getGradientColor');

The getGradientColor function, when resolved to WGSL, includes the definitions of all used external resources:

fn getBlue_1() -> vec4f {
  return vec4f(0.114, 0.447, 0.941, 1);
}

fn getGradientColor_0(ratio: f32) -> vec4f {
  let color = mix(vec4f(0.769, 0.392, 1, 1), getBlue_1(), ratio);
  return color;
}

Entry functions

Defining entry functions is similar to regular ones, but is done through dedicated constructors:

tgpu['~unstable'].vertexFn
tgpu['~unstable'].fragmentFn
tgpu['~unstable'].computeFn

They can be passed to root-defined pipelines and they accept special arguments like builtins (d.builtin) and decorated data (d.location).

const mainVertex = tgpu['~unstable'].vertexFn({
  in: { vertexIndex: d.builtin.vertexIndex },
  out: { outPos: d.builtin.position, uv: d.vec2f },
}) /* wgsl */`{
    var pos = array<vec2f, 3>(
      vec2(0.0, 0.5),
      vec2(-0.5, -0.5),
      vec2(0.5, -0.5)
    );

    var uv = array<vec2f, 3>(
      vec2(0.5, 1.0),
      vec2(0.0, 0.0),
      vec2(1.0, 0.0),
    );

    return Out(vec4f(pos[in.vertexIndex], 0.0, 1.0), uv[in.vertexIndex]);
  }`;

const mainFragment = tgpu['~unstable'].fragmentFn({
  in: { uv: d.vec2f },
  out: d.vec4f,
}) /* wgsl */`{
    return getGradientColor((in.uv[0] + in.uv[1]) / 2);
  }`.$uses({ getGradientColor });

When entry function inputs or outputs are specified as objects containing builtins and inter-stage variables, the WGSL implementations need to access these arguments as passed in via structs. TypeGPU schemas for these structs are created automatically by the library and their definitions are included when resolving the functions. Input values are accessible through the in keyword, while the automatically created structs for input and output shall be referenced in implementation as In and Out respectively.

Usage in pipelines

Typed functions are crucial for simplified pipeline creation offered by TypeGPU. You can define and run pipelines as follows:

const pipeline = root['~unstable']
  .withVertex(mainVertex, {})
  .withFragment(mainFragment, { format: presentationFormat })
  .createPipeline();

pipeline
  .withColorAttachment({
    view: context.getCurrentTexture().createView(),
    clearValue: [0, 0, 0, 0],
    loadOp: 'clear',
    storeOp: 'store',
  })
  .draw(3);

The rendering result looks like this: rendering result - gradient triangle

You can check out the full example on our examples page.