docs: Generator slot

apps/typegpu-docs/astro.config.mjs

@@ -4,13 +4,15 @@ import react from '@astrojs/react';
 import sitemap from '@astrojs/sitemap';
 import starlight from '@astrojs/starlight';
 import tailwindVite from '@tailwindcss/vite';
-import basicSsl from '@vitejs/plugin-basic-ssl';
 import { defineConfig } from 'astro/config';
 import starlightBlog from 'starlight-blog';
 import starlightTypeDoc, { typeDocSidebarGroup } from 'starlight-typedoc';
 import typegpu from 'unplugin-typegpu/rollup';
 import { imagetools } from 'vite-imagetools';
 import wasm from 'vite-plugin-wasm';
+import basicSsl from '@vitejs/plugin-basic-ssl';
+import rehypeMathJax from 'rehype-mathjax';
+import remarkMath from 'remark-math';
 
 /**
  * @template T
@@ -25,6 +27,10 @@ const DEV = import.meta.env.DEV;
 export default defineConfig({
   site: 'https://docs.swmansion.com',
   base: 'TypeGPU',
+  markdown: {
+    remarkPlugins: [remarkMath],
+    rehypePlugins: [rehypeMathJax],
+  },
   vite: {
     // Allowing query params, for invalidation
     plugins: [
@@ -48,7 +54,11 @@ export default defineConfig({
   integrations: [
     starlight({
       title: 'TypeGPU',
-      customCss: ['./src/tailwind.css', './src/fonts/font-face.css'],
+      customCss: [
+        './src/tailwind.css',
+        './src/fonts/font-face.css',
+        './src/mathjax.css',
+      ],
       plugins: stripFalsy([
         starlightBlog({
           navigation: 'none',

apps/typegpu-docs/package.json

@@ -39,10 +39,13 @@
     "lucide-react": "^0.536.0",
     "lz-string": "^1.5.0",
     "monaco-editor": "^0.52.2",
+    "morphcharts": "^1.3.2",
     "motion": "^12.23.7",
     "pathe": "^2.0.3",
     "react": "^19.1.0",
     "react-dom": "^19.1.0",
+    "rehype-mathjax": "^7.1.0",
+    "remark-math": "^6.0.0",
     "remeda": "^2.21.2",
     "sharp": "^0.34.2",
     "starlight-blog": "^0.23.2",

apps/typegpu-docs/src/components/design/Select.tsx

@@ -21,7 +21,7 @@ export function Select({
           position='popper'
           className='-top-4 relative min-w-30'
         >
-          <RadixSelect.Viewport className='rounded bg-tameplum-50'>
+          <RadixSelect.Viewport className='max-h-[50vh] rounded bg-tameplum-50'>
             {options.map((option) => (
               <RadixSelect.Item
                 key={option}

apps/typegpu-docs/src/content/docs/ecosystem/typegpu-noise.mdx

@@ -10,7 +10,7 @@ It also features a [Perlin noise](#perlin-noise) implementation, which is useful
 effects, terrains, and other procedural elements.
 
 :::note
-Threads do not share the generator's `State`. As a result, unless you change the seed or provide thread-dependent variables, each thread will produce the same sequence of sampled values.
+Threads do not share the generator's `State`. As a result, unless you change the seed in each thread, each thread will produce the same sequence of sampled values.
 :::
 
 ## Use with either TypeGPU or WebGPU
@@ -80,6 +80,7 @@ Does this mean we allow object access inside of WGSL shaders?... yes, yes we do
 The `@typegpu/noise` package provides a pseudo-random number generator (PRNG) that generates uniformly distributed random numbers in the range `[0, 1)`.
 Each call to `randf.sample` returns the next random float in the sequence, allowing for predictable and repeatable results. The seed can be set or reset
 using a set of `randf.seedN` functions, where `N` is the number of components our seed has.
+
 ```ts twoslash
 import tgpu from 'typegpu';
 import * as d from 'typegpu/data';
@@ -100,10 +101,43 @@ const main = tgpu['~unstable'].fragmentFn({
 });
 ```
 
-There are higher-level utilities built on top of `randf.sample`:
+:::tip
+Due to limited float precision, for large seeds the samples tend to repeat quickly. To avoid this, keep the seed in the range `[0, 1]`.
+:::
+
+There are higher-level utilities built on top of `randf.sample`. Most of the distributions can be derived from the fact that cumulative distribution functions (CDFs) are uniformly distributed on`[0, 1]` and can be easily inverted to generate random variables. We provide insights for probability enthusiasts. Let's assume $U \sim \mathrm{Uniform}[0, 1) \sim \mathrm{randf.sample}$ :
+
+### Discrete
+- `bernoulli` - returns 1 with probability `p`, and 0 otherwise. You can think of it as the flip of a biased coin. `p` must be in the range `[0, 1]`
+$$
+\mathrm{Bernoulli}(p) = \begin{cases}
+1 & \text{if } U \leq p \\
+0 & \text{otherwise}
+\end{cases}
+$$
+
+### Continuous
+- `sampleExclusive` - works the same as `sample`, but returns a value strictly between 0 and 1
+- `exponential` - returns a number based on the exponential distribution with a given `rate`. The `rate` must be greater than 0. It is commonly used for modeling the time until an event occurs
+$$
+\mathrm{Exponential}(\lambda) = -\frac{1}{\lambda} \ln(1 - U)
+$$
+
+- `normal` - returns a number based on the normal (Gaussian) distribution with a given mean (`mu`) and standard deviation (`sigma`). `sigma` must be `> 0`. Following transformation is called Box-Muller transform
+$$
+\mathrm{Normal}(\mu, \sigma) = \mu + \sigma \sin(2 \pi U_1) \sqrt{-2 \ln(U_2)}
+$$
+
+- `cauchy` - returns a number based on the Cauchy distribution with a given origin point (`x0`, `gamma`). `gamma` must be `> 0`. It can be interpreted as the probability distribution of the x-intercept of a light ray from the origin.
+
+### Geometric
 - `inUnitCircle` - returns a random 2D vector uniformly distributed inside a unit circle
+- `onUnitCircle` - returns a random 2D vector uniformly distributed on the perimeter of a unit circle
 - `inUnitCube` - returns a random 3D vector uniformly distributed inside a unit cube
-- `onHemisphere` - returns a random 3D vector uniformly distributed on the surface of the upper hemisphere oriented accordingly to given normal vector
+- `onUnitCube` - returns a random 3D vector uniformly distributed on the surface of a unit cube
+- `inHemisphere` - returns a random 3D vector uniformly distributed inside an upper hemisphere oriented according to a given normal vector
+- `onHemisphere` - returns a random 3D vector uniformly distributed on the surface of an upper hemisphere oriented according to a given normal vector
+- `inUnitSphere` - returns a random 3D vector uniformly distributed inside a unit sphere. Thanks to the nature of the normal distribution, we can derive this distribution by sampling three coordinates from the $\mathrm{Normal}(0, 1)$ and normalizing the result (this gives us a direction vector). Finally, we choose the radius as $\sqrt[3]{U}$
 - `onUnitSphere` - returns a random 3D vector uniformly distributed on the surface of a unit sphere
 
 ## Perlin noise
@@ -233,3 +267,69 @@ pipeline
 // a different domain size
 bindGroup = initBindGroup(d.vec3u(5, 5, 1));
 ```
+
+## Generator slot
+The package provides convenient way to change PRNG as needed.
+You can easily implement your own PRNG and plug it into the pipeline.
+
+```ts twoslash
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+import { randf } from '@typegpu/noise';
+import * as std from 'typegpu/std';
+
+const root = await tgpu.init();
+
+const b = root.createMutable(d.f32);
+const f = tgpu['~unstable'].computeFn({ workgroupSize: [1] })(() => {
+  b.$ = randf.sample();
+});
+// ---cut---
+import {
+  randomGeneratorShell,
+  randomGeneratorSlot,
+  type StatefulGenerator,
+} from '@typegpu/noise';
+
+const LCG: StatefulGenerator = (() => {
+  const seed = tgpu['~unstable'].privateVar(d.u32);
+
+  const u32ToFloat = tgpu.fn([d.u32], d.f32)`(val){
+    let exponent: u32 = 0x3f800000;
+    let mantissa: u32 = 0x007fffff & val;
+    var ufloat: u32 = (exponent | mantissa);
+    return bitcast<f32>(ufloat) - 1f;
+  }`;
+
+  return {
+    seed: tgpu.fn([d.f32])((value) => {
+      seed.$ = d.u32(value * std.pow(32, 3));
+    }),
+    seed2: tgpu.fn([d.vec2f])((value) => {
+      seed.$ = d.u32(value.x * std.pow(32, 3) + value.y * std.pow(32, 2));
+    }),
+    seed3: tgpu.fn([d.vec3f])((value) => {
+      seed.$ = d.u32(
+        value.x * std.pow(32, 3) + value.y * std.pow(32, 2) +
+          value.z * std.pow(32, 1),
+      );
+    }),
+    seed4: tgpu.fn([d.vec4f])((value) => {
+      seed.$ = d.u32(
+        value.x * std.pow(32, 3) + value.y * std.pow(32, 2) +
+          value.z * std.pow(32, 1) + value.w * std.pow(32, 0),
+      );
+    }),
+    sample: randomGeneratorShell(() => {
+      'kernel';
+      seed.$ = seed.$ * 1664525 + 1013904223; // implicit modulo 2 ^ 32
+      return u32ToFloat(seed.$);
+    }),
+  };
+})();
+
+const pipeline = root['~unstable']
+  .with(randomGeneratorSlot, LCG)
+  .withCompute(f)
+  .createPipeline();
+```

apps/typegpu-docs/src/content/examples/simple/probability/constants.ts

@@ -0,0 +1,10 @@
+import { Distribution } from './types.ts';
+
+export const popupCooldown = 100000;
+export const cameraPositionGeo = [0, 0, 0.5];
+export const cameraPositionHist = [0, 0, 0.2];
+export const initialCameraAngle = 15;
+export const numSamplesOptions = [100, 1000, 2000, 5000, 10000, 50000];
+export const initialNumSamples = numSamplesOptions[2];
+export const initialDistribution: Distribution = Distribution.NORMAL;
+export const distributions: Distribution[] = Object.values(Distribution);

apps/typegpu-docs/src/content/examples/simple/probability/executor.ts

@@ -0,0 +1,134 @@
+import tgpu from 'typegpu';
+import type {
+  StorageFlag,
+  TgpuBindGroup,
+  TgpuBindGroupLayout,
+  TgpuBuffer,
+  TgpuComputeFn,
+  TgpuComputePipeline,
+  TgpuFn,
+  TgpuRoot,
+  TgpuSlot,
+} from 'typegpu';
+import * as d from 'typegpu/data';
+import { randf } from '@typegpu/noise';
+
+export class Executor {
+  // don't exceed max workgroup grid X dimension size
+  #count!: number;
+  #samplesBuffer!:
+    & TgpuBuffer<d.WgslArray<d.Vec3f>>
+    & StorageFlag;
+  #seedBuffer!:
+    & TgpuBuffer<d.WgslArray<d.F32>>
+    & StorageFlag;
+  #bindGroup!: TgpuBindGroup;
+  readonly #root: TgpuRoot;
+  readonly #dataMoreWorkersFunc: TgpuComputeFn;
+  readonly #distributionSlot: TgpuSlot<TgpuFn<() => d.Vec3f>>;
+  readonly #bindGroupLayout: TgpuBindGroupLayout;
+  readonly #bufferCache: Map<
+    number,
+    [
+      & TgpuBuffer<d.WgslArray<d.Vec3f>>
+      & StorageFlag,
+      & TgpuBuffer<d.WgslArray<d.F32>>
+      & StorageFlag,
+    ]
+  >;
+  readonly #pipelineCache: Map<TgpuFn, TgpuComputePipeline>;
+
+  constructor(root: TgpuRoot) {
+    this.#root = root;
+    this.#bufferCache = new Map();
+    this.#pipelineCache = new Map();
+
+    const distributionSlotTempAlias = tgpu.slot<TgpuFn<() => d.Vec3f>>();
+    this.#distributionSlot = distributionSlotTempAlias;
+
+    const bindGroupLayoutTempAlias = tgpu.bindGroupLayout({
+      seedBuffer: { storage: d.arrayOf(d.f32), access: 'readonly' },
+      samplesBuffer: { storage: d.arrayOf(d.vec3f), access: 'mutable' },
+    });
+    this.#bindGroupLayout = bindGroupLayoutTempAlias;
+
+    this.#dataMoreWorkersFunc = tgpu['~unstable'].computeFn({
+      in: { gid: d.builtin.globalInvocationId },
+      workgroupSize: [64],
+    })((input) => {
+      const id = input.gid.x;
+      if (id >= bindGroupLayoutTempAlias.$.samplesBuffer.length) return;
+      randf.seed(bindGroupLayoutTempAlias.$.seedBuffer[id]);
+      bindGroupLayoutTempAlias.$.samplesBuffer[id] = distributionSlotTempAlias
+        .$();
+    });
+  }
+
+  reseed() {
+    this.#seedBuffer.write(
+      Array.from({ length: this.#count }, () => Math.random()),
+    );
+  }
+
+  set count(value: number) {
+    this.#count = value;
+    if (this.#bufferCache.has(value)) {
+      // biome-ignore lint/style/noNonNullAssertion: just checked it above
+      [this.#samplesBuffer, this.#seedBuffer] = this.#bufferCache.get(value)!;
+    } else {
+      this.#samplesBuffer = this.#root
+        .createBuffer(d.arrayOf(d.vec3f, value))
+        .$usage('storage');
+      this.#seedBuffer = this.#root
+        .createBuffer(
+          d.arrayOf(d.f32, value),
+          Array.from({ length: value }, () => Math.random()),
+        ).$usage('storage');
+      this.#bufferCache.set(value, [this.#samplesBuffer, this.#seedBuffer]);
+    }
+
+    this.#bindGroup = this.#root.createBindGroup(this.#bindGroupLayout, {
+      seedBuffer: this.#seedBuffer,
+      samplesBuffer: this.#samplesBuffer,
+    });
+  }
+
+  cachedPipeline(distribution: TgpuFn<() => d.Vec3f>) {
+    if (!import.meta.env.DEV) {
+      throw new Error('Function only for testing purposes');
+    }
+
+    if (!this.#pipelineCache.has(distribution)) {
+      const pipeline = this.#root['~unstable']
+        .with(this.#distributionSlot, distribution)
+        .withCompute(this.#dataMoreWorkersFunc as TgpuComputeFn)
+        .createPipeline();
+      this.#pipelineCache.set(distribution, pipeline);
+    }
+
+    // biome-ignore lint/style/noNonNullAssertion: just checked it above
+    return this.#pipelineCache.get(distribution)!;
+  }
+
+  async executeMoreWorkers(
+    distribution: TgpuFn<() => d.Vec3f>,
+  ): Promise<d.v3f[]> {
+    let pipeline: TgpuComputePipeline;
+    if (this.#pipelineCache.has(distribution)) {
+      // biome-ignore lint/style/noNonNullAssertion: just checked it above
+      pipeline = this.#pipelineCache.get(distribution)!;
+    } else {
+      pipeline = this.#root['~unstable']
+        .with(this.#distributionSlot, distribution)
+        .withCompute(this.#dataMoreWorkersFunc as TgpuComputeFn)
+        .createPipeline();
+      this.#pipelineCache.set(distribution, pipeline);
+    }
+
+    pipeline
+      .with(this.#bindGroupLayout, this.#bindGroup)
+      .dispatchWorkgroups(Math.ceil(this.#count / 64));
+
+    return await this.#samplesBuffer.read();
+  }
+}