uniformity example refactor

Author: cieplypolar

apps/typegpu-docs/src/examples/algorithms/probability/constants.ts

@@ -1,12 +1,10 @@
-import { Distribution, Generator } from './types.ts';
+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 initialGenerator: Generator = Generator.BPETER;
-export const generators: Generator[] = Object.values(Generator);
 export const initialDistribution: Distribution = Distribution.NORMAL;
 export const distributions: Distribution[] = Object.values(Distribution);

apps/typegpu-docs/src/examples/algorithms/probability/executor.ts

@@ -1,4 +1,4 @@
-import { randf, randomGeneratorSlot, type StatefulGenerator } from '@typegpu/noise';
+import { randf } from '@typegpu/noise';
 import type {
   StorageFlag,
   TgpuBindGroup,
@@ -26,8 +26,7 @@ export class Executor {
     number,
     [TgpuBuffer<d.WgslArray<d.Vec3f>> & StorageFlag, TgpuBuffer<d.WgslArray<d.F32>> & StorageFlag]
   >;
-  // they can be WeakMaps, because we always have reference to distribution and PRNG
-  readonly #pipelineCache: WeakMap<TgpuFn, WeakMap<StatefulGenerator, TgpuComputePipeline>>;
+  readonly #pipelineCache: Map<TgpuFn, TgpuComputePipeline>;
 
   constructor(root: TgpuRoot) {
     this.#root = root;
@@ -80,43 +79,21 @@ export class Executor {
     });
   }
 
-  #pipelineCacheSet(
-    distribution: TgpuFn<() => d.Vec3f>,
-    generator: StatefulGenerator,
-    pipeline: TgpuComputePipeline,
-  ) {
-    let distributionMap = this.#pipelineCache.get(distribution);
-    if (!distributionMap) {
-      distributionMap = new WeakMap();
-      this.#pipelineCache.set(distribution, distributionMap);
-    }
-
-    distributionMap.set(generator, pipeline);
-  }
-
-  pipelineCacheGet(
-    distribution: TgpuFn<() => d.Vec3f>,
-    generator: StatefulGenerator,
-  ): TgpuComputePipeline {
-    let pipeline = this.#pipelineCache.get(distribution)?.get(generator);
+  getPipeline(distribution: TgpuFn<() => d.Vec3f>) {
+    let pipeline = this.#pipelineCache.get(distribution);
     if (!pipeline) {
       pipeline = this.#root
-        .with(randomGeneratorSlot, generator)
         .with(this.#distributionSlot, distribution)
         .createComputePipeline({ compute: this.#dataMoreWorkersFunc });
-      this.#pipelineCacheSet(distribution, generator, pipeline);
+      this.#pipelineCache.set(distribution, pipeline);
     }
-
     return pipeline;
   }
 
-  async executeMoreWorkers(
-    distribution: TgpuFn<() => d.Vec3f>,
-    generator: StatefulGenerator,
-  ): Promise<d.v3f[]> {
-    const pipeline = this.pipelineCacheGet(distribution, generator);
-
-    pipeline.with(this.#bindGroup).dispatchWorkgroups(Math.ceil(this.#count / 64));
+  async executeMoreWorkers(distribution: TgpuFn<() => d.Vec3f>): Promise<d.v3f[]> {
+    this.getPipeline(distribution)
+      .with(this.#bindGroup)
+      .dispatchWorkgroups(Math.ceil(this.#count / 64));
 
     return await this.#samplesBuffer.read();
   }

apps/typegpu-docs/src/examples/algorithms/probability/helpers.ts

@@ -1,7 +1,8 @@
+import { randf } from '@typegpu/noise';
 import tgpu, { d } from 'typegpu';
-import { BPETER, LCG, randf, type StatefulGenerator } from '@typegpu/noise';
-import { Distribution, Generator, PlotType, type PRNG } from './types.ts';
+
 import * as c from './constants.ts';
+import { Distribution, PlotType, type PRNG } from './types.ts';
 
 const normal = d.vec3f(1.41, 1.41, 0);
 const z2D = 0.5;
@@ -94,10 +95,3 @@ const distributionCameras = {
 export function getCameraPosition(distribution: Distribution): number[] {
   return distributionCameras[distribution];
 }
-
-const GENERATOR_MAP = {
-  [Generator.BPETER]: BPETER,
-  [Generator.LCG]: LCG,
-};
-
-export const getGenerator = (gen: Generator): StatefulGenerator => GENERATOR_MAP[gen];

apps/typegpu-docs/src/examples/algorithms/probability/index.ts

@@ -2,9 +2,9 @@ import tgpu from 'typegpu';
 
 import { Plotter } from './plotter.ts';
 import { Executor } from './executor.ts';
-import type { Distribution, Generator } from './types.ts';
+import type { Distribution } from './types.ts';
 import * as c from './constants.ts';
-import { getCameraPosition, getGenerator, getPRNG } from './helpers.ts';
+import { getCameraPosition, getPRNG } from './helpers.ts';
 import { defineControls } from '../../common/defineControls.ts';
 
 const root = await tgpu.init();
@@ -13,17 +13,11 @@ const executor = new Executor(root);
 const plotter = new Plotter();
 
 let currentDistribution = c.initialDistribution;
-let currentGenerator = c.initialGenerator;
 
-const replot = async (
-  currentDistribution: Distribution,
-  currentGenerator: Generator,
-  animate = false,
-) => {
+const replot = async (currentDistribution: Distribution, animate = false) => {
   const prng = getPRNG(currentDistribution);
-  const gen = getGenerator(currentGenerator);
 
-  const samples = await executor.executeMoreWorkers(prng.prng, gen);
+  const samples = await executor.executeMoreWorkers(prng.prng);
   await plotter.plot(samples, prng, animate);
 };
 
@@ -70,19 +64,7 @@ export const controls = defineControls({
   Reseed: {
     async onButtonClick() {
       executor.reseed();
-      await replot(currentDistribution, currentGenerator, true);
-      plotter.resetView(getCameraPosition(currentDistribution));
-    },
-  },
-  Generator: {
-    initial: c.initialGenerator,
-    options: c.generators,
-    onSelectChange: async (value: Generator) => {
-      if (currentGenerator === value) {
-        return;
-      }
-      currentGenerator = value;
-      await replot(currentDistribution, currentGenerator, true);
+      await replot(currentDistribution, true);
       plotter.resetView(getCameraPosition(currentDistribution));
     },
   },
@@ -95,7 +77,7 @@ export const controls = defineControls({
       }
 
       currentDistribution = value;
-      await replot(currentDistribution, currentGenerator, true);
+      await replot(currentDistribution, true);
       plotter.resetView(getCameraPosition(currentDistribution));
     },
   },
@@ -104,26 +86,15 @@ export const controls = defineControls({
     options: c.numSamplesOptions,
     async onSelectChange(value) {
       executor.count = value;
-      await replot(currentDistribution, currentGenerator);
+      await replot(currentDistribution);
     },
   },
+  // this is the only place where some niche distributions are tested
   'Test Resolution': import.meta.env.DEV && {
     onButtonClick() {
-      for (const [i, j] of [
-        [0, 0],
-        [8, 1],
-        [10, 0],
-      ]) {
-        const code = tgpu.resolve({
-          externals: {
-            p: executor.pipelineCacheGet(
-              getPRNG(c.distributions[i]).prng,
-              getGenerator(c.generators[j]),
-            ),
-          },
-        });
-        root.device.createShaderModule({ code });
-      }
+      c.distributions
+        .map((dist) => tgpu.resolve([executor.getPipeline(getPRNG(dist).prng)]))
+        .forEach((r) => root.device.createShaderModule({ code: r }));
     },
   },
 });
@@ -132,4 +103,5 @@ export function onCleanup() {
   root.destroy();
   plotter.destroy();
 }
+
 // #endregion

apps/typegpu-docs/src/examples/algorithms/probability/types.ts

@@ -69,10 +69,3 @@ export interface SimplePRNG {
   plotType: PlotType;
   prng: TgpuFn<() => d.Vec3f>;
 }
-
-export const Generator = {
-  BPETER: 'bpeter (default)',
-  LCG: 'lcg',
-} as const;
-
-export type Generator = (typeof Generator)[keyof typeof Generator];

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

@@ -1,6 +1,5 @@
-import { PRNG } from './prngs.ts';
+import type { PRNGKey } 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 initialPRNG: PRNGKey = 'bpeter';

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

@@ -2,7 +2,7 @@ import { randf, randomGeneratorSlot } from '@typegpu/noise';
 import tgpu, { common, d, std, type TgpuRenderPipeline } 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();
@@ -11,29 +11,43 @@ 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,
+  useSeed2: d.u32,
+});
+
+const configUniform = root.createUniform(Config, {
+  gridSize: c.initialGridSize,
+  canvasRatio: canvas.width / canvas.height,
+  useSeed2: d.u32(prngs[initialPRNG].useSeed2),
+});
 
 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 gridSize = configUniform.$.gridSize;
+  const uv = input.uv * d.vec2f(configUniform.$.canvasRatio, 1);
+  const gridedUV = std.floor(uv * gridSize);
 
-  randf.seed2(gridedUV);
+  if (configUniform.$.useSeed2 === 1) {
+    randf.seed2(gridedUV);
+  } else {
+    randf.seed(gridedUV.x * gridSize + gridedUV.y);
+  }
 
   return d.vec4f(d.vec3f(randf.sample()), 1);
 });
 
-const pipelineCache = new Map<PRNG, TgpuRenderPipeline<d.Vec4f>>();
-let prng: PRNG = c.initialPRNG;
+const pipelineCache = new Map<PRNGKey, TgpuRenderPipeline<d.Vec4f>>();
+let prng: PRNGKey = initialPRNG;
 
 const redraw = () => {
   let pipeline = pipelineCache.get(prng);
   if (!pipeline) {
-    pipeline = root.with(randomGeneratorSlot, getPRNG(prng)).createRenderPipeline({
+    pipeline = root.with(randomGeneratorSlot, prngs[prng].generator).createRenderPipeline({
       vertex: common.fullScreenTriangle,
       fragment: fragmentShader,
       targets: { format: presentationFormat },
@@ -47,44 +61,41 @@ const redraw = () => {
 // #region Example controls & Cleanup
 export const controls = defineControls({
   PRNG: {
-    initial: c.initialPRNG,
-    options: c.prngs,
+    initial: initialPRNG,
+    options: prngKeys,
     onSelectChange: (value) => {
       prng = value;
+      configUniform.writePartial({ useSeed2: d.u32(prngs[value].useSeed2) });
       redraw();
     },
   },
   'Grid Size': {
     initial: c.initialGridSize,
     options: c.gridSizes,
     onSelectChange: (value) => {
-      gridSizeUniform.write(value);
+      configUniform.writePartial({ gridSize: value });
       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 },
-          ),
+      prngKeys
+        .map((key) =>
+          tgpu.resolve([
+            root.with(randomGeneratorSlot, prngs[key].generator).createRenderPipeline({
+              vertex: common.fullScreenTriangle,
+              fragment: fragmentShader,
+            }),
+          ]),
         )
-        .map((r) => root.device.createShaderModule({ code: r }));
+        .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,15 +1,19 @@
-import { BPETER, LCG, type StatefulGenerator } from '@typegpu/noise';
+import { BPETER, LCG32, XOROSHIRO64STARSTAR, type StatefulGenerator } from '@typegpu/noise';
 
-export const PRNG = {
-  BPETER: 'bpeter (default)',
-  LCG: 'lcg',
-} as const;
+interface PRNGOptions {
+  name: string;
+  useSeed2: boolean;
+  generator: StatefulGenerator;
+}
 
-export type PRNG = (typeof PRNG)[keyof typeof PRNG];
+export const prngs = {
+  bpeter: { name: 'bpeter (default)', useSeed2: true, generator: BPETER },
+  lcg32: { name: 'lcg32', useSeed2: false, generator: LCG32 },
+  xoroshiro64: { name: 'xoroshiro64', useSeed2: true, generator: XOROSHIRO64STARSTAR },
+} as const satisfies Record<string, PRNGOptions>;
 
-const PRNG_MAP = {
-  [PRNG.BPETER]: BPETER,
-  [PRNG.LCG]: LCG,
-};
+export type PRNGKey = keyof typeof prngs;
 
-export const getPRNG = (prng: PRNG): StatefulGenerator => PRNG_MAP[prng];
+export const prngKeys = Object.keys(prngs) as PRNGKey[];
+
+export const initialPRNG: PRNGKey = 'bpeter';