impr: Small improvements for future concurrent-scan work

apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts

@@ -2,12 +2,37 @@ import tgpu from 'typegpu';
 import * as d from 'typegpu/data';
 import { type BinaryOp, prefixScan, scan } from '@typegpu/sort';
 import * as std from 'typegpu/std';
+import { defineControls } from '../../common/defineControls.ts';
 import { addFn, concat10, isArrayEqual, mulFn, prefixScanJS, scanJS } from './functions.ts';
 
 const root = await tgpu.init({
   device: { requiredFeatures: ['timestamp-query'] },
 });
 
+function compareAndLog(actual: number[], expected: number[]): boolean {
+  if (isArrayEqual(actual, expected)) {
+    return true;
+  }
+
+  if (actual.length !== expected.length) {
+    console.error(`  Mismatch: length ${actual.length} !== ${expected.length}`);
+  } else if (actual.length <= 32) {
+    console.error('  actual:  ', actual);
+    console.error('  expected:', expected);
+  } else {
+    const idx = actual.findIndex((v, i) => v !== expected[i]);
+    const lo = Math.max(0, idx - 2);
+    const hi = Math.min(actual.length, idx + 3);
+    console.error(
+      `  first mismatch at index ${idx} (showing [${lo}..${hi - 1}] of ${actual.length}):`,
+    );
+    console.error('  actual:  ', actual.slice(lo, hi));
+    console.error('  expected:', expected.slice(lo, hi));
+  }
+
+  return false;
+}
+
 async function runAndCompare(arr: number[], op: BinaryOp, scanOnly: boolean) {
   const input = root.createBuffer(d.arrayOf(d.f32, arr.length), arr).$usage('storage');
 
@@ -24,7 +49,9 @@ async function runAndCompare(arr: number[], op: BinaryOp, scanOnly: boolean) {
         identityElement: op.identityElement,
       });
 
-  return isArrayEqual(await output.read(), scanOnly ? scanJS(arr, op) : prefixScanJS(arr, op));
+  const actual = await output.read();
+  const expected = scanOnly ? scanJS(arr, op) : prefixScanJS(arr, op);
+  return compareAndLog(actual, expected);
 }
 
 // single element f32 tests
@@ -88,7 +115,7 @@ async function testDoesNotDestroyBuffer(): Promise<boolean> {
     identityElement: 0,
   });
 
-  return isArrayEqual(await input.read(), [1, 2, 3, 4, 5, 6, 7, 8]);
+  return compareAndLog(await input.read(), [1, 2, 3, 4, 5, 6, 7, 8]);
 }
 
 async function testDoesNotCacheBuffers(): Promise<boolean> {
@@ -115,7 +142,7 @@ async function testDoesNotCacheBuffers(): Promise<boolean> {
     identityElement: op.identityElement,
   });
 
-  return isArrayEqual(await output1.read(), [36]) && isArrayEqual(await output2.read(), [10]);
+  return compareAndLog(await output1.read(), [36]) && compareAndLog(await output2.read(), [10]);
 }
 
 // prefix f32 tests
@@ -179,7 +206,7 @@ async function testPrefixDoesNotDestroyBuffer(): Promise<boolean> {
     operation: addFn,
     identityElement: 0,
   });
-  return isArrayEqual(await input.read(), [1, 2, 3, 4, 5, 6, 7, 8]);
+  return compareAndLog(await input.read(), [1, 2, 3, 4, 5, 6, 7, 8]);
 }
 
 async function testPrefixDoesNotCacheBuffers(): Promise<boolean> {
@@ -206,37 +233,88 @@ async function testPrefixDoesNotCacheBuffers(): Promise<boolean> {
   });
 
   return (
-    isArrayEqual(await output1.read(), prefixScanJS(arr1, op)) &&
-    isArrayEqual(await output2.read(), prefixScanJS(arr2, op))
+    compareAndLog(await output1.read(), prefixScanJS(arr1, op)) &&
+    compareAndLog(await output2.read(), prefixScanJS(arr2, op))
   );
 }
 
+// benchmark
+
+const BENCH_SIZES = [2_048, 65_536, 1_048_576, 16_777_216];
+const BENCH_WARMUP = 3;
+const BENCH_RUNS = 10;
+
+async function benchmarkSize(size: number): Promise<number> {
+  const inputData = Array.from({ length: size }, () => 1);
+  const inputBuffer = root.createBuffer(d.arrayOf(d.f32, size), inputData).$usage('storage');
+  const outputBuffer = root.createBuffer(d.arrayOf(d.f32, size)).$usage('storage');
+
+  for (let i = 0; i < BENCH_WARMUP; i++) {
+    prefixScan(root, { inputBuffer, outputBuffer, operation: addFn, identityElement: 0 });
+    await root.device.queue.onSubmittedWorkDone();
+  }
+
+  let total = 0;
+  for (let i = 0; i < BENCH_RUNS; i++) {
+    const t0 = performance.now();
+    prefixScan(root, { inputBuffer, outputBuffer, operation: addFn, identityElement: 0 });
+    await root.device.queue.onSubmittedWorkDone();
+    total += performance.now() - t0;
+  }
+
+  const avgMs = total / BENCH_RUNS;
+  inputBuffer.destroy();
+  outputBuffer.destroy();
+  return avgMs;
+}
+
+async function runBenchmarks(): Promise<void> {
+  console.log('=== Prefix Scan Benchmark ===');
+  for (const size of BENCH_SIZES) {
+    const avgMs = await benchmarkSize(size);
+    console.log(
+      `  size ${size.toLocaleString().padStart(12)}: ${avgMs.toFixed(2)} ms avg (${BENCH_RUNS} runs)`,
+    );
+  }
+  console.log('==============================');
+}
+
 // running the tests
 
+async function runTest(name: string, fn: () => Promise<boolean>): Promise<boolean> {
+  const passed = await fn();
+  if (!passed) {
+    console.error(`FAILED: ${name}`);
+  }
+  return passed;
+}
+
 async function runTests(): Promise<boolean> {
   let result = true;
 
-  result = (await testAdd8()) && result;
-  result = (await testAdd123()) && result;
-  result = (await testMul()) && result;
-  result = (await testStdMax()) && result;
-  result = (await testConcat()) && result;
-  result = (await testLength1()) && result;
-  result = (await testLength65537()) && result;
-  result = (await testLength16777217()) && result;
-  result = (await testDoesNotDestroyBuffer()) && result;
-  result = (await testDoesNotCacheBuffers()) && result;
-
-  result = (await testPrefixAdd8()) && result;
-  result = (await testPrefixAdd123()) && result;
-  result = (await testPrefixMul()) && result;
-  result = (await testPrefixStdMax()) && result;
-  result = (await testPrefixConcat()) && result;
-  result = (await testPrefixLength1()) && result;
-  result = (await testPrefixLength65537()) && result;
-  result = (await testPrefixLength16777217()) && result;
-  result = (await testPrefixDoesNotDestroyBuffer()) && result;
-  result = (await testPrefixDoesNotCacheBuffers()) && result;
+  result = (await runTest('testAdd8', testAdd8)) && result;
+  result = (await runTest('testAdd123', testAdd123)) && result;
+  result = (await runTest('testMul', testMul)) && result;
+  result = (await runTest('testStdMax', testStdMax)) && result;
+  result = (await runTest('testConcat', testConcat)) && result;
+  result = (await runTest('testLength1', testLength1)) && result;
+  result = (await runTest('testLength65537', testLength65537)) && result;
+  result = (await runTest('testLength16777217', testLength16777217)) && result;
+  result = (await runTest('testDoesNotDestroyBuffer', testDoesNotDestroyBuffer)) && result;
+  result = (await runTest('testDoesNotCacheBuffers', testDoesNotCacheBuffers)) && result;
+
+  result = (await runTest('testPrefixAdd8', testPrefixAdd8)) && result;
+  result = (await runTest('testPrefixAdd123', testPrefixAdd123)) && result;
+  result = (await runTest('testPrefixMul', testPrefixMul)) && result;
+  result = (await runTest('testPrefixStdMax', testPrefixStdMax)) && result;
+  result = (await runTest('testPrefixConcat', testPrefixConcat)) && result;
+  result = (await runTest('testPrefixLength1', testPrefixLength1)) && result;
+  result = (await runTest('testPrefixLength65537', testPrefixLength65537)) && result;
+  result = (await runTest('testPrefixLength16777217', testPrefixLength16777217)) && result;
+  result =
+    (await runTest('testPrefixDoesNotDestroyBuffer', testPrefixDoesNotDestroyBuffer)) && result;
+  result =
+    (await runTest('testPrefixDoesNotCacheBuffers', testPrefixDoesNotCacheBuffers)) && result;
 
   return result;
 }
@@ -245,12 +323,50 @@ const table = document.querySelector<HTMLDivElement>('.result');
 if (!table) {
   throw new Error('Nowhere to display the results');
 }
+
+let testsPassed: boolean | null = null;
+let benchmarkPromise: Promise<void> | null = null;
+
 void runTests().then((result) => {
+  testsPassed = result;
   table.innerText = `Tests ${result ? 'succeeded' : 'failed'}.`;
 });
 
+async function startBenchmarks(): Promise<void> {
+  if (testsPassed === null) {
+    table.innerText = 'Tests are still running.';
+    return;
+  }
+
+  if (!testsPassed) {
+    table.innerText = 'Tests failed. Benchmarks skipped.';
+    return;
+  }
+
+  if (benchmarkPromise) {
+    return benchmarkPromise;
+  }
+
+  table.innerText = 'Tests succeeded. Running benchmarks...';
+  benchmarkPromise = runBenchmarks()
+    .then(() => {
+      table.innerText = 'Tests succeeded. Benchmark complete (see console).';
+    })
+    .finally(() => {
+      benchmarkPromise = null;
+    });
+
+  return benchmarkPromise;
+}
+
 // #region Example controls and cleanup
 
+export const controls = defineControls({
+  'Run benchmarks': {
+    onButtonClick: startBenchmarks,
+  },
+});
+
 export function onCleanup() {
   root.destroy();
 }

packages/typegpu-sort/src/scan/compute/applySums.ts

@@ -1,5 +1,5 @@
-import tgpu, { d } from 'typegpu';
-import { operatorSlot, uniformOpLayout, WORKGROUP_SIZE } from '../schemas.ts';
+import tgpu, { d, std } from 'typegpu';
+import { ELEMENTS_PER_THREAD, operatorSlot, uniformOpLayout, WORKGROUP_SIZE } from '../schemas.ts';
 
 export const uniformOp = tgpu.computeFn({
   workgroupSize: [WORKGROUP_SIZE],
@@ -10,11 +10,10 @@ export const uniformOp = tgpu.computeFn({
 })(({ gid, wid }) => {
   const globalIdx = gid.x;
   const workgroupId = wid.x;
-  const baseIdx = globalIdx * 8;
+  const baseIdx = globalIdx * ELEMENTS_PER_THREAD;
   const opValue = uniformOpLayout.$.sums[workgroupId];
 
-  // TODO: use `tgpu.unroll(8)`
-  for (let i = d.u32(0); i < 8; i++) {
+  for (const i of tgpu.unroll(std.range(ELEMENTS_PER_THREAD))) {
     if (baseIdx + i < uniformOpLayout.$.input.length) {
       (uniformOpLayout.$.input[baseIdx + i] as number) = operatorSlot.$(
         opValue as number,

packages/typegpu-sort/src/scan/compute/scan.ts

@@ -1,5 +1,6 @@
 import tgpu, { d, std } from 'typegpu';
 import {
+  ELEMENTS_PER_THREAD,
   identitySlot,
   onlyGreatestElementSlot,
   operatorSlot,
@@ -8,7 +9,9 @@ import {
 } from '../schemas.ts';
 import { downsweep, upsweep, workgroupMemory } from './shared.ts';
 
-const fillIdentityArray = tgpu.comptime(() => Array.from({ length: 8 }, () => identitySlot.$));
+const fillIdentityArray = tgpu.comptime(() =>
+  Array.from({ length: ELEMENTS_PER_THREAD }, () => identitySlot.$),
+);
 
 export const computeBlock = tgpu.computeFn({
   workgroupSize: [WORKGROUP_SIZE],
@@ -22,19 +25,17 @@ export const computeBlock = tgpu.computeFn({
   const workgroupId = wid.x;
   const localIdx = lid.x;
 
-  // 8 elements per thread
-  const baseIdx = globalIdx * 8;
+  const baseIdx = globalIdx * ELEMENTS_PER_THREAD;
 
-  const partialSums = d.arrayOf(d.f32, 8)(fillIdentityArray());
+  const partialSums = d.arrayOf(d.f32, ELEMENTS_PER_THREAD)(fillIdentityArray());
 
   let prev = identitySlot.$;
   let lastIdx = d.u32(0);
 
-  // TODO: use `tgpu.unroll(8)`
-  for (let i = d.u32(); i < 8; i++) {
+  for (const i of tgpu.unroll(std.range(ELEMENTS_PER_THREAD))) {
     if (baseIdx + i < scanLayout.$.input.length) {
       partialSums[i] = operatorSlot.$(prev, scanLayout.$.input[baseIdx + i] as number);
-      prev = partialSums[i] as number;
+      prev = partialSums[i];
       lastIdx = i;
     }
   }
@@ -56,7 +57,7 @@ export const computeBlock = tgpu.computeFn({
 
     const scannedSum = workgroupMemory.$[localIdx];
 
-    for (let i = d.u32(0); i < 8; i++) {
+    for (const i of tgpu.unroll(std.range(ELEMENTS_PER_THREAD))) {
       if (baseIdx + i < scanLayout.$.input.length) {
         if (i === 0) {
           scanLayout.$.input[baseIdx + i] = scannedSum;

packages/typegpu-sort/src/scan/prefixScan.ts

@@ -8,6 +8,7 @@ import {
 } from 'typegpu';
 import type { BinaryOp } from './types.ts';
 import {
+  ELEMENTS_PER_THREAD,
   identitySlot,
   onlyGreatestElementSlot,
   operatorSlot,
@@ -83,7 +84,7 @@ export class PrefixScanComputer {
     querySet: TgpuQuerySet<'timestamp'> | null,
     isFirstPass: boolean,
   ): TgpuBuffer<d.WgslArray<d.F32>> & StorageFlag {
-    const numWorkgroups = Math.ceil(actualLength / (WORKGROUP_SIZE * 8));
+    const numWorkgroups = Math.ceil(actualLength / (WORKGROUP_SIZE * ELEMENTS_PER_THREAD));
     const scanPipeline = this.getScanPipeline(onlyGreatestElement);
 
     // Base case: single workgroup

packages/typegpu-sort/src/scan/schemas.ts

@@ -1,6 +1,7 @@
 import tgpu, { d } from 'typegpu';
 
 export const WORKGROUP_SIZE = 256;
+export const ELEMENTS_PER_THREAD = 8;
 
 export const scanLayout = tgpu.bindGroupLayout({
   input: { storage: d.arrayOf(d.f32), access: 'mutable' },

packages/typegpu/src/core/slot/accessor.ts

@@ -2,7 +2,7 @@ import { type AnyData, isData } from '../../data/dataTypes.ts';
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
 import { isSnippet, type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import type { BaseData } from '../../data/wgslTypes.ts';
-import { getResolutionCtx, inCodegenMode } from '../../execMode.ts';
+import { getResolutionCtx } from '../../execMode.ts';
 import { getName, hasTinyestMetadata, setName } from '../../shared/meta.ts';
 import type { InferGPU } from '../../shared/repr.ts';
 import {
@@ -174,7 +174,7 @@ export class TgpuAccessorImpl<T extends BaseData>
   }
 
   get $(): InferGPU<T> {
-    if (inCodegenMode()) {
+    if (getResolutionCtx()) {
       return this[$gpuValueOf];
     }
 
@@ -198,7 +198,7 @@ export class TgpuMutableAccessorImpl<T extends BaseData>
   }
 
   get $(): InferGPU<T> {
-    if (inCodegenMode()) {
+    if (getResolutionCtx()) {
       return this[$gpuValueOf];
     }