feat (@typegpu/noise): LCG32 and xoroshiro PRNGs, utilities

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

@@ -277,34 +277,33 @@ const f = tgpu.computeFn({ workgroupSize: [1] })(() => {
 });
 // ---cut---
 import {
+  hash,
+  randomGeneratorShell,
   randomGeneratorSlot,
+  u32To01F32,
   type StatefulGenerator,
 } from '@typegpu/noise';
 
-const LCG: StatefulGenerator = (() => {
+const LCG32: StatefulGenerator = (() => {
   const seed = tgpu.privateVar(d.u32);
 
-  const u32To01Float = tgpu.fn([d.u32], d.f32)((value) => {
-    const mantissa = value >> 9;
-    const bits = 0x3F800000 | mantissa;
-    const f = std.bitcastU32toF32(bits);
-    return f - 1;
-  });
-
-  return {
-    seed2: (value: d.v2f) => {
-      'use gpu';
-      seed.$ = d.u32(value.x * std.pow(32, 3) + value.y * std.pow(32, 2));
-    },
-    sample: () => {
-      'use gpu';
-      seed.$ = seed.$ * 1664525 + 1013904223; // % 2 ^ 32
-      return u32To01Float(seed.$);
-    },
-  };
+ const multiplier = 1664525;
+ const increment = 1013904223;
+
+ return {
+   seed: tgpu.fn([d.f32])((value) => {
+     seed.$ = hash(d.u32(value));
+   }),
+
+   sample: randomGeneratorShell(() => {
+     'use gpu';
+     seed.$ = multiplier * seed.$ + increment; // % 2 ^ 32
+     return u32To01F32(seed.$);
+   }).$name('sample'),
+ };
 })();
 
 const pipeline = root
-  .with(randomGeneratorSlot, LCG)
+  .with(randomGeneratorSlot, LCG32)
   .createComputePipeline({ compute: f });
 ```

apps/typegpu-docs/src/examples/tests/uniformity/constants.ts

@@ -1,6 +1,5 @@
-import { PRNG } from './prngs.ts';
-
 export const gridSizes = [8, 16, 32, 64, 128, 256, 512, 1024];
 export const initialGridSize = gridSizes[4];
-export const initialPRNG = PRNG.BPETER;
-export const prngs: PRNG[] = Object.values(PRNG);
+export const samplesPerThread = [1, 8, 16, 64, 256, 1024, 131072, 262144];
+export const initialSamplesPerThread = samplesPerThread[0];
+export const initialTakeAverage = false;

apps/typegpu-docs/src/examples/tests/uniformity/index.ts

@@ -1,90 +1,162 @@
 import { randf, randomGeneratorSlot } from '@typegpu/noise';
-import tgpu, { common, d, std, type TgpuRenderPipeline } from 'typegpu';
+import tgpu, { common, d, std, type TgpuGuardedComputePipeline } from 'typegpu';
 
 import * as c from './constants.ts';
-import { getPRNG, type PRNG } from './prngs.ts';
+import { initialPRNG, prngKeys, prngs, type PRNGKey } from './prngs.ts';
 import { defineControls } from '../../common/defineControls.ts';
 
-const root = await tgpu.init();
+const root = await tgpu.init({ device: { requiredFeatures: ['timestamp-query'] } });
 
 const canvas = document.querySelector('canvas') as HTMLCanvasElement;
 const context = root.configureContext({ canvas, alphaMode: 'premultiplied' });
 const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
 
-const gridSizeUniform = root.createUniform(d.f32, c.initialGridSize);
-const canvasRatioUniform = root.createUniform(d.f32, canvas.width / canvas.height);
+const Config = d.struct({
+  gridSize: d.f32,
+  canvasRatio: d.f32,
+  samplesPerThread: d.u32,
+  takeAverage: d.u32,
+});
 
-const fragmentShader = tgpu.fragmentFn({
-  in: { uv: d.vec2f },
-  out: d.vec4f,
-})((input) => {
-  'use gpu';
-  const uv = ((input.uv + 1) / 2) * d.vec2f(canvasRatioUniform.$, 1);
-  const gridedUV = std.floor(uv * gridSizeUniform.$);
+const configUniform = root.createUniform(Config, {
+  gridSize: c.initialGridSize,
+  canvasRatio: canvas.width / canvas.height,
+  samplesPerThread: c.initialSamplesPerThread,
+  takeAverage: d.u32(c.initialTakeAverage),
+});
 
-  randf.seed2(gridedUV);
+const layouts = {
+  compute: tgpu.bindGroupLayout({
+    texture: { storageTexture: d.textureStorage2d('r32float', 'write-only') },
+  }),
+  display: tgpu.bindGroupLayout({
+    texture: { storageTexture: d.textureStorage2d('r32float', 'read-only') },
+  }),
+};
 
-  return d.vec4f(d.vec3f(randf.sample()), 1);
+const bindGroups = Object.fromEntries(
+  c.gridSizes.map((size) => {
+    const texture = root
+      .createTexture({ size: [size, size], format: 'r32float' })
+      .$usage('storage', 'sampled');
+    return [
+      size,
+      {
+        compute: root.createBindGroup(layouts.compute, { texture }),
+        display: root.createBindGroup(layouts.display, { texture }),
+      },
+    ];
+  }),
+);
+
+const displayPipeline = root.createRenderPipeline({
+  vertex: common.fullScreenTriangle,
+  fragment: ({ uv }) => {
+    'use gpu';
+    const adjustedUv = uv * d.vec2f(configUniform.$.canvasRatio, 1);
+    const gridSize = configUniform.$.gridSize;
+    const coords = d.vec2u(std.floor(adjustedUv * gridSize));
+    const value = std.textureLoad(layouts.display.$.texture, coords).r;
+    return d.vec4f(d.vec3f(value), 1);
+  },
+  targets: { format: presentationFormat },
 });
 
-const pipelineCache = new Map<PRNG, TgpuRenderPipeline<d.Vec4f>>();
-let prng: PRNG = c.initialPRNG;
+const computeFn = (x: number, y: number) => {
+  'use gpu';
+  const gridSize = configUniform.$.gridSize;
 
-const redraw = () => {
-  let pipeline = pipelineCache.get(prng);
+  if (!randomGeneratorSlot.$.seed2) {
+    randf.seed(d.f32(x + 1) * gridSize + d.f32(y + 1));
+  } else {
+    randf.seed2(d.vec2f(x, y) + 1);
+  }
+
+  let i = d.u32(0);
+  const samplesPerThread = configUniform.$.samplesPerThread;
+  let samples = d.f32(0);
+  while (i < samplesPerThread - 1) {
+    samples += randf.sample();
+    i += 1;
+  }
+
+  let result = randf.sample();
+  if (configUniform.$.takeAverage === 1) {
+    result = (result + samples) / samplesPerThread;
+  }
+
+  std.textureStore(layouts.compute.$.texture, d.vec2u(x, y), d.vec4f(result, 0, 0, 0));
+};
+
+const computePipelineCache = new Map<PRNGKey, TgpuGuardedComputePipeline<[number, number]>>();
+const getComputePipeline = (key: PRNGKey) => {
+  let pipeline = computePipelineCache.get(key);
   if (!pipeline) {
-    pipeline = root.with(randomGeneratorSlot, getPRNG(prng)).createRenderPipeline({
-      vertex: common.fullScreenTriangle,
-      fragment: fragmentShader,
-      targets: { format: presentationFormat },
-    });
-    pipelineCache.set(prng, pipeline);
+    pipeline = root
+      .with(randomGeneratorSlot, prngs[key].generator)
+      .createGuardedComputePipeline(computeFn)
+      .withPerformanceCallback((start, end) => {
+        console.log(`[${key}] - ${Number(end - start) / 1_000_000} ms.`);
+      });
+    computePipelineCache.set(key, pipeline);
   }
+  return pipeline;
+};
 
-  pipeline.withColorAttachment({ view: context }).draw(3);
+let prng = initialPRNG;
+let gridSize = c.initialGridSize;
+
+const redraw = () => {
+  getComputePipeline(prng).with(bindGroups[gridSize].compute).dispatchThreads(gridSize, gridSize);
+  displayPipeline.withColorAttachment({ view: context }).with(bindGroups[gridSize].display).draw(3);
 };
 
 // #region Example controls & Cleanup
 export const controls = defineControls({
   PRNG: {
-    initial: c.initialPRNG,
-    options: c.prngs,
+    initial: initialPRNG,
+    options: prngKeys,
     onSelectChange: (value) => {
       prng = value;
       redraw();
     },
   },
+  'Samples per thread': {
+    initial: c.initialSamplesPerThread,
+    options: c.samplesPerThread,
+    onSelectChange: (value) => {
+      configUniform.writePartial({ samplesPerThread: value });
+      redraw();
+    },
+  },
+  'Take Average': {
+    initial: c.initialTakeAverage,
+    onToggleChange: (value) => {
+      configUniform.writePartial({ takeAverage: d.u32(value) });
+      redraw();
+    },
+  },
   'Grid Size': {
     initial: c.initialGridSize,
     options: c.gridSizes,
     onSelectChange: (value) => {
-      gridSizeUniform.write(value);
+      gridSize = value;
+      configUniform.writePartial({ gridSize });
       redraw();
     },
   },
+  // this is the only place where some niche prngs are tested
   'Test Resolution': import.meta.env.DEV && {
     onButtonClick: () => {
-      const namespace = tgpu['~unstable'].namespace();
-      c.prngs
-        .map((prng) =>
-          tgpu.resolve(
-            [
-              root.with(randomGeneratorSlot, getPRNG(prng)).createRenderPipeline({
-                vertex: common.fullScreenTriangle,
-                fragment: fragmentShader,
-                targets: { format: presentationFormat },
-              }),
-            ],
-            { names: namespace },
-          ),
-        )
-        .map((r) => root.device.createShaderModule({ code: r }));
+      prngKeys
+        .map((key) => tgpu.resolve([getComputePipeline(key).pipeline]))
+        .forEach((r) => root.device.createShaderModule({ code: r }));
     },
   },
 });
 
 const resizeObserver = new ResizeObserver(() => {
-  canvasRatioUniform.write(canvas.width / canvas.height);
+  configUniform.writePartial({ canvasRatio: canvas.width / canvas.height });
   redraw();
 });
 resizeObserver.observe(canvas);

apps/typegpu-docs/src/examples/tests/uniformity/prngs.ts

@@ -1,17 +1,17 @@
-import { BPETER, type StatefulGenerator } from '@typegpu/noise';
+import { BPETER, LCG32, XOROSHIRO64STARSTAR, type StatefulGenerator } from '@typegpu/noise';
 
-import { LCG } from './lcg.ts';
+interface PRNGOptions {
+  name: string;
+  generator: StatefulGenerator;
+}
 
-export const PRNG = {
-  BPETER: 'bpeter (default)',
-  LCG: 'lcg',
-} as const;
+export const prngs = {
+  bpeter: { name: 'bpeter (default)', generator: BPETER },
+  lcg32: { name: 'lcg32', generator: LCG32 },
+  xoroshiro64: { name: 'xoroshiro64', generator: XOROSHIRO64STARSTAR },
+} as const satisfies Record<string, PRNGOptions>;
 
-export type PRNG = (typeof PRNG)[keyof typeof PRNG];
+export type PRNGKey = keyof typeof prngs;
 
-const PRNG_MAP = {
-  [PRNG.BPETER]: BPETER,
-  [PRNG.LCG]: LCG,
-};
-
-export const getPRNG = (prng: PRNG): StatefulGenerator => PRNG_MAP[prng];
+export const prngKeys = Object.keys(prngs) as PRNGKey[];
+export const initialPRNG: PRNGKey = prngKeys[0];