Skip to content

Commit bc070f0

Browse files
iwoplazaiwoplaza
authored
docs: Stable tgpu.fn guide (#1409)
1 parent 93a3a99 commit bc070f0

File tree

1 file changed

+234
-46
lines changed
  • apps/typegpu-docs/src/content/docs/fundamentals/functions

1 file changed

+234
-46
lines changed

‎apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx

Lines changed: 234 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -3,84 +3,214 @@ title: Functions
33
description: A guide on how to create and use the TypeGPU typed functions.
44
---
55

6-
:::caution[Experimental]
7-
Functions are an *unstable* feature. The API may be subject to change in the near future.
8-
:::
9-
106
:::note[Recommended reading]
117
We assume that you are familiar with the following concepts:
128
- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-fundamentals.html" target="_blank" rel="noopener noreferrer">WebGPU Fundamentals</a>
139
- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-wgsl.html" target="_blank" rel="noopener noreferrer">WebGPU Shading Language</a>
1410
:::
1511

16-
TypeGPU allows writing shaders by composing typed functions, which are special wrappers around WGSL code.
17-
These functions can reference outside resources, like other user-defined or helper functions, buffers, bind group layouts etc.
12+
TypeGPU functions let you define shader logic in a modular and type-safe way.
13+
Their signatures are fully visible to TypeScript, enabling tooling and static checks.
14+
Dependencies, including GPU resources or other functions, are resolved automatically, with no duplication or name clashes.
15+
This also supports distributing shader logic across multiple modules or packages.
16+
Imported functions from external sources are automatically resolved and embedded into the final shader when referenced.
1817

19-
## Creating a function
18+
## Defining a function
2019

21-
Functions are constructed by first defining their shells, which specify their inputs and outputs.
22-
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.
20+
In order to construct a TypeGPU function, you need to start by defining its shell, an object holding only the input and output types.
21+
The shell constructor `tgpu.fn` relies on [TypeGPU schemas](/TypeGPU/fundamentals/data-schemas), objects that represent WGSL data types and assist in generating shader code at runtime.
22+
It accepts two arguments:
2323

24-
The following code defines a function that accepts one argument and returns one value.
24+
- An array of schemas representing argument types,
25+
- (Optionally) a schema representing the return type.
2526

26-
```ts
27-
const getGradientColor = tgpu.fn([d.f32], d.vec4f)(/* wgsl */ `(ratio: f32) -> vec4f {
28-
let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
29-
return color;
30-
}`);
27+
Then the actual WGSL implementation is passed in to a shell invocation using the tagged template literal.
28+
29+
The following code defines a function that accepts one argument of type `f32` and returns a `vec4f`.
3130

32-
// or
31+
```ts twoslash
32+
import tgpu from 'typegpu';
33+
import * as d from 'typegpu/data';
3334

34-
const getGradientColor = tgpu.fn([d.f32], d.vec4f) /* wgsl */`(ratio: f32) -> vec4f {
35-
let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
36-
return color;
37-
};
35+
const getGradientColor = tgpu.fn(
36+
[d.f32],
37+
d.vec4f
38+
) /* wgsl */`(ratio: f32) -> vec4f {
39+
var purple = vec4f(0.769, 0.392, 1.0, 1);
40+
var blue = vec4f(0.114, 0.447, 0.941, 1);
41+
return mix(purple, blue, ratio);
42+
}`;
3843
```
3944

40-
If you're using Visual Studio Code, you can use an [extension](https://marketplace.visualstudio.com/items?itemName=ggsimm.wgsl-literal) that brings syntax highlighting to the code fragments marked with `/* wgsl */` comments.
45+
:::tip
46+
If you're using Visual Studio Code, you can use [this extension](https://marketplace.visualstudio.com/items?itemName=ggsimm.wgsl-literal) that brings syntax highlighting to the code fragments marked with `/* wgsl */` comments.
47+
:::
48+
49+
Since type information is already present in the shell, the WGSL header can be simplified to include only the argument names.
50+
51+
```ts twoslash
52+
import tgpu from 'typegpu';
53+
import * as d from 'typegpu/data';
54+
55+
// ---cut---
56+
const getGradientColor = tgpu.fn([d.f32], d.vec4f) /* wgsl */`(ratio) {
57+
var purple = vec4f(0.769, 0.392, 1.0, 1);
58+
var blue = vec4f(0.114, 0.447, 0.941, 1);
59+
return mix(purple, blue, ratio);
60+
}`;
61+
```
4162

4263
## External resources
4364

44-
Functions can use external resources passed inside a record via the `$uses` method.
45-
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.).
65+
Functions can use external resources passed via the `$uses` method.
66+
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.).
4667

47-
```ts
48-
const getBlue = tgpu.fn([], d.vec4f)`() -> vec4f {
68+
```ts twoslash
69+
import tgpu from 'typegpu';
70+
import * as d from 'typegpu/data';
71+
72+
// ---cut---
73+
const getBlueFunction = tgpu.fn([], d.vec4f)`() {
4974
return vec4f(0.114, 0.447, 0.941, 1);
5075
}`;
5176

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

54-
const getGradientColor = tgpu.fn([d.f32], d.vec4f)`(ratio: f32) -> vec4f {
55-
let color = mix(purple, getBlue(), ratio);
56-
return color;
80+
const getGradientColor = tgpu.fn([d.f32], d.vec4f)`(ratio) {
81+
return mix(purple, getBlue(), ratio);;
5782
}
58-
`.$uses({ purple, getBlue });
83+
`.$uses({ purple, getBlue: getBlueFunction });
5984
```
6085

61-
The `getGradientColor` function, when resolved to WGSL, includes the definitions of all used external resources:
86+
You can check yourself what `getGradientColor` resolves to by calling [`tgpu.resolve`](/TypeGPU/fundamentals/resolve), all relevant definitions will be automatically included:
6287

6388
```wgsl
64-
fn getBlue_1() -> vec4f {
89+
// results of calling tgpu.resolve({ externals: { getGradientColor } })
90+
fn getBlueFunction_1() -> vec4f{
6591
return vec4f(0.114, 0.447, 0.941, 1);
6692
}
6793
68-
fn getGradientColor_0(ratio: f32) -> vec4f {
69-
let color = mix(vec4f(0.769, 0.392, 1, 1), getBlue_1(), ratio);
70-
return color;
94+
fn getGradientColor_0(ratio: f32) -> vec4f{
95+
return mix(vec4f(0.769, 0.392, 1, 1), getBlueFunction_1(), ratio);;
7196
}
7297
```
7398

7499
## Entry functions
75100

76-
Defining entry functions is similar to regular ones, but is done through dedicated constructors:
77-
- `tgpu['~unstable'].vertexFn`
78-
- `tgpu['~unstable'].fragmentFn`
79-
- `tgpu['~unstable'].computeFn`
101+
:::caution[Experimental]
102+
Entry functions are an *unstable* feature. The API may be subject to change in the near future.
103+
:::
104+
105+
Instead of annotating a `TgpuFn` with attributes, entry functions are defined using dedicated shell constructors:
106+
107+
- `tgpu['~unstable'].computeFn`,
108+
- `tgpu['~unstable'].vertexFn`,
109+
- `tgpu['~unstable'].fragmentFn`.
110+
111+
### Entry point function I/O
80112

81-
They can be passed to root-defined pipelines and they accept special arguments like builtins (`d.builtin`) and decorated data (`d.location`).
113+
To describe the input and output of an entry point function, we use `IORecord`s, JavaScript objects that map argument names to their types.
82114

83115
```ts
116+
const vertexInput = {
117+
idx: d.builtin.vertexIndex,
118+
position: d.vec4f,
119+
color: d.vec4f
120+
}
121+
```
122+
123+
As you may note, builtin inter-stage inputs and outputs are available on the `d.builtin` object,
124+
and require no further type clarification.
125+
126+
Another thing to note is that there is no need to specify locations of the arguments,
127+
as TypeGPU tries to assign locations automatically.
128+
If you wish to, you can assign the locations manually with the `d.location` decorator.
129+
130+
During WGSL generation, TypeGPU automatically generates structs corresponding to the passed `IORecord`s.
131+
In WGSL implementation, input and output structs of the given function can be referenced as `In` and `Out` respectively.
132+
Headers in WGSL implementations must be omitted, all input values are accessible through the struct named `in`.
133+
134+
:::note
135+
Schemas used in `d.struct` can be wrapped in `d.size` and `d.align` decorators,
136+
corresponding to `@size` and `@align` WGSL attributes.
137+
138+
Since TypeGPU wraps `IORecord`s into automatically generated structs, you can also use those decorators in `IOStruct`s.
139+
:::
140+
141+
### Compute
142+
143+
`TgpuComputeFn` accepts an object with two properties:
144+
145+
- `in` -- an `IORecord` describing the input of the function,
146+
- `workgroupSize` -- a JS array of 1-3 numbers that corresponds to the `@workgroup_size` attribute.
147+
148+
```ts twoslash
149+
import tgpu from 'typegpu';
150+
import * as d from 'typegpu/data';
151+
152+
const root = await tgpu.init();
153+
154+
const particleDataBuffer = root
155+
.createBuffer(d.arrayOf(d.u32, 100))
156+
.$usage('storage', 'uniform', 'vertex');
157+
158+
const deltaTime = root.createUniform(d.f32);
159+
const time = root.createMutable(d.f32);
160+
const particleDataStorage = particleDataBuffer.as('mutable');
161+
// ---cut---
162+
const mainCompute = tgpu['~unstable'].computeFn({
163+
in: { gid: d.builtin.globalInvocationId },
164+
workgroupSize: [1],
165+
}) /* wgsl */`{
166+
let index = in.gid.x;
167+
if index == 0 {
168+
time += deltaTime;
169+
}
170+
let phase = (time / 300) + particleData[index].seed;
171+
particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
172+
}`.$uses({ particleData: particleDataStorage, deltaTime, time });
173+
```
174+
175+
Resolved WGSL for the compute function above is equivalent (with respect to some cleanup) to the following:
176+
177+
```wgsl
178+
@group(0) @binding(0) var<storage, read_write> particleData: array<u32, 100>;
179+
@group(0) @binding(1) var<uniform> deltaTime: f32;
180+
@group(0) @binding(2) var<storage, read_write> time: f32;
181+
182+
struct mainCompute_Input {
183+
@builtin(global_invocation_id) gid: vec3u,
184+
}
185+
186+
@compute @workgroup_size(1) fn mainCompute(in: mainCompute_Input) {
187+
let index = in.gid.x;
188+
if index == 0 {
189+
time += deltaTime;
190+
}
191+
let phase = (time / 300) + particleData[index].seed;
192+
particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
193+
}
194+
```
195+
196+
### Vertex and fragment
197+
198+
`TgpuVertexFn` accepts an object with two properties:
199+
200+
- `in` -- an `IORecord` describing the input of the function,
201+
- `out` -- an `IORecord` describing the output of the function.
202+
203+
`TgpuFragment` accepts an object with two properties:
204+
205+
- `in` -- an `IORecord` describing the input of the function,
206+
- `out` -- `d.vec4f`, or an `IORecord` describing the output of the function.
207+
208+
```ts twoslash
209+
import tgpu from 'typegpu';
210+
import * as d from 'typegpu/data';
211+
212+
const getGradientColor = tgpu.fn([d.f32], d.vec4f)``;
213+
// ---cut---
84214
const mainVertex = tgpu['~unstable'].vertexFn({
85215
in: { vertexIndex: d.builtin.vertexIndex },
86216
out: { outPos: d.builtin.position, uv: d.vec2f },
@@ -108,15 +238,75 @@ const mainFragment = tgpu['~unstable'].fragmentFn({
108238
}`.$uses({ getGradientColor });
109239
```
110240

111-
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.
112-
TypeGPU schemas for these structs are created automatically by the library and their definitions are included when resolving the functions.
113-
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.
241+
Resolved WGSL for the pipeline including the two entry point functions above is equivalent (with respect to some cleanup) to the following:
242+
243+
```wgsl
244+
struct mainVertex_Input {
245+
@builtin(vertex_index) vertexIndex: u32,
246+
}
247+
248+
struct mainVertex_Output {
249+
@builtin(position) outPos: vec4f,
250+
@location(0) uv: vec2f,
251+
}
252+
253+
@vertex fn mainVertex(in: mainVertex_Input) -> mainVertex_Output {
254+
var pos = array<vec2f, 3>(
255+
vec2(0.0, 0.5),
256+
vec2(-0.5, -0.5),
257+
vec2(0.5, -0.5)
258+
);
259+
260+
var uv = array<vec2f, 3>(
261+
vec2(0.5, 1.0),
262+
vec2(0.0, 0.0),
263+
vec2(1.0, 0.0),
264+
);
265+
266+
return mainVertex_Output(vec4f(pos[in.vertexIndex], 0.0, 1.0), uv[in.vertexIndex]);
267+
}
268+
269+
fn getGradientColor(ratio: f32) -> vec4f{
270+
return mix(vec4f(0.769, 0.392, 1, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
271+
}
272+
273+
struct mainFragment_Input {
274+
@location(0) uv: vec2f,
275+
}
276+
277+
@fragment fn mainFragment(in: mainFragment_Input) -> @location(0) vec4f {
278+
return getGradientColor((in.uv[0] + in.uv[1]) / 2);
279+
}
280+
```
114281

115282
## Usage in pipelines
116283

284+
:::caution[Experimental]
285+
Pipelines are an *unstable* feature. The API may be subject to change in the near future.
286+
:::
287+
117288
Typed functions are crucial for simplified pipeline creation offered by TypeGPU. You can define and run pipelines as follows:
118289

119-
```ts
290+
```ts twoslash
291+
import tgpu from 'typegpu';
292+
import * as d from 'typegpu/data';
293+
294+
const context = undefined as any;
295+
const presentationFormat = "rgba8unorm";
296+
const root = await tgpu.init();
297+
298+
const getGradientColor = tgpu.fn([d.f32], d.vec4f)/* wgsl */``;
299+
300+
const mainVertex = tgpu['~unstable'].vertexFn({
301+
in: { vertexIndex: d.builtin.vertexIndex },
302+
out: { outPos: d.builtin.position, uv: d.vec2f },
303+
})``;
304+
305+
const mainFragment = tgpu['~unstable'].fragmentFn({
306+
in: { uv: d.vec2f },
307+
out: d.vec4f,
308+
})``;
309+
// ---cut---
120310
const pipeline = root['~unstable']
121311
.withVertex(mainVertex, {})
122312
.withFragment(mainFragment, { format: presentationFormat })
@@ -136,5 +326,3 @@ The rendering result looks like this:
136326
![rendering result - gradient triangle](./triangle-result.png)
137327

138328
You can check out the full example on [our examples page](/TypeGPU/examples#example=simple--triangle).
139-
140-

0 commit comments

Comments
 (0)