docs: TypeGPU functions

apps/typegpu-docs/astro.config.mjs

@@ -132,7 +132,7 @@ export default defineConfig({
               slug: 'fundamentals/functions',
             },
             {
-              label: 'TGSL',
+              label: 'JS Kernels',
               slug: 'fundamentals/tgsl',
               badge: { text: 'new' },
             },

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

@@ -3,20 +3,186 @@ title: Functions
 description: A guide on how to create and use the TypeGPU typed functions.
 ---
 
-:::note[Recommended reading]
-We assume that you are familiar with the following concepts:
-- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-fundamentals.html" target="_blank" rel="noopener noreferrer">WebGPU Fundamentals</a>
-- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-wgsl.html" target="_blank" rel="noopener noreferrer">WebGPU Shading Language</a>
+:::caution[May require unplugin-typegpu]
+To write TypeGPU functions in JavaScript/TypeScript, you need to install and configure [unplugin-typegpu](/TypeGPU/tooling/unplugin-typegpu).
+If you're planning on only using WGSL, you can skip this step.
 :::
 
-TypeGPU functions let you define shader logic in a modular and type-safe way.
+**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
 
+:::note[WGSL enthusiasts!]
+Don't let the JavaScript scare you! TypeGPU functions can be implemented using either WGSL or JS, both being able to call one another.
+If you're planning on only using WGSL, you can skip right over to [Implementing functions in WGSL](#implementing-functions-in-wgsl),
+though we recommend reading through anyway.
+:::
+
+The simplest and most powerful way to define TypeGPU functions is to just place `'use gpu'` at the beginning of the function body.
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+// ---cut---
+const neighborhood = (a: number, r: number) => {
+  'use gpu';
+  return d.vec2f(a - r, a + r);
+};
+```
+
+The `'use gpu'` directive allows the function to be picked up by our dedicated build plugin -- [unplugin-typegpu](/TypeGPU/tooling/unplugin-typegpu)
+and transformed into a format TypeGPU can understand. This doesn't alter the fact that the function is still callable from JavaScript, and behaves
+the same on the CPU and GPU.
+
+There are three main ways to use TypeGPU functions.
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+const root = await tgpu.init();
+
+const neighborhood = (a: number, r: number) => {
+  'use gpu';
+  return d.vec2f(a - r, a + r);
+};
+
+// ---cut---
+const main = () => {
+  'use gpu';
+  // Call from another function
+  return neighborhood(1.1, 0.5);
+};
+
+// #1) Can be called in JS
+const range = main();
+//    ^?
+
+// #2) Used to generate WGSL
+const wgsl = tgpu.resolve({ externals: { main } });
+//    ^?
+
+// #3) Executed on the GPU (generates WGSL underneath)
+root['~unstable']
+  .prepareDispatch(main)
+  .dispatchThreads();
+````
+
+The contents of the `wgsl` variable would contain the following:
+
+```wgsl
+// Generated WGSL
+fn neighborhood(a: f32, r: f32) -> vec2f {
+  return vec2f(a - r, a + r);
+}
+
+fn main() -> vec2f {
+  return neighborhood(1.1, 0.5);
+}
+
+// ...
+````
+
+You can already notice a few things about TypeGPU functions:
+- Using operators like `+`, `-`, `*`, `/`, etc. is perfectly valid on numbers.
+- TS types are properly inferred, feel free to hover over the variables to see their types.
+- The generated code closely matches your source code.
+
+## Type inference
+
+Let's take a closer look at `neighborhood` versus the WGSL it generates.
+
+```ts
+// TS
+const neighborhood = (a: number, r: number) => {
+  'use gpu';
+  return d.vec2f(a - r, a + r);
+};
+```
+```wgsl
+// WGSL
+fn neighborhood(a: f32, r: f32) -> vec2f {
+  return vec2f(a - r, a + r);
+}
+```
+
+How does TypeGPU determine that `a` and `r` are of type `f32`, and that the return type is `vec2f`? You might think that we parse the TypeScript source file and use the types
+that the user provided in the function signature, **but that's not the case**.
+
+While generating WGSL, TypeGPU infers the type of each expression, which means it knows the types of values passed in at each call site.
+
+```ts twoslash "1.1, 0.5"
+import * as d from 'typegpu/data';
+const neighborhood = (a: number, r: number) => {
+  'use gpu';
+  return d.vec2f(a - r, a + r);
+};
+// ---cut---
+const main = () => {
+  'use gpu';
+  // A very easy case, just floating point literals, so f32 by default
+  return neighborhood(1.1, 0.5);
+};
+```
+
+TypeGPU then propagates those types into the function body and analyses the types returned by the function.
+If it cannot unify them into a single type, it will throw an error.
+
+### Polymorphism
+
+For each set of input types, TypeGPU generates a specialized version of the function.
+
+```ts twoslash
+import * as d from 'typegpu/data';
+const neighborhood = (a: number, r: number) => {
+  'use gpu';
+  return d.vec2f(a - r, a + r);
+};
+// ---cut---
+const main = () => {
+  'use gpu';
+  const a = neighborhood(0, 1);
+  // We can also use casts to coerce values into a specific type.
+  const b = neighborhood(d.u32(1), d.f16(5.25));
+};
+```
+
+```wgsl
+// WGSL
+fn neighborhood(a: i32, r: i32) -> vec2f {
+  return vec2f(f32(a - r), f32(a + r));
+}
+
+fn neighborhood2(a: u32, r: f16) -> vec2f {
+  return vec2f(f32(f16(a) - r), f32(f16(a) + r));
+}
+
+fn main() {
+  var a = neighborhood(0, 1);
+  var b = neighborhood2(1, 5.25);
+}
+```
+
+You can limit the types that a function can accept by using wrapping it in a shell.
+
+### Generics
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+// ---cut---
+const double = <T extends d.v2f | d.v3f | d.v4f>(a: T): T => {
+  'use gpu';
+  return std.mul(a, a);
+};
+```
+
+## Bruh
+
 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](/TypeGPU/fundamentals/data-schemas), objects that represent WGSL data types and assist in generating shader code at runtime.
 It accepts two arguments:
@@ -33,7 +199,7 @@ import tgpu from 'typegpu';
 import * as d from 'typegpu/data';
 
 const getGradientColor = tgpu.fn(
-  [d.f32], 
+  [d.f32],
   d.vec4f
 ) /* wgsl */`(ratio: f32) -> vec4f {
   var purple = vec4f(0.769, 0.392, 1.0, 1);
@@ -60,6 +226,34 @@ const getGradientColor = tgpu.fn([d.f32], d.vec4f) /* wgsl */`(ratio) {
 }`;
 ```
 
+## Implementing functions in WGSL
+
+:::note[Recommended reading]
+We assume that you are familiar with the following concepts:
+- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-fundamentals.html" target="_blank" rel="noopener noreferrer">WebGPU Fundamentals</a>
+- <a href="https://webgpufundamentals.org/webgpu/lessons/webgpu-wgsl.html" target="_blank" rel="noopener noreferrer">WebGPU Shading Language</a>
+:::
+
+```ts
+// "shelled" function:
+//   Same as a "bare" function, but with restricted argument
+//   and return types.
+//
+const rangeB = tgpu.fn([d.f32, d.f32], d.vec2f)((a, b) => {
+  return d.vec2f(a + b, a - b).mul(multiplier);
+});
+
+// "shelled wgsl" function:
+//   A shell defining the function's signature, and the
+//   implementation in WGSL.
+//
+const rangeC = tgpu.fn([d.f32, d.f32], d.vec2f)`(a, b) {
+  return vec2f(a + b, a - b) * multiplier;
+}
+`.$uses({ multiplier });
+// ^ Giving the function access to the 'multiplier' external value
+````
+
 ## External resources
 
 Functions can use external resources passed via the `$uses` method.
@@ -71,11 +265,11 @@ import * as d from 'typegpu/data';
 
 // ---cut---
 const getBlueFunction = tgpu.fn([], d.vec4f)`() {
-  return vec4f(0.114, 0.447, 0.941, 1); 
+  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 purple = d.vec4f(0.769, 0.392, 1.0, 1);
 
 const getGradientColor = tgpu.fn([d.f32], d.vec4f)`(ratio) {
   return mix(purple, getBlue(), ratio);;
@@ -88,7 +282,7 @@ You can check yourself what `getGradientColor` resolves to by calling [`tgpu.res
 ```wgsl
 // results of calling tgpu.resolve({ externals: { getGradientColor } })
 fn getBlueFunction_1() -> vec4f{
-  return vec4f(0.114, 0.447, 0.941, 1); 
+  return vec4f(0.114, 0.447, 0.941, 1);
 }
 
 fn getGradientColor_0(ratio: f32) -> vec4f{
@@ -113,9 +307,9 @@ Instead of annotating a `TgpuFn` with attributes, entry functions are defined us
 To describe the input and output of an entry point function, we use `IORecord`s, JavaScript objects that map argument names to their types.
 
 ```ts
-const vertexInput = { 
+const vertexInput = {
   idx: d.builtin.vertexIndex,
-  position: d.vec4f, 
+  position: d.vec4f,
   color: d.vec4f
 }
 ```
@@ -124,15 +318,15 @@ As you may note, builtin inter-stage inputs and outputs are available on the `d.
 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. 
+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 `IORecord`s.
 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, 
+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 `IORecord`s into automatically generated structs, you can also use those decorators in `IOStruct`s.

apps/typegpu-docs/src/content/docs/fundamentals/tgsl.mdx

@@ -1,19 +1,94 @@
 ---
-title: TGSL
+title: TS Kernels
 description: Guide on using JavaScript for WebGPU function definitions.
 ---
 
 :::caution[Experimental]
 This feature is under heavy development and is yet to reach stability.
 :::
 
-**TGSL (TypeGPU Shading Language)** is a subset of JavaScript used to define functions that run on the GPU via TypeGPU.
-It works by transpiling JavaScript into a compact AST format, called [tinyest](https://www.npmjs.com/package/tinyest),
-which is then used to generate equivalent WGSL.
+:::note[Requires unplugin-typegpu]
+To use TS kernels, you need to install and configure [unplugin-typegpu](/TypeGPU/tooling/unplugin-typegpu).
+:::
+
+**TS Kernels** are functions marked with the `'kernel'` directive, picked up by our dedicated build plugin -- [unplugin-typegpu](/TypeGPU/tooling/unplugin-typegpu) and transformed into a format TypeGPU can understand.
 
-You can check the current state of supported JavaScript syntax in the [tinyest-for-wgsl repository](https://github.com/software-mansion/TypeGPU/blob/release/packages/tinyest-for-wgsl/src/parsers.ts).
+```ts twoslash
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+function manhattanDistance(a: d.v3f, b: d.v3f) {
+  'tgpu';
+  const dx = std.abs(a.x - b.x);
+  const dy = std.abs(a.y - b.y);
+  const dz = std.abs(a.z - b.z);
+
+  return std.max(dx, std.max(dy, dz));
+}
+```
 
-For the TGSL functions to work, you need to use the dedicated build plugin -- [unplugin-typegpu](/TypeGPU/tooling/unplugin-typegpu).
+TypeGPU already provides a set of standard kernels under `typegpu/std`, which you can use in your own kernels. Our goal is for all kernels to have matching
+behavior on the CPU and GPU, which unlocks many possibilities.
+
+There are three main ways to use kernels.
+
+```ts twoslash
+import tgpu, { prepareDispatch } from 'typegpu';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+function manhattanDistance(a: d.v3f, b: d.v3f) {
+  'kernel';
+  const dx = std.abs(a.x - b.x);
+  const dy = std.abs(a.y - b.y);
+  const dz = std.abs(a.z - b.z);
+
+  return std.max(dx, std.max(dy, dz));
+}
+
+// ---cut---
+// #1) Can be called in JS
+const dist = manhattanDistance(d.vec3f(1, 2, 3), d.vec3f(6, 5, 4));
+
+// #2) Executed on the GPU
+const root = await tgpu.init();
+prepareDispatch(root, () => {
+  'kernel';
+  manhattanDistance(d.vec3f(1, 2, 3), d.vec3f(6, 5, 4));
+}).dispatch();
+
+// #3) Used to generate WGSL
+const wgsl = tgpu.resolve({ externals: { manhattanDistance } });
+````
+
+The contents of the `wgsl` variable will be the following:
+
+```wgsl
+// Generated WGSL
+fn manhattanDistance(a: vec3f, b: vec3f) -> f32 {
+  let dx = abs(a.x - b.x);
+  let dy = abs(a.y - b.y);
+  let dz = abs(a.z - b.z);
+
+  return max(dx, max(dy, dz));
+}
+````
+
+You can already notice a few things about TS Kernels:
+- Using operators like `+`, `-`, `*`, `/`, etc. is perfectly valid on numbers.
+- TS types are properly inferred, feel free to hover over the variables to see their types.
+- The generated code closely matches your source code.
+
+## Anatomy of a kernel
+
+:::note
+TS kernels are picked up by `unplugin-typegpu` during bundling, their AST is compacted into our custom format called [tinyest](https://www.npmjs.com/package/tinyest) and injected
+into the user code. This metadata is then used to generate equivalent WGSL at runtime efficiently.
+:::
+
+:::tip
+If you wish for all WGSL generation to happen at build time, combine [unplugin-macros](https://github.com/unplugin/unplugin-macros) and our [resolve API](/TypeGPU/fundamentals/resolve).
+:::
 
 ## Usage