benchmarking prngs

Author: cieplypolar

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

@@ -3,3 +3,6 @@ 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: PRNGKey = 'bpeter';
+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,11 +1,11 @@
 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 { 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' });
@@ -15,47 +15,102 @@ const Config = d.struct({
   gridSize: d.f32,
   canvasRatio: d.f32,
   useSeed2: d.u32,
+  samplesPerThread: d.u32,
+  takeAverage: d.u32,
 });
 
 const configUniform = root.createUniform(Config, {
   gridSize: c.initialGridSize,
   canvasRatio: canvas.width / canvas.height,
   useSeed2: d.u32(prngs[initialPRNG].useSeed2),
+  samplesPerThread: c.initialSamplesPerThread,
+  takeAverage: d.u32(c.initialTakeAverage),
 });
 
-const fragmentShader = tgpu.fragmentFn({
-  in: { uv: d.vec2f },
-  out: d.vec4f,
-})((input) => {
+const layouts = {
+  compute: tgpu.bindGroupLayout({
+    texture: { storageTexture: d.textureStorage2d('r32float', 'write-only') },
+  }),
+  display: tgpu.bindGroupLayout({
+    texture: { storageTexture: d.textureStorage2d('r32float', 'read-only') },
+  }),
+};
+
+const bindGroups = Object.fromEntries(
+  c.gridSizes.map((size) => {
+    const texture = root['~unstable']
+      .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 computeFn = (x: number, y: number) => {
   'use gpu';
   const gridSize = configUniform.$.gridSize;
-  const uv = input.uv * d.vec2f(configUniform.$.canvasRatio, 1);
-  const gridedUV = std.floor(uv * gridSize);
 
   if (configUniform.$.useSeed2 === 1) {
-    randf.seed2(gridedUV);
+    randf.seed2(d.vec2f(x, y));
   } else {
-    randf.seed(gridedUV.x * gridSize + gridedUV.y);
+    randf.seed(d.f32(x) * gridSize + d.f32(y));
   }
 
-  return d.vec4f(d.vec3f(randf.sample()), 1);
-});
+  let i = d.u32(0);
+  const samplesPerThread = configUniform.$.samplesPerThread;
+  let samples = d.f32(0);
+  while (i < samplesPerThread - 1) {
+    samples += randf.sample();
+    i += 1;
+  }
 
-const pipelineCache = new Map<PRNGKey, TgpuRenderPipeline<d.Vec4f>>();
-let prng: PRNGKey = initialPRNG;
+  let result = randf.sample();
+  if (configUniform.$.takeAverage === 1) {
+    result = (result + samples) / samplesPerThread;
+  }
 
-const redraw = () => {
-  let pipeline = pipelineCache.get(prng);
+  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, prngs[prng].generator).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) / 1000} ms.`);
+      });
+    computePipelineCache.set(key, pipeline);
   }
+  return pipeline;
+};
+
+let prng = initialPRNG;
+let gridSize = c.initialGridSize;
 
-  pipeline.withColorAttachment({ view: context }).draw(3);
+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
@@ -69,26 +124,35 @@ export const controls = defineControls({
       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) => {
-      configUniform.writePartial({ gridSize: 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: () => {
       prngKeys
-        .map((key) =>
-          tgpu.resolve([
-            root.with(randomGeneratorSlot, prngs[key].generator).createRenderPipeline({
-              vertex: common.fullScreenTriangle,
-              fragment: fragmentShader,
-            }),
-          ]),
-        )
+        .map((key) => tgpu.resolve([getComputePipeline(key).pipeline]))
         .forEach((r) => root.device.createShaderModule({ code: r }));
     },
   },

apps/typegpu-docs/tests/individual-example-tests/probability.test.ts

@@ -15,7 +15,7 @@ describe('probability distribution plot example', () => {
         category: 'algorithms',
         name: 'probability',
         controlTriggers: ['Test Resolution'],
-        expectedCalls: 3,
+        expectedCalls: 13,
       },
       device,
     );