Functions

Recommended reading

We assume that you are familiar with the following concepts:

• WebGPU Fundamentals
• WebGPU Shading Language

TypeGPU functions let you define shader logic in a modular and type-safe way. Their signatures are fully visible to TypeScript, enabling tooling and static checks. Dependencies, including GPU resources or other functions, are resolved automatically, with no duplication or name clashes. This also supports distributing shader logic across multiple modules or packages. Imported functions from external sources are automatically resolved and embedded into the final shader when referenced.

Defining a function

In order to construct a TypeGPU function, you need to start by defining its shell, an object holding only the input and output types. The shell constructor tgpu.fn relies on TypeGPU schemas, objects that represent WGSL data types and assist in generating shader code at runtime. It accepts two arguments:

• An array of schemas representing argument types,
• (Optionally) a schema representing the return type.

Then the actual WGSL implementation is passed in to a shell invocation using the tagged template literal.

The following code defines a function that accepts one argument of type f32 and returns a vec4f.

import tgpu from 'typegpu';
import * as d from 'typegpu/data';

const getGradientColor = tgpu.fn(
  [d.f32],
  d.vec4f
) /* wgsl */`(ratio: f32) -> vec4f {
  var purple = vec4f(0.769, 0.392, 1.0, 1);
  var blue = vec4f(0.114, 0.447, 0.941, 1);
  return mix(purple, blue, ratio);
}`;

Tip

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

Since type information is already present in the shell, the WGSL header can be simplified to include only the argument names.

const getGradientColor = tgpu.fn([d.f32], d.vec4f) /* wgsl */`(ratio) {
  var purple = vec4f(0.769, 0.392, 1.0, 1);
  var blue = vec4f(0.114, 0.447, 0.941, 1);
  return mix(purple, blue, ratio);
}`;

External resources

Functions can use external resources passed via the $uses method. Externals can include anything that can be resolved to WGSL by TypeGPU (numbers, vectors, matrices, constants, TypeGPU functions, buffer usages, textures, samplers, slots, accessors etc.).

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

// calling a schema to create a value on the JS side
const purple = d.vec4f(0.769, 0.392, 1.0, 1);

const getGradientColor = tgpu.fn([d.f32], d.vec4f)`(ratio) {
  return mix(purple, getBlue(), ratio);;
}
`.$uses({ purple, getBlue: getBlueFunction });

You can check yourself what getGradientColor resolves to by calling tgpu.resolve, all relevant definitions will be automatically included:

// results of calling tgpu.resolve({ externals: { getGradientColor } })
fn getBlueFunction_1() -> vec4f{
  return vec4f(0.114, 0.447, 0.941, 1);
}

fn getGradientColor_0(ratio: f32) -> vec4f{
  return mix(vec4f(0.769, 0.392, 1, 1), getBlueFunction_1(), ratio);;
}

Entry functions

Experimental

Entry functions are an unstable feature. The API may be subject to change in the near future.

Instead of annotating a TgpuFn with attributes, entry functions are defined using dedicated shell constructors:

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

Entry point function I/O

To describe the input and output of an entry point function, we use IORecords, JavaScript objects that map argument names to their types.

const vertexInput = {
  idx: d.builtin.vertexIndex,
  position: d.vec4f,
  color: d.vec4f
}

As you may note, builtin inter-stage inputs and outputs are available on the d.builtin object, and require no further type clarification.

Another thing to note is that there is no need to specify locations of the arguments, as TypeGPU tries to assign locations automatically. If you wish to, you can assign the locations manually with the d.location decorator.

During WGSL generation, TypeGPU automatically generates structs corresponding to the passed IORecords. In WGSL implementation, input and output structs of the given function can be referenced as In and Out respectively. Headers in WGSL implementations must be omitted, all input values are accessible through the struct named in.

Note

Schemas used in d.struct can be wrapped in d.size and d.align decorators, corresponding to @size and @align WGSL attributes.

Since TypeGPU wraps IORecords into automatically generated structs, you can also use those decorators in IOStructs.

Compute

TgpuComputeFn accepts an object with two properties:

• in — an IORecord describing the input of the function,
• workgroupSize — a JS array of 1-3 numbers that corresponds to the @workgroup_size attribute.

const mainCompute = tgpu['~unstable'].computeFn({
  in: { gid: d.builtin.globalInvocationId },
  workgroupSize: [1],
}) /* wgsl */`{
  let index = in.gid.x;
  if index == 0 {
    time += deltaTime;
  }
  let phase = (time / 300) + particleData[index].seed;
  particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
}`.$uses({ particleData: particleDataStorage, deltaTime, time });

Resolved WGSL for the compute function above is equivalent (with respect to some cleanup) to the following:

@group(0) @binding(0) var<storage, read_write> particleData: array<u32, 100>;
@group(0) @binding(1) var<uniform> deltaTime: f32;
@group(0) @binding(2) var<storage, read_write> time: f32;

struct mainCompute_Input {
  @builtin(global_invocation_id) gid: vec3u,
}

@compute @workgroup_size(1) fn mainCompute(in: mainCompute_Input)  {
  let index = in.gid.x;
  if index == 0 {
    time += deltaTime;
  }
  let phase = (time / 300) + particleData[index].seed;
  particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
}

Vertex and fragment

TgpuVertexFn accepts an object with two properties:

• in — an IORecord describing the input of the function,
• out — an IORecord describing the output of the function.

TgpuFragment accepts an object with two properties:

• in — an IORecord describing the input of the function,
• out — d.vec4f, or an IORecord describing the output of the function.

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 });

Resolved WGSL for the pipeline including the two entry point functions above is equivalent (with respect to some cleanup) to the following:

struct mainVertex_Input {
  @builtin(vertex_index) vertexIndex: u32,
}

struct mainVertex_Output {
  @builtin(position) outPos: vec4f,
  @location(0) uv: vec2f,
}

@vertex fn mainVertex(in: mainVertex_Input) -> mainVertex_Output {
  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 mainVertex_Output(vec4f(pos[in.vertexIndex], 0.0, 1.0), uv[in.vertexIndex]);
}

fn getGradientColor(ratio: f32) -> vec4f{
  return mix(vec4f(0.769, 0.392, 1, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
}

struct mainFragment_Input {
  @location(0) uv: vec2f,
}

@fragment fn mainFragment(in: mainFragment_Input) -> @location(0) vec4f {
  return getGradientColor((in.uv[0] + in.uv[1]) / 2);
}

Usage in pipelines

Experimental

Pipelines are an unstable feature. The API may be subject to change in the near future.

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:

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