diff --git a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts b/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts deleted file mode 100644 index 737c1e69d7..0000000000 --- a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts +++ /dev/null @@ -1,75 +0,0 @@ -import { createPrefixScanComputer, prefixScan } from '@typegpu/sort'; -import type { TgpuRoot } from 'typegpu'; -import { d, std } from 'typegpu'; - -type SumResult = { - success: boolean; - jsTime: number; - gpuTime: number; - gpuShaderTime: number; -}; - -function prefixSumOnJS(arr: number[]) { - for (let i = 1; i < arr.length; i++) { - arr[i] += arr[i - 1]; - } - // In Blelloch scan, the result starts with identity element - arr.unshift(0); - arr.pop(); - return arr; -} - -function arraysEqual(a: number[], b: number[]): boolean { - if (a.length !== b.length) { - return false; - } - for (let i = 0; i < a.length; i++) { - if (a[i] !== b[i]) { - return false; - } - } - return true; -} - -export async function performCalculationsWithTime( - root: TgpuRoot, - inputArray: number[], -): Promise { - const arraySize = inputArray.length; - const inputBuffer = root.createBuffer(d.arrayOf(d.f32, arraySize)).$usage('storage'); - inputBuffer.write(inputArray); - - // JS version - const jsStartTime = performance.now(); - const jsResult = prefixSumOnJS(inputArray); - const jsTime = performance.now() - jsStartTime; - - // GPU version - createPrefixScanComputer(root, { operation: std.add, identityElement: 0 }); - const querySet = root.createQuerySet('timestamp', 2); - const gpuStartTime = performance.now(); - const calcResult = prefixScan( - root, - { - inputBuffer: inputBuffer, - outputBuffer: inputBuffer, - operation: std.add, - identityElement: 0, - }, - querySet, - ); - querySet.resolve(); - await root.device.queue.onSubmittedWorkDone(); - const gpuTime = performance.now() - gpuStartTime; - - const gpuResult = await calcResult.read(); - const timestamps = await querySet.read(); - const gpuShaderTime = Number(timestamps[1] - timestamps[0]) / 1_000_000; - - return { - success: arraysEqual(jsResult, gpuResult), - jsTime, - gpuTime, - gpuShaderTime, - }; -} diff --git a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.html b/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.html deleted file mode 100644 index 74f59eb01c..0000000000 --- a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.html +++ /dev/null @@ -1,317 +0,0 @@ -
-
- -
Time [ms]
- - -
- - - - - 0 -
- - -
- -
-
-
-
-
-
-
-
-
-
-
-
-
-
21K
-
- - -
-
-
-
-
-
-
-
-
-
-
-
-
-
131K
-
- - -
-
-
-
-
-
-
-
-
-
-
-
-
-
1M
-
- - -
-
-
-
-
-
-
-
-
-
-
-
-
-
4M
-
- - -
-
-
-
-
-
-
-
-
-
-
-
-
-
8M
-
-
-
-
- - diff --git a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.ts b/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.ts deleted file mode 100644 index a85e1aa83f..0000000000 --- a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.ts +++ /dev/null @@ -1,81 +0,0 @@ -import { tgpu } from 'typegpu'; -import { defineControls } from '../../common/defineControls.ts'; -import { performCalculationsWithTime } from './calculator.ts'; - -const SIZES = [21037, 131072, 1048576, 4194304, 8388608] as const; - -const root = await tgpu.init({ - device: { - requiredFeatures: ['timestamp-query'], - }, -}); - -const dataGroups = Array.from(document.querySelectorAll('.data-group')); -const yAxisLabels = Array.from(document.querySelectorAll('.y-axis-labels span')); - -const results = SIZES.map(() => ({ jsTime: 0, gpuTime: 0, gpuShaderTime: 0 })); - -function drawCharts() { - const overallMax = Math.max( - ...results.map((r) => Math.max(r.jsTime, r.gpuTime, r.gpuShaderTime)), - ); - - // Update y-axis - const ticks = - overallMax <= 0 ? [0, 0, 0, 0, 0] : Array.from({ length: 5 }, (_, i) => (i / 4) * overallMax); - for (const [i, label] of yAxisLabels.toReversed().entries()) { - label.textContent = ticks[i].toFixed(1); - } - - const metrics = [ - { cls: '.bar-js', key: 'jsTime', label: 'JS' }, - { cls: '.bar-gpu-total', key: 'gpuTime', label: 'Total GPU' }, - { cls: '.bar-gpu-shader', key: 'gpuShaderTime', label: 'GPU shader' }, - ] as const; - - for (const [i, group] of dataGroups.entries()) { - const r = results[i]; - - // Update speedup label - const speedup = r.gpuShaderTime > 0 ? (r.jsTime / r.gpuShaderTime).toFixed(1) : '-'; - (group.querySelector('.speedup-label') as HTMLDivElement).textContent = `${speedup}x`; - - // Update bars and tooltips - for (const m of metrics) { - const bar = group.querySelector(m.cls) as HTMLDivElement; - const value = r[m.key]; - const height = overallMax > 0 ? value / overallMax : 0; - bar.style.setProperty('--bar-height', `${height}`); - - const tooltip = bar.querySelector('.bar-tooltip') as HTMLDivElement; - tooltip.textContent = `${m.label}: ${value.toFixed(2)}ms`; - } - } -} - -async function runBenchmarks() { - for (const [i, size] of SIZES.entries()) { - const input = Array(size).fill(1); - const result = await performCalculationsWithTime(root, input); - if (result.success) { - results[i] = result; - } - } - drawCharts(); -} - -void runBenchmarks(); - -// #region Example controls & Cleanup - -export const controls = defineControls({ - Recalculate: { - onButtonClick: runBenchmarks, - }, -}); - -export function onCleanup() { - root.destroy(); -} - -// #endregion diff --git a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/meta.json b/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/meta.json deleted file mode 100644 index eb745e2c29..0000000000 --- a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/meta.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "title": "Concurrent Chart", - "category": "algorithms", - "tags": ["concurrent"], - "dev": true, - "coolFactor": 4 -} diff --git a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/thumbnail.png b/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/thumbnail.png deleted file mode 100644 index de4e7cb329..0000000000 Binary files a/apps/typegpu-docs/src/examples/algorithms/concurrent-chart/thumbnail.png and /dev/null differ diff --git a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.html b/apps/typegpu-docs/src/examples/algorithms/sort/index.html similarity index 100% rename from apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.html rename to apps/typegpu-docs/src/examples/algorithms/sort/index.html diff --git a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.ts b/apps/typegpu-docs/src/examples/algorithms/sort/index.ts similarity index 53% rename from apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.ts rename to apps/typegpu-docs/src/examples/algorithms/sort/index.ts index 459d3d296f..dc0ec5bacc 100644 --- a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.ts +++ b/apps/typegpu-docs/src/examples/algorithms/sort/index.ts @@ -1,9 +1,11 @@ -import { tgpu, d, std } from 'typegpu'; +import { tgpu, d, std, type TgpuQuerySet } from 'typegpu'; import { type BitonicSorter, type BitonicSorterOptions, createBitonicSorter, + createRadixSorter, decomposeWorkgroups, + type Sorter, } from '@typegpu/sort'; import { randf } from '@typegpu/noise'; import { fullScreenTriangle } from 'typegpu/common'; @@ -26,8 +28,9 @@ const root = await tgpu.init({ }, }, }); -const hasTimestampQuery = root.enabledFeatures.has('timestamp-query'); -const querySet = hasTimestampQuery ? root.createQuerySet('timestamp', 2) : null; +const querySet = root.enabledFeatures.has('timestamp-query') + ? root.createQuerySet('timestamp', 2) + : null; const canvas = document.querySelector('canvas') as HTMLCanvasElement; const context = root.configureContext({ canvas }); @@ -35,13 +38,14 @@ const context = root.configureContext({ canvas }); const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); const maxSide = Math.floor(Math.sqrt(maxBufferSize / 4)); -const minLog = 2; // log_2(4) +const minLog = 2; const maxLog = Math.floor(Math.log2(maxSide)); const arraySizeOptions = Array.from({ length: 8 }, (_, i) => { const side = Math.round(2 ** (minLog + (i * (maxLog - minLog)) / 7)); return side * side; }); +type AlgorithmKey = 'bitonic' | 'radix'; type SortOrderKey = 'ascending' | 'descending' | 'bit-reversed' | 'xor-scatter'; const sortOrders: Record = { @@ -68,6 +72,7 @@ const sortOrders: Record = { }; const state = { + algorithm: 'bitonic' as AlgorithmKey, arraySize: arraySizeOptions[2], sortOrder: 'ascending' as SortOrderKey, }; @@ -144,48 +149,56 @@ const initPipeline = root.createComputePipeline({ compute: initKernel }); let buffer = root.createBuffer(d.arrayOf(d.u32, state.arraySize)).$usage('storage'); -let bindGroup = root.createBindGroup(renderLayout, { - data: buffer, -}); -let initBindGroup = root.createBindGroup(initLayout, { - data: buffer, -}); +let bindGroup = root.createBindGroup(renderLayout, { data: buffer }); -function createSorters(buf: typeof buffer) { +function createBitonicSorters(buf: typeof buffer) { return Object.fromEntries( Object.entries(sortOrders).map(([key, opts]) => [key, createBitonicSorter(root, buf, opts)]), ) as Record; } -let sorters = createSorters(buffer); +function createRadixSorters(buf: typeof buffer) { + return { + ascending: createRadixSorter(root, buf), + descending: createRadixSorter(root, buf, { direction: 'descending' }), + }; +} -function recreateBuffer() { - for (const s of Object.values(sorters)) { +let bitonicSorters = createBitonicSorters(buffer); +let radixSorters = createRadixSorters(buffer); + +function destroySorters() { + for (const s of Object.values(bitonicSorters)) { + s.destroy(); + } + for (const s of Object.values(radixSorters)) { s.destroy(); } +} + +function recreateBuffer() { + destroySorters(); buffer.destroy(); buffer = root.createBuffer(d.arrayOf(d.u32, state.arraySize)).$usage('storage'); + bindGroup = root.createBindGroup(renderLayout, { data: buffer }); - bindGroup = root.createBindGroup(renderLayout, { - data: buffer, - }); - - initBindGroup = root.createBindGroup(initLayout, { - data: buffer, - }); - - sorters = createSorters(buffer); + bitonicSorters = createBitonicSorters(buffer); + radixSorters = createRadixSorters(buffer); } -function generateRandomArray() { - const workgroupsTotal = Math.ceil(state.arraySize / WORKGROUP_SIZE); +function fillRandom(buf: typeof buffer, size: number) { + const workgroupsTotal = Math.ceil(size / WORKGROUP_SIZE); const [workgroupsX, workgroupsY, workgroupsZ] = decomposeWorkgroups(workgroupsTotal); initSeed.write(Math.random()); + initPipeline + .with(root.createBindGroup(initLayout, { data: buf })) + .dispatchWorkgroups(workgroupsX, workgroupsY, workgroupsZ); +} - initPipeline.with(initBindGroup).dispatchWorkgroups(workgroupsX, workgroupsY, workgroupsZ); - +function generateRandomArray() { + fillRandom(buffer, state.arraySize); render(); } @@ -221,36 +234,150 @@ function hideOverlay(delayMs = 1500) { }, delayMs); } +function pickSorter(): { sorter: Sorter; note: string } { + if (state.algorithm === 'radix') { + if (state.sortOrder === 'ascending' || state.sortOrder === 'descending') { + return { sorter: radixSorters[state.sortOrder], note: '' }; + } + return { + sorter: bitonicSorters[state.sortOrder], + note: ' (custom orders need the bitonic sorter)', + }; + } + return { sorter: bitonicSorters[state.sortOrder], note: '' }; +} + +async function runSorterTimed(sorter: Sorter): Promise { + if (!querySet?.available) { + sorter.run(); + return null; + } + + const encoder = root.device.createCommandEncoder(); + const pass = encoder.beginComputePass({ + timestampWrites: { + querySet: querySet.querySet, + beginningOfPassWriteIndex: 0, + endOfPassWriteIndex: 1, + }, + }); + sorter.run({ pass }); + pass.end(); + root.device.queue.submit([encoder.finish()]); + + querySet.resolve(); + const [start, end] = await querySet.read(); + return Number(end - start) / 1_000_000; +} + async function sort() { - const sorter = sorters[state.sortOrder]; + const { sorter, note } = pickSorter(); showOverlay('Sorting...'); - sorter.run({ querySet: querySet ?? undefined }); - - let gpuTimeMs: number | null = null; - if (querySet?.available) { - querySet.resolve(); - const timestamps = await querySet.read(); - gpuTimeMs = Number(timestamps[1] - timestamps[0]) / 1_000_000; - } + const gpuTimeMs = await runSorterTimed(sorter); render(); - const timeStr = - gpuTimeMs !== null - ? ` in ${ - gpuTimeMs >= 1000 ? `${(gpuTimeMs / 1000).toFixed(2)}s` : `${gpuTimeMs.toFixed(2)}ms` - }` - : ''; - showOverlay(`\u2714 Sorted${timeStr}`, false); + const timeStr = gpuTimeMs !== null ? ` in ${formatMs(gpuTimeMs)}` : ''; + showOverlay(`✔ Sorted${timeStr}${note}`, false); hideOverlay(); } +// #region Benchmark + +const BENCH_WARMUP = 3; +const BENCH_RUNS = 10; + +function formatMs(milliseconds: number): string { + return milliseconds >= 1000 + ? `${(milliseconds / 1000).toFixed(2)}s` + : `${milliseconds.toFixed(2)}ms`; +} + +async function benchmarkSorter( + sorter: Sorter, + timestamps: TgpuQuerySet<'timestamp'>, +): Promise { + for (let i = 0; i < BENCH_WARMUP; i++) { + sorter.run(); + } + await root.device.queue.onSubmittedWorkDone(); + + let total = 0; + for (let i = 0; i < BENCH_RUNS; i++) { + const encoder = root.device.createCommandEncoder(); + const pass = encoder.beginComputePass({ + timestampWrites: { + querySet: timestamps.querySet, + beginningOfPassWriteIndex: 0, + endOfPassWriteIndex: 1, + }, + }); + sorter.run({ pass }); + pass.end(); + root.device.queue.submit([encoder.finish()]); + + timestamps.resolve(); + const [start, end] = await timestamps.read(); + total += Number(end - start) / 1_000_000; + } + return total / BENCH_RUNS; +} + +async function runBenchmark() { + if (!querySet) { + showOverlay('Benchmark requires timestamp-query', false); + hideOverlay(); + return; + } + + const sizes = [2 ** 12, 2 ** 16, 2 ** 20, 2 ** 22, 2 ** 24].filter( + (size) => size * 4 <= maxBufferSize, + ); + + console.log(`=== Sort benchmark (avg GPU time, ${BENCH_RUNS} runs) ===`); + for (const size of sizes) { + showOverlay(`Benchmarking ${size.toLocaleString()} keys...`); + + const benchBuffer = root.createBuffer(d.arrayOf(d.u32, size)).$usage('storage'); + fillRandom(benchBuffer, size); + + const bitonic = createBitonicSorter(root, benchBuffer); + const bitonicMs = await benchmarkSorter(bitonic, querySet); + bitonic.destroy(); + + fillRandom(benchBuffer, size); + const radix = createRadixSorter(root, benchBuffer); + const radixMs = await benchmarkSorter(radix, querySet); + radix.destroy(); + + benchBuffer.destroy(); + + console.log( + ` ${size.toLocaleString().padStart(12)} keys: bitonic ${formatMs(bitonicMs)}, radix ${formatMs(radixMs)}`, + ); + } + console.log('==============================================='); + + showOverlay('✔ Benchmark complete (see console)', false); + hideOverlay(3000); +} + +// #endregion + // #region Example controls & Cleanup +const algorithmKeys: AlgorithmKey[] = ['bitonic', 'radix']; const sortOrderKeys = Object.keys(sortOrders) as SortOrderKey[]; export const controls = defineControls({ + Algorithm: { + initial: 'bitonic' as AlgorithmKey, + options: algorithmKeys, + onSelectChange: (value) => { + state.algorithm = value; + }, + }, 'Array Size': { initial: arraySizeOptions[2], options: arraySizeOptions, @@ -261,7 +388,7 @@ export const controls = defineControls({ }, }, 'Sort Order': { - initial: 'ascending', + initial: 'ascending' as SortOrderKey, options: sortOrderKeys, onSelectChange: (value) => { state.sortOrder = value; @@ -269,12 +396,12 @@ export const controls = defineControls({ }, Reshuffle: { onButtonClick: generateRandomArray }, Sort: { onButtonClick: sort }, + Benchmark: { onButtonClick: runBenchmark }, }); export function onCleanup() { - for (const s of Object.values(sorters)) { - s.destroy(); - } + destroySorters(); + querySet?.destroy(); root.destroy(); } diff --git a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/meta.json b/apps/typegpu-docs/src/examples/algorithms/sort/meta.json similarity index 76% rename from apps/typegpu-docs/src/examples/algorithms/bitonic-sort/meta.json rename to apps/typegpu-docs/src/examples/algorithms/sort/meta.json index 5c3697752f..534e32771b 100644 --- a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/meta.json +++ b/apps/typegpu-docs/src/examples/algorithms/sort/meta.json @@ -1,5 +1,5 @@ { - "title": "Bitonic Sort", + "title": "Sort", "category": "algorithms", "tags": ["compute"], "dev": true, diff --git a/apps/typegpu-docs/src/examples/algorithms/bitonic-sort/thumbnail.png b/apps/typegpu-docs/src/examples/algorithms/sort/thumbnail.png similarity index 100% rename from apps/typegpu-docs/src/examples/algorithms/bitonic-sort/thumbnail.png rename to apps/typegpu-docs/src/examples/algorithms/sort/thumbnail.png diff --git a/apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts b/apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts index 9049fd9982..86364d08ea 100644 --- a/apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts +++ b/apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts @@ -1,6 +1,6 @@ import { tgpu } from 'typegpu'; import * as d from 'typegpu/data'; -import { type BinaryOp, prefixScan, scan } from '@typegpu/sort'; +import { type BinaryOp, prefixScan, reduce } from '@typegpu/sort'; import * as std from 'typegpu/std'; import { addFn, concat10, isArrayEqual, mulFn, prefixScanJS, scanJS } from './functions.ts'; @@ -12,7 +12,7 @@ async function runAndCompare(arr: number[], op: BinaryOp, scanOnly: boolean) { const input = root.createBuffer(d.arrayOf(d.f32, arr.length), arr).$usage('storage'); const output = scanOnly - ? scan(root, { + ? reduce(root, { inputBuffer: input, operation: op.operation, identityElement: op.identityElement, @@ -82,7 +82,7 @@ async function testLength16777217(): Promise { async function testDoesNotDestroyBuffer(): Promise { const input = root.createBuffer(d.arrayOf(d.f32, 8), [1, 2, 3, 4, 5, 6, 7, 8]).$usage('storage'); - scan(root, { + reduce(root, { inputBuffer: input, operation: addFn, identityElement: 0, @@ -96,7 +96,7 @@ async function testDoesNotCacheBuffers(): Promise { const input1 = root.createBuffer(d.arrayOf(d.f32, 8), [1, 2, 3, 4, 5, 6, 7, 8]).$usage('storage'); - const output1 = scan(root, { + const output1 = reduce(root, { inputBuffer: input1, operation: op.operation, identityElement: op.identityElement, @@ -109,7 +109,7 @@ async function testDoesNotCacheBuffers(): Promise { ) .$usage('storage'); - const output2 = scan(root, { + const output2 = reduce(root, { inputBuffer: input2, operation: op.operation, identityElement: op.identityElement, diff --git a/apps/typegpu-docs/tests/individual-example-tests/bitonic-sort.test.ts b/apps/typegpu-docs/tests/individual-example-tests/bitonic-sort.test.ts deleted file mode 100644 index b8d6405c40..0000000000 --- a/apps/typegpu-docs/tests/individual-example-tests/bitonic-sort.test.ts +++ /dev/null @@ -1,144 +0,0 @@ -/** - * @vitest-environment jsdom - */ - -import { describe, expect } from 'vitest'; -import { it } from 'typegpu-testing-utility'; -import { runExampleTest, setupCommonMocks } from './utils/baseTest.ts'; - -describe('bitonic sort example', () => { - setupCommonMocks(); - - it('should produce valid code', async ({ device }) => { - const shaderCodes = await runExampleTest( - { - category: 'algorithms', - name: 'bitonic-sort', - controlTriggers: ['Sort'], - expectedCalls: 4, - }, - device, - ); - - expect(shaderCodes).toMatchInlineSnapshot(` - "struct copyParamsType { - srcLength: u32, - dstLength: u32, - paddingValue: u32, - } - - @group(0) @binding(2) var params: copyParamsType; - - @group(0) @binding(1) var dst: array; - - @group(0) @binding(0) var src: array; - - @compute @workgroup_size(256) fn copyPadKernel(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { - let spanX = (numWorkgroups.x * 256u); - let spanY = (numWorkgroups.y * spanX); - let idx = ((gid.x + (gid.y * spanX)) + (gid.z * spanY)); - let dstLength = params.dstLength; - let srcLength = params.srcLength; - if ((idx >= dstLength)) { - return; - } - dst[idx] = select(params.paddingValue, src[idx], (idx < srcLength)); - } - - struct sortUniformsType { - k: u32, - jShift: u32, - } - - @group(0) @binding(1) var uniforms: sortUniformsType; - - @group(0) @binding(0) var data: array; - - fn defaultCompare(a: u32, b: u32) -> bool { - return (a < b); - } - - @compute @workgroup_size(256) fn bitonicStepKernel(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { - let spanX = (numWorkgroups.x * 256u); - let spanY = (numWorkgroups.y * spanX); - let tid = ((gid.x + (gid.y * spanX)) + (gid.z * spanY)); - let k = uniforms.k; - let shift = uniforms.jShift; - let dataLength = arrayLength(&data); - let stride = (1u << shift); - let maskBelow = (stride - 1u); - let below = (tid & maskBelow); - let above = (tid >> shift); - let i = (below + (above * (stride << 1u))); - let ixj = (i + stride); - if ((ixj >= dataLength)) { - return; - } - let ascending = ((i & k) == 0u); - let left = data[i]; - let right = data[ixj]; - let leftFirst = defaultCompare(left, right); - let shouldSwap = select(leftFirst, !leftFirst, ascending); - if (shouldSwap) { - data[i] = right; - data[ixj] = left; - } - } - - struct copyParamsType { - srcLength: u32, - dstLength: u32, - paddingValue: u32, - } - - @group(0) @binding(2) var params: copyParamsType; - - @group(0) @binding(1) var dst: array; - - @group(0) @binding(0) var src: array; - - @compute @workgroup_size(256) fn copyBackKernel(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { - let spanX = (numWorkgroups.x * 256u); - let spanY = (numWorkgroups.y * spanX); - let idx = ((gid.x + (gid.y * spanX)) + (gid.z * spanY)); - if ((idx < params.srcLength)) { - dst[idx] = src[idx]; - } - } - - struct fullScreenTriangle_Output { - @builtin(position) pos: vec4f, - @location(0) uv: vec2f, - } - - @vertex fn fullScreenTriangle(@builtin(vertex_index) vertexIndex: u32) -> fullScreenTriangle_Output { - const pos = array(vec2f(-1, -1), vec2f(3, -1), vec2f(-1, 3)); - const uv = array(vec2f(0, 1), vec2f(2, 1), vec2f(0, -1)); - - return fullScreenTriangle_Output(vec4f(pos[vertexIndex], 0, 1), uv[vertexIndex]); - } - - @group(0) @binding(0) var data_1: array; - - struct fragmentFn_Input { - @location(0) uv: vec2f, - } - - @fragment fn fragmentFn(_arg_0: fragmentFn_Input) -> @location(0) vec4f { - let data = (&data_1); - let arrayLength_1 = arrayLength(&(*data)); - let cols = u32(round(sqrt(f32(arrayLength_1)))); - let rows = u32(round((f32(arrayLength_1) / f32(cols)))); - let col = u32(floor((_arg_0.uv.x * f32(cols)))); - let row = u32(floor((_arg_0.uv.y * f32(rows)))); - let idx = ((row * cols) + col); - if ((idx >= arrayLength_1)) { - return vec4f(0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 1); - } - let value = (*data)[idx]; - let normalized = (f32(value) / 255f); - return vec4f(normalized, normalized, normalized, 1f); - }" - `); - }); -}); diff --git a/apps/typegpu-docs/tests/individual-example-tests/sort.test.ts b/apps/typegpu-docs/tests/individual-example-tests/sort.test.ts new file mode 100644 index 0000000000..3c9988ea30 --- /dev/null +++ b/apps/typegpu-docs/tests/individual-example-tests/sort.test.ts @@ -0,0 +1,264 @@ +/** + * @vitest-environment jsdom + */ + +import { describe, expect } from 'vitest'; +import { it } from 'typegpu-testing-utility'; +import { runExampleTest, setupCommonMocks } from './utils/baseTest.ts'; + +describe('sort example', () => { + setupCommonMocks(); + + it('should produce valid code', async ({ device }) => { + const shaderCodes = await runExampleTest( + { + category: 'algorithms', + name: 'sort', + controlTriggers: ['Sort'], + expectedCalls: 6, + }, + device, + ); + + expect(shaderCodes).toMatchInlineSnapshot(` + "fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(1) var dst: array; + + @group(0) @binding(0) var src: array; + + @group(0) @binding(2) var padding: u32; + + @compute @workgroup_size(256) fn pad(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let idx = ((flatWorkgroupIndex(wid, numWorkgroups) * 256u) + lid.x); + if ((idx >= 1024u)) { + return; + } + if ((idx < 841u)) { + dst[idx] = src[idx]; + } + else { + dst[idx] = padding; + } + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(0) var data: array; + + var localKeys: array; + + fn loadShared(base: u32, tid: u32) { + localKeys[tid] = data[(base + tid)]; + localKeys[(tid + 256u)] = data[((base + tid) + 256u)]; + } + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn swapLocalAt(a: u32, b: u32, left: u32, right: u32) { + localKeys[a] = right; + localKeys[b] = left; + } + + fn exchangeLocal(base: u32, iLocal: u32, stride: u32, k: u32) { + let jLocal = (iLocal + stride); + let left = localKeys[iLocal]; + let right = localKeys[jLocal]; + let ascending = (((base + iLocal) & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapLocalAt(iLocal, jLocal, left, right); + } + } + + fn mergeDown(base: u32, tid: u32, startShift: u32, k: u32) { + for (var jShift = startShift; (jShift > 0u); jShift--) { + workgroupBarrier(); + let stride = (1u << (jShift - 1u)); + let below = (tid & (stride - 1u)); + let above = (tid >> (jShift - 1u)); + exchangeLocal(base, (below + (above * (stride << 1u))), stride, k); + } + } + + fn storeShared(base: u32, tid: u32) { + data[(base + tid)] = localKeys[tid]; + data[((base + tid) + 256u)] = localKeys[(tid + 256u)]; + } + + @compute @workgroup_size(256) fn localSort(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let base = (flatWorkgroupIndex(wid, numWorkgroups) * 512u); + if ((base >= arrayLength(&data))) { + return; + } + loadShared(base, lid.x); + for (var kShift = 1u; (kShift <= 9u); kShift++) { + mergeDown(base, lid.x, kShift, (1u << kShift)); + } + workgroupBarrier(); + storeShared(base, lid.x); + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + struct sortUniformsType { + k: u32, + jShift: u32, + } + + @group(0) @binding(1) var uniforms: sortUniformsType; + + @group(0) @binding(0) var data: array; + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn swapAt(i: u32, j: u32, left: u32, right: u32) { + data[i] = right; + data[j] = left; + } + + @compute @workgroup_size(256) fn item(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let tid = ((flatWorkgroupIndex(wid, numWorkgroups) * 256u) + lid.x); + let k = uniforms.k; + let shift = uniforms.jShift; + let stride = (1u << shift); + let below = (tid & (stride - 1u)); + let above = (tid >> shift); + let i = (below + (above * (stride << 1u))); + let ixj = (i + stride); + if ((ixj >= arrayLength(&data))) { + return; + } + let left = data[i]; + let right = data[ixj]; + let ascending = ((i & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapAt(i, ixj, left, right); + } + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(0) var data: array; + + var localKeys: array; + + fn loadShared(base: u32, tid: u32) { + localKeys[tid] = data[(base + tid)]; + localKeys[(tid + 256u)] = data[((base + tid) + 256u)]; + } + + struct sortUniformsType { + k: u32, + jShift: u32, + } + + @group(0) @binding(1) var uniforms: sortUniformsType; + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn swapLocalAt(a: u32, b: u32, left: u32, right: u32) { + localKeys[a] = right; + localKeys[b] = left; + } + + fn exchangeLocal(base: u32, iLocal: u32, stride: u32, k: u32) { + let jLocal = (iLocal + stride); + let left = localKeys[iLocal]; + let right = localKeys[jLocal]; + let ascending = (((base + iLocal) & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapLocalAt(iLocal, jLocal, left, right); + } + } + + fn mergeDown(base: u32, tid: u32, startShift: u32, k: u32) { + for (var jShift = startShift; (jShift > 0u); jShift--) { + workgroupBarrier(); + let stride = (1u << (jShift - 1u)); + let below = (tid & (stride - 1u)); + let above = (tid >> (jShift - 1u)); + exchangeLocal(base, (below + (above * (stride << 1u))), stride, k); + } + } + + fn storeShared(base: u32, tid: u32) { + data[(base + tid)] = localKeys[tid]; + data[((base + tid) + 256u)] = localKeys[(tid + 256u)]; + } + + @compute @workgroup_size(256) fn localMerge(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let base = (flatWorkgroupIndex(wid, numWorkgroups) * 512u); + if ((base >= arrayLength(&data))) { + return; + } + loadShared(base, lid.x); + mergeDown(base, lid.x, 9u, uniforms.k); + workgroupBarrier(); + storeShared(base, lid.x); + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(1) var dst: array; + + @group(0) @binding(0) var src: array; + + @compute @workgroup_size(256) fn unpad(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let idx = ((flatWorkgroupIndex(wid, numWorkgroups) * 256u) + lid.x); + if ((idx < 841u)) { + dst[idx] = src[idx]; + } + } + + struct fullScreenTriangle_Output { + @builtin(position) pos: vec4f, + @location(0) uv: vec2f, + } + + @vertex fn fullScreenTriangle(@builtin(vertex_index) vertexIndex: u32) -> fullScreenTriangle_Output { + const pos = array(vec2f(-1, -1), vec2f(3, -1), vec2f(-1, 3)); + const uv = array(vec2f(0, 1), vec2f(2, 1), vec2f(0, -1)); + + return fullScreenTriangle_Output(vec4f(pos[vertexIndex], 0, 1), uv[vertexIndex]); + } + + @group(0) @binding(0) var data_1: array; + + struct fragmentFn_Input { + @location(0) uv: vec2f, + } + + @fragment fn fragmentFn(_arg_0: fragmentFn_Input) -> @location(0) vec4f { + let data = (&data_1); + let arrayLength_1 = arrayLength(&(*data)); + let cols = u32(round(sqrt(f32(arrayLength_1)))); + let rows = u32(round((f32(arrayLength_1) / f32(cols)))); + let col = u32(floor((_arg_0.uv.x * f32(cols)))); + let row = u32(floor((_arg_0.uv.y * f32(rows)))); + let idx = ((row * cols) + col); + if ((idx >= arrayLength_1)) { + return vec4f(0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 1); + } + let value = (*data)[idx]; + let normalized = (f32(value) / 255f); + return vec4f(normalized, normalized, normalized, 1f); + }" + `); + }); +}); diff --git a/packages/typegpu-sort/README.md b/packages/typegpu-sort/README.md index 217b269409..acee954c20 100644 --- a/packages/typegpu-sort/README.md +++ b/packages/typegpu-sort/README.md @@ -4,54 +4,131 @@ -GPU sorting and scanning algorithms for TypeGPU. +GPU sorting and scanning algorithms for TypeGPU. Sorts and scans `u32`, `i32` and +`f32` storage buffers, optionally reordering a payload buffer alongside the keys, +and composes with your own command encoders and compute passes. -## Bitonic Sort +## Radix Sort -Sorts a `u32` storage buffer in-place. Arrays with non-power-of-2 lengths are padded automatically. +A stable LSD radix sort, and the fastest option here. Keys are ordered by the +natural order of their type. ```ts import { tgpu, d } from 'typegpu'; -import { createBitonicSorter } from '@typegpu/sort'; +import { createRadixSorter } from '@typegpu/sort'; const root = await tgpu.init(); -const buffer = root.createBuffer(d.arrayOf(d.u32, 1024), data).$usage('storage'); +const keys = root.createBuffer(d.arrayOf(d.f32, 100_000), data).$usage('storage'); -const sorter = createBitonicSorter(root, buffer); +const sorter = createRadixSorter(root, keys); sorter.run(); sorter.destroy(); ``` -Custom comparator (descending): +With a payload (e.g. sorting indices by distance) and descending order: ```ts -const sorter = createBitonicSorter(root, buffer, { - compare: (a, b) => { 'use gpu'; return a > b; }, - paddingValue: 0, // must sort to the end — use 0 for descending +const indices = root.createBuffer(d.arrayOf(d.u32, 100_000), idx).$usage('storage'); + +const sorter = createRadixSorter(root, keys, { + direction: 'descending', + values: indices, }); ``` -With GPU timing (`timestamp-query` feature required): +All GPU resources are created once in `createRadixSorter`, so `run()` only +records dispatches and is cheap to call every frame. + +For `f32` keys sorted ascending, NaNs with a cleared sign bit sort after ++Infinity and NaNs with a set sign bit sort before -Infinity. + +## Bitonic Sort + +Sorts with an arbitrary comparator, which radix sort cannot do. Slower than +radix sort. Arrays with non-power-of-2 lengths are padded automatically. + +```ts +import { createBitonicSorter } from '@typegpu/sort'; + +const sorter = createBitonicSorter(root, keys); +sorter.run(); +``` + +Custom comparator (descending): ```ts -const querySet = root.createQuerySet('timestamp', 2); -sorter.run({ querySet }); -querySet.resolve(); -const [start, end] = await querySet.read(); -const gpuTimeMs = Number(end - start) / 1_000_000; +const sorter = createBitonicSorter(root, keys, { + compare: (a, b) => { 'use gpu'; return a > b; }, + paddingValue: 0, // must sort to the end, so the minimum value for descending +}); ``` +The bitonic sorter also accepts a `values` payload buffer for power-of-two input +sizes, swapped alongside the keys. Use radix sort for arbitrary-length numeric +key/payload pairs. + ## Prefix Scan +An exclusive work-efficient prefix scan over `f32` (default), `u32` or `i32` +buffers, with any associative operation. + ```ts -import { prefixScan, scan } from '@typegpu/sort'; +import { prefixScan, reduce } from '@typegpu/sort'; import * as std from 'typegpu/std'; -// Full prefix scan +// Full prefix scan (in place) const result = prefixScan(root, { inputBuffer, operation: std.add, identityElement: 0 }); -// Reduction only (returns the final aggregated value) -const total = scan(root, { inputBuffer, operation: std.add, identityElement: 0 }); +// Reduction only (returns a single-element buffer with the aggregate) +const total = reduce(root, { inputBuffer, operation: std.add, identityElement: 0 }); + +// Integer scan +const sums = prefixScan(root, { + inputBuffer: u32Buffer, + operation: std.add, + identityElement: 0, +}); +``` + +For repeated scans of the same buffer, prepare a plan once. All scratch buffers +and bind groups are allocated up front and `run()` only records dispatches: + +```ts +import { createPrefixScanComputer } from '@typegpu/sort'; + +const computer = createPrefixScanComputer(root, { + operation: std.add, + identityElement: 0, + dataType: d.u32, +}); +const plan = computer.prepare(inputBuffer); + +plan.run(); // any number of times +plan.destroy(); +``` + +`computer.scan(buffer)` and `computer.reduce(buffer)` do the same through a plan +cached per buffer, which is what `prefixScan` and `reduce` use internally. + +Note: passing `-2147483648` (i32 minimum) as `identityElement` currently +generates WGSL that does not compile. Use `-2147483647` instead. + +## Composing with your own passes + +Every `run()`, on sorters and scan plans alike, accepts an `encoder` or `pass` +to record the work into instead of submitting on its own: + +```ts +const encoder = root.device.createCommandEncoder(); + +sorter.run({ encoder }); // records, does not submit +// ... encode more work ... +root.device.queue.submit([encoder.finish()]); + +// or straight into an open compute pass: +const pass = encoder.beginComputePass(); +sorter.run({ pass }); +pass.end(); ``` ## TypeGPU is created by Software Mansion diff --git a/packages/typegpu-sort/package.json b/packages/typegpu-sort/package.json index a95844dc0e..ed7c6d29fa 100644 --- a/packages/typegpu-sort/package.json +++ b/packages/typegpu-sort/package.json @@ -41,6 +41,7 @@ "@typegpu/tgpu-dev-cli": "workspace:*", "@webgpu/types": "catalog:types", "typegpu": "workspace:*", + "typegpu-testing-utility": "workspace:*", "typescript": "catalog:types", "unbuild": "catalog:build", "unplugin-typegpu": "workspace:*" diff --git a/packages/typegpu-sort/src/bitonic/bitonicSort.ts b/packages/typegpu-sort/src/bitonic/bitonicSort.ts index a4bba2eade..96aa93581a 100644 --- a/packages/typegpu-sort/src/bitonic/bitonicSort.ts +++ b/packages/typegpu-sort/src/bitonic/bitonicSort.ts @@ -3,291 +3,386 @@ import { d, std, type StorageFlag, - type TgpuBindGroup, type TgpuBuffer, + type TgpuComputeFn, + type TgpuComputePipeline, type TgpuRoot, - type UniformFlag, } from 'typegpu'; -import { compareSlot, defaultCompare } from './slots.ts'; -import type { BitonicSorter, BitonicSorterOptions, BitonicSorterRunOptions } from './types.ts'; -import { decomposeWorkgroups, nextPowerOf2 } from './utils.ts'; +import { decomposeWorkgroups, dispatchIn, flatWorkgroupIndex } from '../dispatch.ts'; +import { beginRunPass, bindPass } from '../runPass.ts'; +import type { RunOptions } from '../types.ts'; +import { compareSlot, defaultCompare, defaultPaddingValues } from './slots.ts'; +import type { BitonicSorter, BitonicSorterOptions } from './types.ts'; const WORKGROUP_SIZE = 256; +const LOCAL_BLOCK = WORKGROUP_SIZE * 2; +const LOCAL_BLOCK_LOG2 = Math.log2(LOCAL_BLOCK); -const copyParamsType = d.struct({ - srcLength: d.u32, - dstLength: d.u32, - paddingValue: d.u32, -}); +export type BitonicKeyType = d.U32 | d.I32 | d.F32; + +type KeyBuffer = TgpuBuffer> & StorageFlag; +type ValueBuffer = TgpuBuffer> & StorageFlag; const sortUniformsType = d.struct({ k: d.u32, jShift: d.u32, }); -const sortLayout = tgpu.bindGroupLayout({ - data: { - storage: d.arrayOf(d.u32), - access: 'mutable', - }, - uniforms: { - uniform: sortUniformsType, - }, -}); - -const copyLayout = tgpu.bindGroupLayout({ - src: { - storage: d.arrayOf(d.u32), - access: 'readonly', - }, - dst: { - storage: d.arrayOf(d.u32), - access: 'mutable', - }, - params: { - uniform: copyParamsType, - }, -}); +function nextPowerOf2(n: number): number { + let p = 1; + while (p < n) { + p <<= 1; + } + return p; +} -const copyPadKernel = tgpu.computeFn({ - workgroupSize: [WORKGROUP_SIZE], - in: { - gid: d.builtin.globalInvocationId, - numWorkgroups: d.builtin.numWorkgroups, - }, -})((input) => { - const spanX = input.numWorkgroups.x * WORKGROUP_SIZE; - const spanY = input.numWorkgroups.y * spanX; +function makeBitonicSchemas(keyType: BitonicKeyType, valueType: d.AnyWgslData | undefined) { + const sortLayout = tgpu.bindGroupLayout({ + data: { storage: d.arrayOf(keyType), access: 'mutable' }, + uniforms: { uniform: sortUniformsType }, + }); - const idx = input.gid.x + input.gid.y * spanX + input.gid.z * spanY; + const hasPayload = valueType !== undefined; + const payloadType = valueType ?? d.u32; - const dstLength = copyLayout.$.params.dstLength; - const srcLength = copyLayout.$.params.srcLength; + const valsLayout = tgpu.bindGroupLayout({ + vals: { storage: d.arrayOf(payloadType), access: 'mutable' }, + }); - if (idx >= dstLength) { - return; + function swapAt(i: number, j: number, left: number, right: number) { + 'use gpu'; + sortLayout.$.data[i] = right; + sortLayout.$.data[j] = left; + if (hasPayload) { + const tmp = std.copy(valsLayout.$.vals[i] as number); + (valsLayout.$.vals[i] as number) = std.copy(valsLayout.$.vals[j] as number); + (valsLayout.$.vals[j] as number) = std.copy(tmp); + } } - copyLayout.$.dst[idx] = std.select( - copyLayout.$.params.paddingValue, - copyLayout.$.src[idx] as number, - idx < srcLength, - ); -}); + return { keyType, valueType, hasPayload, payloadType, sortLayout, valsLayout, swapAt }; +} -const copyBackKernel = tgpu.computeFn({ - workgroupSize: [WORKGROUP_SIZE], - in: { - gid: d.builtin.globalInvocationId, - numWorkgroups: d.builtin.numWorkgroups, - }, -})((input) => { - const spanX = input.numWorkgroups.x * WORKGROUP_SIZE; - const spanY = input.numWorkgroups.y * spanX; +type BitonicSchemas = ReturnType; - const idx = input.gid.x + input.gid.y * spanX + input.gid.z * spanY; +function makePaddingKernels(keyType: BitonicKeyType, size: number, paddedSize: number) { + const copyLayout = tgpu.bindGroupLayout({ + src: { storage: d.arrayOf(keyType), access: 'readonly' }, + dst: { storage: d.arrayOf(keyType), access: 'mutable' }, + padding: { uniform: keyType }, + }); - if (idx < copyLayout.$.params.srcLength) { - copyLayout.$.dst[idx] = copyLayout.$.src[idx] as number; - } -}); + const pad = tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })(({ + lid, + wid, + numWorkgroups, + }) => { + const idx = flatWorkgroupIndex(wid, numWorkgroups) * WORKGROUP_SIZE + lid.x; + if (idx >= paddedSize) { + return; + } -const bitonicStepKernel = tgpu.computeFn({ - workgroupSize: [WORKGROUP_SIZE], - in: { - gid: d.builtin.globalInvocationId, - numWorkgroups: d.builtin.numWorkgroups, - }, -})((input) => { - const spanX = input.numWorkgroups.x * WORKGROUP_SIZE; - const spanY = input.numWorkgroups.y * spanX; + if (idx < size) { + copyLayout.$.dst[idx] = copyLayout.$.src[idx] as number; + } else { + copyLayout.$.dst[idx] = copyLayout.$.padding; + } + }); - const tid = input.gid.x + input.gid.y * spanX + input.gid.z * spanY; + const unpad = tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })(({ + lid, + wid, + numWorkgroups, + }) => { + const idx = flatWorkgroupIndex(wid, numWorkgroups) * WORKGROUP_SIZE + lid.x; + if (idx < size) { + (copyLayout.$.dst[idx] as number) = copyLayout.$.src[idx] as number; + } + }); - const k = sortLayout.$.uniforms.k; - const shift = sortLayout.$.uniforms.jShift; - const dataLength = d.u32(sortLayout.$.data.length); - const stride = d.u32(1) << shift; + return { copyLayout, pad, unpad }; +} - const maskBelow = stride - 1; - const below = tid & maskBelow; - const above = tid >> shift; +function makeGlobalStepKernel(schemas: BitonicSchemas) { + const { sortLayout, swapAt } = schemas; - const i = below + above * (stride << 1); - const ixj = i + stride; + return tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })( + ({ lid, wid, numWorkgroups }) => { + const tid = flatWorkgroupIndex(wid, numWorkgroups) * WORKGROUP_SIZE + lid.x; - if (ixj >= dataLength) { - return; - } + const k = sortLayout.$.uniforms.k; + const shift = sortLayout.$.uniforms.jShift; + const stride = d.u32(1) << shift; - const ascending = (i & k) === 0; - const left = sortLayout.$.data[i] as number; - const right = sortLayout.$.data[ixj] as number; + const below = tid & (stride - 1); + const above = tid >> shift; + const i = below + above * (stride << 1); + const ixj = i + stride; - const leftFirst = compareSlot.$(left, right); - const shouldSwap = std.select(leftFirst, !leftFirst, ascending); + if (ixj >= d.u32(sortLayout.$.data.length)) { + return; + } - if (shouldSwap) { - sortLayout.$.data[i] = right; - sortLayout.$.data[ixj] = left; + const left = sortLayout.$.data[i] as number; + const right = sortLayout.$.data[ixj] as number; + const ascending = (i & k) === 0; + + if (std.select(compareSlot.$(left, right), compareSlot.$(right, left), ascending)) { + swapAt(i, ixj, left, right); + } + }, + ); +} + +function makeLocalKernels(schemas: BitonicSchemas) { + const { keyType, hasPayload, payloadType, sortLayout, valsLayout } = schemas; + + const localKeys = tgpu.workgroupVar(d.arrayOf(keyType, LOCAL_BLOCK)); + const localVals = tgpu.workgroupVar(d.arrayOf(payloadType, LOCAL_BLOCK)); + + function loadShared(base: number, tid: number) { + 'use gpu'; + (localKeys.$[tid] as number) = sortLayout.$.data[base + tid] as number; + (localKeys.$[tid + WORKGROUP_SIZE] as number) = sortLayout.$.data[ + base + tid + WORKGROUP_SIZE + ] as number; + if (hasPayload) { + (localVals.$[tid] as number) = std.copy(valsLayout.$.vals[base + tid] as number); + (localVals.$[tid + WORKGROUP_SIZE] as number) = std.copy( + valsLayout.$.vals[base + tid + WORKGROUP_SIZE] as number, + ); + } } -}); -export function createBitonicSorter( - root: TgpuRoot, - data: TgpuBuffer> & StorageFlag, - options?: BitonicSorterOptions, -): BitonicSorter { - const originalSize = data.dataType.elementCount; - const paddedSize = nextPowerOf2(originalSize); - const wasPadded = paddedSize !== originalSize; - - const paddingValue = options?.paddingValue ?? 0xffffffff; - const compareFunc = options?.compare ?? defaultCompare; - - let paddingResources: { - workBuffer: TgpuBuffer> & StorageFlag; - copyPadParams: TgpuBuffer & UniformFlag; - copyBackParams: TgpuBuffer & UniformFlag; - copyPadBindGroup: TgpuBindGroup<(typeof copyLayout)['entries']>; - copyBackBindGroup: TgpuBindGroup<(typeof copyLayout)['entries']>; - } | null = null; - let workBuffer: TgpuBuffer> & StorageFlag; - - const sortWorkgroupsTotal = Math.ceil(paddedSize / 2 / WORKGROUP_SIZE); - const [sortWorkgroupsX, sortWorkgroupsY, sortWorkgroupsZ] = - decomposeWorkgroups(sortWorkgroupsTotal); - - const padWorkgroupsTotal = Math.ceil(paddedSize / WORKGROUP_SIZE); - const [padWorkgroupsX, padWorkgroupsY, padWorkgroupsZ] = decomposeWorkgroups(padWorkgroupsTotal); - - const copyBackWorkgroupsTotal = Math.ceil(originalSize / WORKGROUP_SIZE); - const [copyBackWorkgroupsX, copyBackWorkgroupsY, copyBackWorkgroupsZ] = - decomposeWorkgroups(copyBackWorkgroupsTotal); - - if (wasPadded) { - const paddedWorkBuffer = root.createBuffer(d.arrayOf(d.u32, paddedSize)).$usage('storage'); - - const copyPadParams = root - .createBuffer(copyParamsType, { - srcLength: originalSize, - dstLength: paddedSize, - paddingValue, - }) - .$usage('uniform'); + function storeShared(base: number, tid: number) { + 'use gpu'; + (sortLayout.$.data[base + tid] as number) = localKeys.$[tid] as number; + (sortLayout.$.data[base + tid + WORKGROUP_SIZE] as number) = localKeys.$[ + tid + WORKGROUP_SIZE + ] as number; + if (hasPayload) { + (valsLayout.$.vals[base + tid] as number) = std.copy(localVals.$[tid] as number); + (valsLayout.$.vals[base + tid + WORKGROUP_SIZE] as number) = std.copy( + localVals.$[tid + WORKGROUP_SIZE] as number, + ); + } + } - const copyBackParams = root - .createBuffer(copyParamsType, { - srcLength: originalSize, - dstLength: originalSize, - paddingValue: 0, - }) - .$usage('uniform'); + function swapLocalAt(a: number, b: number, left: number, right: number) { + 'use gpu'; + localKeys.$[a] = right; + localKeys.$[b] = left; + if (hasPayload) { + const tmp = std.copy(localVals.$[a] as number); + (localVals.$[a] as number) = std.copy(localVals.$[b] as number); + (localVals.$[b] as number) = std.copy(tmp); + } + } - paddingResources = { - workBuffer: paddedWorkBuffer, - copyPadParams, - copyBackParams, - copyPadBindGroup: root.createBindGroup(copyLayout, { - src: data, - dst: paddedWorkBuffer, - params: copyPadParams, - }), - copyBackBindGroup: root.createBindGroup(copyLayout, { - src: paddedWorkBuffer, - dst: data, - params: copyBackParams, - }), - }; + function exchangeLocal(base: number, iLocal: number, stride: number, k: number) { + 'use gpu'; + const jLocal = iLocal + stride; + const left = localKeys.$[iLocal] as number; + const right = localKeys.$[jLocal] as number; + const ascending = ((base + iLocal) & k) === 0; - workBuffer = paddedWorkBuffer; - } else { - workBuffer = data; + if (std.select(compareSlot.$(left, right), compareSlot.$(right, left), ascending)) { + swapLocalAt(iLocal, jLocal, left, right); + } + } + + function mergeDown(base: number, tid: number, startShift: number, k: number) { + 'use gpu'; + for (let jShift = d.u32(startShift); jShift > 0; jShift--) { + std.workgroupBarrier(); + const stride = d.u32(1) << (jShift - 1); + const below = tid & (stride - 1); + const above = tid >> (jShift - 1); + exchangeLocal(base, below + above * (stride << 1), stride, k); + } } - const uniformBuffer = root.createBuffer(sortUniformsType).$usage('uniform'); + const localSort = tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })(({ + lid, + wid, + numWorkgroups, + }) => { + const base = flatWorkgroupIndex(wid, numWorkgroups) * LOCAL_BLOCK; + if (base >= sortLayout.$.data.length) { + return; + } - const sortBindGroup = root.createBindGroup(sortLayout, { - data: workBuffer, - uniforms: uniformBuffer, + loadShared(base, lid.x); + for (let kShift = d.u32(1); kShift <= LOCAL_BLOCK_LOG2; kShift++) { + mergeDown(base, lid.x, kShift, d.u32(1) << kShift); + } + std.workgroupBarrier(); + storeShared(base, lid.x); }); - const sortPipeline = root.with(compareSlot, compareFunc).createComputePipeline({ - compute: bitonicStepKernel, + const localMerge = tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })(({ + lid, + wid, + numWorkgroups, + }) => { + const base = flatWorkgroupIndex(wid, numWorkgroups) * LOCAL_BLOCK; + if (base >= sortLayout.$.data.length) { + return; + } + + loadShared(base, lid.x); + mergeDown(base, lid.x, d.u32(LOCAL_BLOCK_LOG2), sortLayout.$.uniforms.k); + std.workgroupBarrier(); + storeShared(base, lid.x); }); - const copyPadPipeline = root.createComputePipeline({ compute: copyPadKernel }); + return { localSort, localMerge }; +} + +interface SortStep { + pipeline: TgpuComputePipeline; + workgroups: [number, number, number]; +} + +/** + * Creates a bitonic sorter for a `u32`, `i32` or `f32` key buffer, optionally reordering + * a payload buffer alongside the keys. The order is defined by an arbitrary comparator. + * All GPU resources are created up front, so `run` only records dispatches. + */ +export function createBitonicSorter< + TKey extends BitonicKeyType, + TValue extends d.AnyWgslData = d.AnyWgslData, +>( + root: TgpuRoot, + data: TgpuBuffer> & StorageFlag, + options?: BitonicSorterOptions, +): BitonicSorter { + const keyBuffer = data as KeyBuffer; + const valueBuffer = options?.values as ValueBuffer | undefined; + + const keyType = keyBuffer.dataType.elementType; + const size = keyBuffer.dataType.elementCount; + const paddedSize = nextPowerOf2(size); + + if (size === 0) { + throw new Error('Cannot create a bitonic sorter for an empty buffer.'); + } + if (valueBuffer && valueBuffer.dataType.elementCount !== size) { + throw new Error( + `The values buffer (${valueBuffer.dataType.elementCount} elements) must match the key buffer (${size} elements).`, + ); + } + if (valueBuffer && paddedSize !== size) { + throw new Error('Bitonic sorting with a values buffer requires a power-of-two element count.'); + } + + const schemas = makeBitonicSchemas(keyType, valueBuffer?.dataType.elementType); + const owned: { destroy(): void }[] = []; + const steps: SortStep[] = []; + + let workBuffer = keyBuffer; + let unpadStep: SortStep | undefined; + + if (paddedSize !== size) { + const { copyLayout, pad, unpad } = makePaddingKernels(keyType, size, paddedSize); + const padding = root + .createBuffer(keyType, options?.paddingValue ?? defaultPaddingValues[keyType.type]) + .$usage('uniform'); + workBuffer = root.createBuffer(d.arrayOf(keyType, paddedSize)).$usage('storage') as KeyBuffer; + owned.push(padding, workBuffer); + + steps.push({ + pipeline: root + .createComputePipeline({ compute: pad }) + .with(root.createBindGroup(copyLayout, { src: keyBuffer, dst: workBuffer, padding })), + workgroups: decomposeWorkgroups(Math.ceil(paddedSize / WORKGROUP_SIZE)), + }); + unpadStep = { + pipeline: root + .createComputePipeline({ compute: unpad }) + .with(root.createBindGroup(copyLayout, { src: workBuffer, dst: keyBuffer, padding })), + workgroups: decomposeWorkgroups(Math.ceil(size / WORKGROUP_SIZE)), + }; + } + + const valsBindGroup = valueBuffer + ? root.createBindGroup(schemas.valsLayout, { vals: valueBuffer }) + : undefined; - const copyBackPipeline = root.createComputePipeline({ compute: copyBackKernel }); + const compare = options?.compare ?? defaultCompare; - const log2N = Math.log2(paddedSize); - const totalSteps = (log2N * (log2N + 1)) / 2; + function createSortPipeline(compute: TgpuComputeFn): TgpuComputePipeline { + const pipeline = root.with(compareSlot, compare).createComputePipeline({ compute }); + return valsBindGroup ? pipeline.with(valsBindGroup) : pipeline; + } - function run(runOptions?: BitonicSorterRunOptions): void { - const querySet = runOptions?.querySet; + function pushStep( + pipeline: TgpuComputePipeline, + k: number, + jShift: number, + workgroups: [number, number, number], + ): void { + const uniforms = root.createBuffer(sortUniformsType, { k, jShift }).$usage('uniform'); + owned.push(uniforms); + + steps.push({ + pipeline: pipeline.with( + root.createBindGroup(schemas.sortLayout, { data: workBuffer, uniforms }), + ), + workgroups, + }); + } - if (paddingResources) { - let pipeline = copyPadPipeline.with(paddingResources.copyPadBindGroup); - if (querySet) { - pipeline = pipeline.withTimestampWrites({ - querySet, - beginningOfPassWriteIndex: 0, - }); + const payloadSize = schemas.valueType ? d.sizeOf(schemas.valueType) : 0; + const sharedMemoryBytes = LOCAL_BLOCK * (d.sizeOf(keyType) + payloadSize); + const useLocalKernels = + paddedSize >= LOCAL_BLOCK && + sharedMemoryBytes <= root.device.limits.maxComputeWorkgroupStorageSize; + + const globalWorkgroups = decomposeWorkgroups(Math.ceil(paddedSize / 2 / WORKGROUP_SIZE)); + const globalPipeline = createSortPipeline(makeGlobalStepKernel(schemas)); + + if (useLocalKernels) { + const { localSort, localMerge } = makeLocalKernels(schemas); + const localSortPipeline = createSortPipeline(localSort); + const localMergePipeline = createSortPipeline(localMerge); + const blockWorkgroups = decomposeWorkgroups(paddedSize / LOCAL_BLOCK); + + pushStep(localSortPipeline, 0, 0, blockWorkgroups); + for (let k = LOCAL_BLOCK * 2; k <= paddedSize; k <<= 1) { + for (let j = k >> 1; j >= LOCAL_BLOCK; j >>= 1) { + pushStep(globalPipeline, k, Math.log2(j), globalWorkgroups); } - pipeline.dispatchWorkgroups(padWorkgroupsX, padWorkgroupsY, padWorkgroupsZ); + pushStep(localMergePipeline, k, 0, blockWorkgroups); } - - let stepIndex = 0; + } else { for (let k = 2; k <= paddedSize; k <<= 1) { for (let j = k >> 1; j > 0; j >>= 1) { - const jShift = 31 - Math.clz32(j); - uniformBuffer.write({ k, jShift }); - - let pipeline = sortPipeline.with(sortBindGroup); - - if (querySet && !paddingResources) { - const isFirst = stepIndex === 0; - const isLast = stepIndex === totalSteps - 1; - if (isFirst || isLast) { - pipeline = pipeline.withTimestampWrites({ - querySet, - ...(isFirst && { beginningOfPassWriteIndex: 0 }), - ...(isLast && { endOfPassWriteIndex: 1 }), - }); - } - } - - pipeline.dispatchWorkgroups(sortWorkgroupsX, sortWorkgroupsY, sortWorkgroupsZ); - stepIndex++; - } - } - - if (paddingResources) { - let pipeline = copyBackPipeline.with(paddingResources.copyBackBindGroup); - if (querySet) { - pipeline = pipeline.withTimestampWrites({ - querySet, - endOfPassWriteIndex: 1, - }); + pushStep(globalPipeline, k, Math.log2(j), globalWorkgroups); } - pipeline.dispatchWorkgroups(copyBackWorkgroupsX, copyBackWorkgroupsY, copyBackWorkgroupsZ); } } - function destroy(): void { - uniformBuffer.destroy(); - if (paddingResources) { - paddingResources.workBuffer.destroy(); - paddingResources.copyPadParams.destroy(); - paddingResources.copyBackParams.destroy(); - } + if (unpadStep) { + steps.push(unpadStep); } return { - originalSize, + size, paddedSize, - wasPadded, - run, - destroy, + + run(runOptions?: RunOptions): void { + const recording = beginRunPass(root.device, runOptions); + for (const step of steps) { + bindPass(step.pipeline, recording.pass).dispatchWorkgroups(...step.workgroups); + } + recording.finish(); + }, + + destroy(): void { + for (const buffer of owned) { + buffer.destroy(); + } + }, }; } diff --git a/packages/typegpu-sort/src/bitonic/index.ts b/packages/typegpu-sort/src/bitonic/index.ts index 5e28296880..0c463d9f40 100644 --- a/packages/typegpu-sort/src/bitonic/index.ts +++ b/packages/typegpu-sort/src/bitonic/index.ts @@ -1,4 +1,2 @@ -export { createBitonicSorter } from './bitonicSort.ts'; -export { compareSlot, defaultCompare } from './slots.ts'; -export type { BitonicSorter, BitonicSorterOptions, BitonicSorterRunOptions } from './types.ts'; -export { decomposeWorkgroups } from './utils.ts'; +export { type BitonicKeyType, createBitonicSorter } from './bitonicSort.ts'; +export type { BitonicSorter, BitonicSorterOptions } from './types.ts'; diff --git a/packages/typegpu-sort/src/bitonic/slots.ts b/packages/typegpu-sort/src/bitonic/slots.ts index 4a91f26cd2..ef8f228a48 100644 --- a/packages/typegpu-sort/src/bitonic/slots.ts +++ b/packages/typegpu-sort/src/bitonic/slots.ts @@ -1,8 +1,14 @@ -import { tgpu, d } from 'typegpu'; +import { tgpu } from 'typegpu'; -/** Default comparison function: ascending order (a < b means a comes before b) */ -export const defaultCompare = tgpu.fn([d.u32, d.u32], d.bool)((a, b) => a < b); +export function defaultCompare(a: number, b: number): boolean { + 'use gpu'; + return a < b; +} -/** Slot for customizing the comparison function in bitonic sort. - * The function should return true if the first argument should come before the second. */ -export const compareSlot = tgpu.slot<(a: number, b: number) => boolean>(defaultCompare); +export const defaultPaddingValues = { + u32: 0xffffffff, + i32: 2147483647, + f32: Number.POSITIVE_INFINITY, +} as const; + +export const compareSlot = tgpu.slot<(a: number, b: number) => boolean>(); diff --git a/packages/typegpu-sort/src/bitonic/types.ts b/packages/typegpu-sort/src/bitonic/types.ts index 17fa03b41a..63e09e1073 100644 --- a/packages/typegpu-sort/src/bitonic/types.ts +++ b/packages/typegpu-sort/src/bitonic/types.ts @@ -1,36 +1,24 @@ -import type { TgpuQuerySet } from 'typegpu'; +import type { d, StorageFlag, TgpuBuffer } from 'typegpu'; +import type { Sorter } from '../types.ts'; -export interface BitonicSorterOptions { +export interface BitonicSorterOptions { /** Custom comparison function. Returns true if first argument should come before second. * Default: ascending order (a < b) */ compare?: (a: number, b: number) => boolean; /** * Value used to pad arrays to power-of-2 length. Must sort to the end with your comparator. - * Default: `0xFFFFFFFF` (works for ascending). For descending order, use `0`. + * Defaults to the maximum value of the key type, which works for ascending. For descending + * order, use the minimum value of the key type. */ paddingValue?: number; -} - -export interface BitonicSorterRunOptions { /** - * Optional timestamp query set for GPU timing. Must have at least 2 entries. - * Timestamps are written to indices 0 and 1. For non-power-of-2 arrays, timing - * includes the padding copy passes. + * Payload buffer reordered alongside the keys, e.g. indices into another data structure. + * Must have the same power-of-two element count as the key buffer. */ - querySet?: TgpuQuerySet<'timestamp'>; + values?: TgpuBuffer> & StorageFlag; } -export interface BitonicSorter { - /** Original size of the input array */ - readonly originalSize: number; - /** Size after padding to power of 2 */ +export interface BitonicSorter extends Sorter { + /** Size the keys are padded to, a power of two */ readonly paddedSize: number; - /** Whether the array was padded */ - readonly wasPadded: boolean; - - /** Execute the sort. Can be called repeatedly. */ - run(options?: BitonicSorterRunOptions): void; - - /** Clean up all GPU resources. */ - destroy(): void; } diff --git a/packages/typegpu-sort/src/bitonic/utils.ts b/packages/typegpu-sort/src/bitonic/utils.ts deleted file mode 100644 index 8961c7507c..0000000000 --- a/packages/typegpu-sort/src/bitonic/utils.ts +++ /dev/null @@ -1,39 +0,0 @@ -/** - * Returns the next power of 2 greater than or equal to n. - * If n is already a power of 2, returns n. - */ -export function nextPowerOf2(n: number): number { - if (n <= 0) return 1; - if ((n & (n - 1)) === 0) return n; - let p = 1; - while (p < n) p <<= 1; - return p; -} - -const MAX_WORKGROUPS_PER_DIMENSION = 65535; - -/** - * Decomposes a total workgroup count into a 3D dispatch grid (x, y, z), - * respecting the WebGPU limit of 65535 workgroups per dimension. - */ -export function decomposeWorkgroups(total: number): [number, number, number] { - if (total <= 0) { - return [1, 1, 1]; - } - - const x = Math.min(total, MAX_WORKGROUPS_PER_DIMENSION); - const remainingAfterX = Math.ceil(total / x); - - const y = Math.min(remainingAfterX, MAX_WORKGROUPS_PER_DIMENSION); - const remainingAfterY = Math.ceil(remainingAfterX / y); - - const z = Math.min(remainingAfterY, MAX_WORKGROUPS_PER_DIMENSION); - - if (Math.ceil(total / (x * y * z)) > 1) { - throw new Error( - `Required workgroups (${total}) exceed device dispatch limits (${MAX_WORKGROUPS_PER_DIMENSION} per dimension)`, - ); - } - - return [x, y, z]; -} diff --git a/packages/typegpu-sort/src/dispatch.ts b/packages/typegpu-sort/src/dispatch.ts new file mode 100644 index 0000000000..d4c5d767ee --- /dev/null +++ b/packages/typegpu-sort/src/dispatch.ts @@ -0,0 +1,31 @@ +import { d } from 'typegpu'; + +const MAX_WORKGROUPS_PER_DIMENSION = 65535; + +/** + * Decomposes a total workgroup count into a 3D dispatch grid (x, y, z), + * respecting the WebGPU limit of 65535 workgroups per dimension. The grid can cover more + * workgroups than requested, so kernels have to guard against running past their data. + */ +export function decomposeWorkgroups(total: number): [number, number, number] { + if (total <= 1) { + return [1, 1, 1]; + } + + const x = Math.min(total, MAX_WORKGROUPS_PER_DIMENSION); + const rows = Math.ceil(total / x); + const y = Math.min(rows, MAX_WORKGROUPS_PER_DIMENSION); + + return [x, y, Math.ceil(rows / y)]; +} + +export const dispatchIn = { + lid: d.builtin.localInvocationId, + wid: d.builtin.workgroupId, + numWorkgroups: d.builtin.numWorkgroups, +} as const; + +export function flatWorkgroupIndex(wid: d.v3u, numWorkgroups: d.v3u): number { + 'use gpu'; + return wid.x + wid.y * numWorkgroups.x + wid.z * numWorkgroups.x * numWorkgroups.y; +} diff --git a/packages/typegpu-sort/src/index.ts b/packages/typegpu-sort/src/index.ts index 6268f407c2..f9b86189a0 100644 --- a/packages/typegpu-sort/src/index.ts +++ b/packages/typegpu-sort/src/index.ts @@ -1,14 +1,18 @@ -export { - compareSlot, - createBitonicSorter, - defaultCompare, - decomposeWorkgroups, -} from './bitonic/index.ts'; -export type { - BitonicSorter, - BitonicSorterOptions, - BitonicSorterRunOptions, -} from './bitonic/index.ts'; +export { decomposeWorkgroups } from './dispatch.ts'; +export type { RunOptions, Sorter } from './types.ts'; + +export { type BitonicKeyType, createBitonicSorter } from './bitonic/index.ts'; +export type { BitonicSorter, BitonicSorterOptions } from './bitonic/index.ts'; -export { prefixScan, scan, createPrefixScanComputer, PrefixScanComputer } from './scan/index.ts'; -export type { BinaryOp } from './scan/index.ts'; +export { createRadixSorter } from './radix/index.ts'; +export type { RadixSorterOptions } from './radix/index.ts'; + +export { + createPrefixScanComputer, + type PrefixScanComputer, + type PrefixScanPlan, + prefixScan, + reduce, + type ScanBuffer, +} from './scan/index.ts'; +export type { BinaryOp, ScanElementType } from './scan/index.ts'; diff --git a/packages/typegpu-sort/src/radix/count.ts b/packages/typegpu-sort/src/radix/count.ts new file mode 100644 index 0000000000..c8115e6a56 --- /dev/null +++ b/packages/typegpu-sort/src/radix/count.ts @@ -0,0 +1,45 @@ +import { tgpu, d, std } from 'typegpu'; +import { dispatchIn, flatWorkgroupIndex } from '../dispatch.ts'; +import { + histLayout, + KEYS_PER_THREAD, + type RadixSchemas, + shiftLayout, + TILE_SIZE, + TILE_THREADS, + wgHist, +} from './schemas.ts'; + +export function makeCountKernel(schemas: RadixSchemas, elementCount: number, numTiles: number) { + const { ioLayout, digitFn } = schemas; + const needsBoundsCheck = elementCount % TILE_SIZE !== 0; + const lastIndex = elementCount - 1; + + return tgpu.computeFn({ workgroupSize: [TILE_THREADS], in: dispatchIn })( + ({ lid, wid, numWorkgroups }) => { + const localIdx = lid.x; + const tileId = flatWorkgroupIndex(wid, numWorkgroups); + if (tileId >= numTiles) { + return; + } + + const tileBase = tileId * TILE_SIZE; + const shift = shiftLayout.$.shift; + + for (const k of tgpu.unroll(std.range(KEYS_PER_THREAD))) { + const globalIdx = tileBase + k * TILE_THREADS + localIdx; + const loadIdx = needsBoundsCheck ? std.min(globalIdx, lastIndex) : globalIdx; + const digit = digitFn(ioLayout.$.src[loadIdx] as number, shift); + + if (needsBoundsCheck ? globalIdx < elementCount : true) { + std.atomicAdd(wgHist.$[digit] as d.atomicU32, 1); + } + } + std.workgroupBarrier(); + + histLayout.$.hist[localIdx * numTiles + tileId] = std.atomicLoad( + wgHist.$[localIdx] as d.atomicU32, + ); + }, + ); +} diff --git a/packages/typegpu-sort/src/radix/index.ts b/packages/typegpu-sort/src/radix/index.ts new file mode 100644 index 0000000000..18660852cf --- /dev/null +++ b/packages/typegpu-sort/src/radix/index.ts @@ -0,0 +1,2 @@ +export { createRadixSorter } from './radixSort.ts'; +export type { RadixSorterOptions } from './types.ts'; diff --git a/packages/typegpu-sort/src/radix/radixSort.ts b/packages/typegpu-sort/src/radix/radixSort.ts new file mode 100644 index 0000000000..c71225d001 --- /dev/null +++ b/packages/typegpu-sort/src/radix/radixSort.ts @@ -0,0 +1,148 @@ +import { d, std, type StorageFlag, type TgpuBuffer, type TgpuRoot } from 'typegpu'; +import { decomposeWorkgroups } from '../dispatch.ts'; +import { beginRunPass, bindPass } from '../runPass.ts'; +import { createPrefixScanComputer } from '../scan/index.ts'; +import type { RunOptions, Sorter } from '../types.ts'; +import { makeCountKernel } from './count.ts'; +import { makeScatterKernel } from './scatter.ts'; +import { + histLayout, + makeRadixSchemas, + NUM_PASSES, + RADIX_BITS, + RADIX_SIZE, + type RadixKeyType, + shiftLayout, + TILE_SIZE, +} from './schemas.ts'; +import type { RadixSorterOptions } from './types.ts'; + +type KeyBuffer = TgpuBuffer> & StorageFlag; +type ValueBuffer = TgpuBuffer> & StorageFlag; + +/** + * Creates a stable LSD radix sorter for a `u32`, `i32` or `f32` key buffer, optionally + * reordering a payload buffer alongside the keys. Keys are ordered by the natural order + * of their type. All GPU resources are created up front, so `run` only records dispatches. + * + * For `f32` keys sorted ascending, NaNs with a cleared sign bit sort after +Infinity and + * NaNs with a set sign bit sort before -Infinity. -0 and +0 compare equal. + */ +export function createRadixSorter< + TKey extends RadixKeyType, + TValue extends d.AnyWgslData = d.AnyWgslData, +>( + root: TgpuRoot, + keys: TgpuBuffer> & StorageFlag, + options?: RadixSorterOptions, +): Sorter { + const keyBuffer = keys as KeyBuffer; + const valueBuffer = options?.values as ValueBuffer | undefined; + const keyType = keyBuffer.dataType.elementType; + const size = keyBuffer.dataType.elementCount; + + if (size === 0) { + throw new Error('Cannot create a radix sorter for an empty buffer.'); + } + if (valueBuffer && valueBuffer.dataType.elementCount !== size) { + throw new Error( + `The values buffer (${valueBuffer.dataType.elementCount} elements) must match the key buffer (${size} elements).`, + ); + } + + const schemas = makeRadixSchemas( + keyType, + options?.direction ?? 'ascending', + valueBuffer?.dataType.elementType, + ); + + const numTiles = Math.ceil(size / TILE_SIZE); + const dispatch = decomposeWorkgroups(numTiles); + + const histBuffer = root.createBuffer(d.arrayOf(d.u32, numTiles * RADIX_SIZE)).$usage('storage'); + const tempBuffer = root.createBuffer(d.arrayOf(keyType, size)).$usage('storage') as KeyBuffer; + const owned: { destroy(): void }[] = [histBuffer, tempBuffer]; + + const scanPlan = createPrefixScanComputer(root, { + operation: std.add, + identityElement: 0, + dataType: d.u32, + }).prepare(histBuffer); + + const tempValues = + valueBuffer && + (root + .createBuffer(d.arrayOf(valueBuffer.dataType.elementType, size)) + .$usage('storage') as ValueBuffer); + if (tempValues) { + owned.push(tempValues); + } + + const histBg = root.createBindGroup(histLayout, { hist: histBuffer }); + const ioBgKeysToTemp = root.createBindGroup(schemas.ioLayout, { + src: keyBuffer, + dst: tempBuffer, + }); + const ioBgTempToKeys = root.createBindGroup(schemas.ioLayout, { + src: tempBuffer, + dst: keyBuffer, + }); + + const valuesBgs = + valueBuffer && tempValues + ? { + keysToTemp: root.createBindGroup(schemas.valuesLayout, { + srcVals: valueBuffer, + dstVals: tempValues, + }), + tempToKeys: root.createBindGroup(schemas.valuesLayout, { + srcVals: tempValues, + dstVals: valueBuffer, + }), + } + : undefined; + + const countPipeline = root + .createComputePipeline({ compute: makeCountKernel(schemas, size, numTiles) }) + .with(histBg); + const scatterPipeline = root + .createComputePipeline({ compute: makeScatterKernel(schemas, size, numTiles) }) + .with(histBg); + + const passes = Array.from({ length: NUM_PASSES }, (_, pass) => { + const shift = root.createBuffer(d.u32, pass * RADIX_BITS).$usage('uniform'); + owned.push(shift); + + const shiftBg = root.createBindGroup(shiftLayout, { shift }); + const forward = pass % 2 === 0; + const ioBg = forward ? ioBgKeysToTemp : ioBgTempToKeys; + const valuesBg = forward ? valuesBgs?.keysToTemp : valuesBgs?.tempToKeys; + + const scatter = scatterPipeline.with(ioBg).with(shiftBg); + return { + count: countPipeline.with(ioBg).with(shiftBg), + scatter: valuesBg ? scatter.with(valuesBg) : scatter, + }; + }); + + return { + size, + + run(runOptions?: RunOptions): void { + const recording = beginRunPass(root.device, runOptions); + for (const { count, scatter } of passes) { + bindPass(count, recording.pass).dispatchWorkgroups(...dispatch); + scanPlan.run({ pass: recording.pass }); + bindPass(scatter, recording.pass).dispatchWorkgroups(...dispatch); + } + recording.finish(); + }, + + destroy(): void { + scanPlan.destroy(); + for (const buffer of owned) { + buffer.destroy(); + } + }, + }; +} diff --git a/packages/typegpu-sort/src/radix/scatter.ts b/packages/typegpu-sort/src/radix/scatter.ts new file mode 100644 index 0000000000..0eb5c61f44 --- /dev/null +++ b/packages/typegpu-sort/src/radix/scatter.ts @@ -0,0 +1,85 @@ +import { tgpu, d, std, type TgpuComputeFn } from 'typegpu'; +import { dispatchIn, flatWorkgroupIndex } from '../dispatch.ts'; +import { + histLayout, + KEYS_PER_THREAD, + RADIX_SIZE, + type RadixSchemas, + shiftLayout, + TILE_SIZE, + TILE_THREADS, +} from './schemas.ts'; + +const BITSET_WORD_BITS = 32; +const BITSET_WORD_SHIFT = Math.log2(BITSET_WORD_BITS); +const BITSET_WORDS = TILE_THREADS / BITSET_WORD_BITS; + +const runningTotal = tgpu.workgroupVar(d.arrayOf(d.u32, RADIX_SIZE)); +const digitBits = tgpu.workgroupVar(d.arrayOf(d.atomic(d.u32), BITSET_WORDS * RADIX_SIZE)); + +export function makeScatterKernel( + schemas: RadixSchemas, + elementCount: number, + numTiles: number, +): TgpuComputeFn { + const { ioLayout, digitFn, writeOutput } = schemas; + const needsBoundsCheck = elementCount % TILE_SIZE !== 0; + const lastIndex = elementCount - 1; + + return tgpu.computeFn({ workgroupSize: [TILE_THREADS], in: dispatchIn })( + ({ lid, wid, numWorkgroups }) => { + const localIdx = lid.x; + const tileId = flatWorkgroupIndex(wid, numWorkgroups); + if (tileId >= numTiles) { + return; + } + + const tileBase = tileId * TILE_SIZE; + const shift = shiftLayout.$.shift; + const bitsetWord = localIdx >> BITSET_WORD_SHIFT; + const bitsetMask = d.u32(1) << (localIdx & (BITSET_WORD_BITS - 1)); + const earlierBits = bitsetMask - 1; + + runningTotal.$[localIdx] = histLayout.$.hist[localIdx * numTiles + tileId] as number; + + for (const k of tgpu.unroll(std.range(KEYS_PER_THREAD))) { + const globalIdx = tileBase + k * TILE_THREADS + localIdx; + const loadIdx = needsBoundsCheck ? std.min(globalIdx, lastIndex) : globalIdx; + const key = ioLayout.$.src[loadIdx] as number; + const digit = digitFn(key, shift); + const inBounds = needsBoundsCheck ? globalIdx < elementCount : true; + + if (inBounds) { + std.atomicOr(digitBits.$[bitsetWord * RADIX_SIZE + digit] as d.atomicU32, bitsetMask); + } + std.workgroupBarrier(); + + let rank = d.u32(0); + let digitTotal = d.u32(0); + if (inBounds) { + for (const word of tgpu.unroll(std.range(BITSET_WORDS))) { + const bits = std.atomicLoad(digitBits.$[word * RADIX_SIZE + digit] as d.atomicU32); + const mask = std.select( + std.select(d.u32(0), earlierBits, word === bitsetWord), + d.u32(0xffffffff), + word < bitsetWord, + ); + rank = rank + std.countOneBits(bits & mask); + digitTotal = digitTotal + std.countOneBits(bits); + } + + writeOutput(key, globalIdx, (runningTotal.$[digit] as number) + rank); + } + std.workgroupBarrier(); + + if (inBounds) { + std.atomicStore(digitBits.$[bitsetWord * RADIX_SIZE + digit] as d.atomicU32, 0); + if (rank === 0) { + runningTotal.$[digit] = (runningTotal.$[digit] as number) + digitTotal; + } + } + std.workgroupBarrier(); + } + }, + ); +} diff --git a/packages/typegpu-sort/src/radix/schemas.ts b/packages/typegpu-sort/src/radix/schemas.ts new file mode 100644 index 0000000000..bb8d5d1c77 --- /dev/null +++ b/packages/typegpu-sort/src/radix/schemas.ts @@ -0,0 +1,99 @@ +import { tgpu, d, std } from 'typegpu'; + +export const RADIX_BITS = 8; +export const RADIX_SIZE = 1 << RADIX_BITS; +export const NUM_PASSES = 32 / RADIX_BITS; +export const TILE_THREADS = RADIX_SIZE; +export const KEYS_PER_THREAD = 8; +export const TILE_SIZE = TILE_THREADS * KEYS_PER_THREAD; + +export type RadixKeyType = d.U32 | d.I32 | d.F32; +export type SortDirection = 'ascending' | 'descending'; + +export const histLayout = tgpu.bindGroupLayout({ + hist: { storage: d.arrayOf(d.u32), access: 'mutable' }, +}); + +export const shiftLayout = tgpu.bindGroupLayout({ + shift: { uniform: d.u32 }, +}); + +export const wgHist = tgpu.workgroupVar(d.arrayOf(d.atomic(d.u32), RADIX_SIZE)); + +function digitOfU32(v: number, shift: number): number { + 'use gpu'; + return (v >> shift) & (RADIX_SIZE - 1); +} + +function digitOfI32(v: number, shift: number): number { + 'use gpu'; + const raw = d.u32((v >> shift) & (RADIX_SIZE - 1)); + return raw ^ std.select(d.u32(0), d.u32(RADIX_SIZE / 2), shift === 24); +} + +function digitOfF32(v: number, shift: number): number { + 'use gpu'; + // -0 and +0 must map to the same bits, otherwise they sort apart + const bits = std.select(std.bitcastF32toU32(v), d.u32(0), v === 0); + const mask = std.select(d.u32(0x80000000), d.u32(0xffffffff), bits >> 31 === 1); + return ((bits ^ mask) >> shift) & (RADIX_SIZE - 1); +} + +const ascendingDigits = { + u32: digitOfU32, + i32: digitOfI32, + f32: digitOfF32, +} as const; + +export function makeDigitFn(keyType: RadixKeyType, direction: SortDirection) { + const ascending = ascendingDigits[keyType.type]; + if (direction === 'ascending') { + return ascending; + } + + function descendingDigit(v: number, shift: number) { + 'use gpu'; + return RADIX_SIZE - 1 - ascending(v, shift); + } + + return descendingDigit; +} + +export function makeRadixSchemas( + keyType: RadixKeyType, + direction: SortDirection, + valueType?: d.AnyWgslData, +) { + const ioLayout = tgpu.bindGroupLayout({ + src: { storage: d.arrayOf(keyType), access: 'readonly' }, + dst: { storage: d.arrayOf(keyType), access: 'mutable' }, + }); + + const hasPayload = valueType !== undefined; + const payloadType = valueType ?? d.u32; + + const valuesLayout = tgpu.bindGroupLayout({ + srcVals: { storage: d.arrayOf(payloadType), access: 'readonly' }, + dstVals: { storage: d.arrayOf(payloadType), access: 'mutable' }, + }); + + function writeOutput(key: number, srcIdx: number, dstIdx: number) { + 'use gpu'; + (ioLayout.$.dst[dstIdx] as number) = key; + if (hasPayload) { + (valuesLayout.$.dstVals[dstIdx] as number) = std.copy( + valuesLayout.$.srcVals[srcIdx] as number, + ); + } + } + + return { + keyType, + ioLayout, + valuesLayout, + writeOutput, + digitFn: makeDigitFn(keyType, direction), + }; +} + +export type RadixSchemas = ReturnType; diff --git a/packages/typegpu-sort/src/radix/types.ts b/packages/typegpu-sort/src/radix/types.ts new file mode 100644 index 0000000000..b8adaae1ef --- /dev/null +++ b/packages/typegpu-sort/src/radix/types.ts @@ -0,0 +1,12 @@ +import type { d, StorageFlag, TgpuBuffer } from 'typegpu'; +import type { SortDirection } from './schemas.ts'; + +export interface RadixSorterOptions { + /** Sort order. Defaults to `'ascending'` */ + direction?: SortDirection; + /** + * Payload buffer reordered alongside the keys, e.g. indices into another data structure. + * Must have the same element count as the key buffer. + */ + values?: TgpuBuffer> & StorageFlag; +} diff --git a/packages/typegpu-sort/src/runPass.ts b/packages/typegpu-sort/src/runPass.ts new file mode 100644 index 0000000000..14d1e7dcbc --- /dev/null +++ b/packages/typegpu-sort/src/runPass.ts @@ -0,0 +1,45 @@ +import type { TgpuComputePass, TgpuComputePipeline } from 'typegpu'; +import type { RunOptions } from './types.ts'; + +export type RunPass = GPUComputePassEncoder | TgpuComputePass; + +export interface RunRecording { + pass: RunPass; + finish(): void; +} + +export function bindPass(pipeline: TgpuComputePipeline, pass: RunPass): TgpuComputePipeline { + if ('resourceType' in pass) { + return pipeline.with(pass); + } + return pipeline.with(pass); +} + +const noop = () => {}; + +export function beginRunPass(device: GPUDevice, options?: RunOptions): RunRecording { + if (options?.pass) { + return { pass: options.pass, finish: noop }; + } + + const externalEncoder = options?.encoder; + if (externalEncoder) { + const pass = externalEncoder.beginComputePass(); + return { + pass, + finish() { + pass.end(); + }, + }; + } + + const encoder = device.createCommandEncoder(); + const pass = encoder.beginComputePass(); + return { + pass, + finish() { + pass.end(); + device.queue.submit([encoder.finish()]); + }, + }; +} diff --git a/packages/typegpu-sort/src/scan/compute/applySums.ts b/packages/typegpu-sort/src/scan/compute/applySums.ts deleted file mode 100644 index 64b7db6578..0000000000 --- a/packages/typegpu-sort/src/scan/compute/applySums.ts +++ /dev/null @@ -1,25 +0,0 @@ -import { tgpu, d } from 'typegpu'; -import { operatorSlot, uniformOpLayout, WORKGROUP_SIZE } from '../schemas.ts'; - -export const uniformOp = tgpu.computeFn({ - workgroupSize: [WORKGROUP_SIZE], - in: { - gid: d.builtin.globalInvocationId, - wid: d.builtin.workgroupId, - }, -})(({ gid, wid }) => { - const globalIdx = gid.x; - const workgroupId = wid.x; - const baseIdx = globalIdx * 8; - const opValue = uniformOpLayout.$.sums[workgroupId]; - - // TODO: use `tgpu.unroll(8)` - for (let i = d.u32(0); i < 8; i++) { - if (baseIdx + i < uniformOpLayout.$.input.length) { - (uniformOpLayout.$.input[baseIdx + i] as number) = operatorSlot.$( - opValue as number, - uniformOpLayout.$.input[baseIdx + i] as number, - ); - } - } -}); diff --git a/packages/typegpu-sort/src/scan/compute/scan.ts b/packages/typegpu-sort/src/scan/compute/scan.ts deleted file mode 100644 index fca23d6c0c..0000000000 --- a/packages/typegpu-sort/src/scan/compute/scan.ts +++ /dev/null @@ -1,72 +0,0 @@ -import { tgpu, d, std } from 'typegpu'; -import { - identitySlot, - onlyGreatestElementSlot, - operatorSlot, - scanLayout, - WORKGROUP_SIZE, -} from '../schemas.ts'; -import { downsweep, upsweep, workgroupMemory } from './shared.ts'; - -const fillIdentityArray = tgpu.comptime(() => Array.from({ length: 8 }, () => identitySlot.$)); - -export const computeBlock = tgpu.computeFn({ - workgroupSize: [WORKGROUP_SIZE], - in: { - gid: d.builtin.globalInvocationId, - lid: d.builtin.localInvocationId, - wid: d.builtin.workgroupId, - }, -})(({ gid, lid, wid }) => { - const globalIdx = gid.x; - const workgroupId = wid.x; - const localIdx = lid.x; - - // 8 elements per thread - const baseIdx = globalIdx * 8; - - const partialSums = d.arrayOf(d.f32, 8)(fillIdentityArray()); - - let prev = d.f32(identitySlot.$); - let lastIdx = d.u32(0); - - // TODO: use `tgpu.unroll(8)` - for (let i = d.u32(); i < 8; i++) { - if (baseIdx + i < scanLayout.$.input.length) { - partialSums[i] = operatorSlot.$(prev, scanLayout.$.input[baseIdx + i] as number); - prev = partialSums[i] as number; - lastIdx = i; - } - } - workgroupMemory.$[localIdx] = partialSums[lastIdx] as number; - - upsweep(localIdx); - - if (localIdx === 0) { - scanLayout.$.sums[workgroupId] = workgroupMemory.$[WORKGROUP_SIZE - 1] as number; - if (!onlyGreatestElementSlot.$) { - workgroupMemory.$[WORKGROUP_SIZE - 1] = d.f32(identitySlot.$); - } - } - - if (!onlyGreatestElementSlot.$) { - downsweep(localIdx); - - std.workgroupBarrier(); - - const scannedSum = workgroupMemory.$[localIdx]; - - for (let i = d.u32(0); i < 8; i++) { - if (baseIdx + i < scanLayout.$.input.length) { - if (i === 0) { - scanLayout.$.input[baseIdx + i] = scannedSum; - } else { - scanLayout.$.input[baseIdx + i] = operatorSlot.$( - scannedSum, - partialSums[i - 1] as number, - ); - } - } - } - } -}); diff --git a/packages/typegpu-sort/src/scan/compute/shared.ts b/packages/typegpu-sort/src/scan/compute/shared.ts deleted file mode 100644 index adbdbe048f..0000000000 --- a/packages/typegpu-sort/src/scan/compute/shared.ts +++ /dev/null @@ -1,35 +0,0 @@ -import { tgpu, d, std } from 'typegpu'; -import { operatorSlot, WORKGROUP_SIZE } from '../schemas.ts'; - -export const workgroupMemory = tgpu.workgroupVar(d.arrayOf(d.f32, WORKGROUP_SIZE)); - -export const upsweep = tgpu.fn([d.u32])((localIdx) => { - let offset = d.u32(1); - for (let d_val = d.u32(WORKGROUP_SIZE / 2); d_val > 0; d_val >>= 1) { - std.workgroupBarrier(); - if (localIdx < d_val) { - const ai = offset * (2 * localIdx + 1) - 1; - const bi = offset * (2 * localIdx + 2) - 1; - workgroupMemory.$[bi] = operatorSlot.$( - workgroupMemory.$[ai] as number, - workgroupMemory.$[bi] as number, - ); - } - offset <<= 1; - } -}); - -export const downsweep = tgpu.fn([d.u32])((localIdx) => { - let offset = d.u32(WORKGROUP_SIZE); - for (let d_val = d.u32(1); d_val < WORKGROUP_SIZE; d_val <<= 1) { - offset >>= 1; - std.workgroupBarrier(); - if (localIdx < d_val) { - const ai = offset * (2 * localIdx + 1) - 1; - const bi = offset * (2 * localIdx + 2) - 1; - const t = workgroupMemory.$[ai] as number; - workgroupMemory.$[ai] = workgroupMemory.$[bi] as number; - workgroupMemory.$[bi] = operatorSlot.$(workgroupMemory.$[bi] as number, t); - } - } -}); diff --git a/packages/typegpu-sort/src/scan/index.ts b/packages/typegpu-sort/src/scan/index.ts index 20fb09d5ed..259afd4e9e 100644 --- a/packages/typegpu-sort/src/scan/index.ts +++ b/packages/typegpu-sort/src/scan/index.ts @@ -1,2 +1,10 @@ -export { prefixScan, scan, createPrefixScanComputer, PrefixScanComputer } from './prefixScan.ts'; +export { + createPrefixScanComputer, + type PrefixScanComputer, + type PrefixScanPlan, + prefixScan, + reduce, + type ScanBuffer, +} from './prefixScan.ts'; +export type { ScanElementType } from './schemas.ts'; export type { BinaryOp } from './types.ts'; diff --git a/packages/typegpu-sort/src/scan/kernels.ts b/packages/typegpu-sort/src/scan/kernels.ts new file mode 100644 index 0000000000..e379ff6be9 --- /dev/null +++ b/packages/typegpu-sort/src/scan/kernels.ts @@ -0,0 +1,116 @@ +import { tgpu, d, std } from 'typegpu'; +import { dispatchIn, flatWorkgroupIndex } from '../dispatch.ts'; +import { ELEMENTS_PER_THREAD, type ScanSchemas, WORKGROUP_SIZE } from './schemas.ts'; + +export function makeScanKernel(schemas: ScanSchemas, identityElement: number) { + const { elementType, scanLayout, identitySlot, reduceOnlySlot, operatorSlot, workgroupMemory } = + schemas; + + function upsweep(localIdx: number) { + 'use gpu'; + let offset = d.u32(1); + for (let span = d.u32(WORKGROUP_SIZE / 2); span > 0; span >>= 1) { + std.workgroupBarrier(); + if (localIdx < span) { + const ai = offset * (2 * localIdx + 1) - 1; + const bi = offset * (2 * localIdx + 2) - 1; + workgroupMemory.$[bi] = operatorSlot.$( + workgroupMemory.$[ai] as number, + workgroupMemory.$[bi] as number, + ); + } + offset <<= 1; + } + } + + function downsweep(localIdx: number) { + 'use gpu'; + let offset = d.u32(WORKGROUP_SIZE); + for (let span = d.u32(1); span < WORKGROUP_SIZE; span <<= 1) { + offset >>= 1; + std.workgroupBarrier(); + if (localIdx < span) { + const ai = offset * (2 * localIdx + 1) - 1; + const bi = offset * (2 * localIdx + 2) - 1; + const t = workgroupMemory.$[ai] as number; + workgroupMemory.$[ai] = workgroupMemory.$[bi] as number; + workgroupMemory.$[bi] = operatorSlot.$(workgroupMemory.$[bi] as number, t); + } + } + } + + const identityArray = Array.from({ length: ELEMENTS_PER_THREAD }, () => identityElement); + + return tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })( + ({ lid, wid, numWorkgroups }) => { + const workgroupId = flatWorkgroupIndex(wid, numWorkgroups); + const localIdx = lid.x; + const baseIdx = (workgroupId * WORKGROUP_SIZE + localIdx) * ELEMENTS_PER_THREAD; + + const partialSums = d.arrayOf(elementType, ELEMENTS_PER_THREAD)(identityArray); + + let prev = identitySlot.$; + let lastIdx = d.u32(0); + + 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]; + lastIdx = i; + } + } + workgroupMemory.$[localIdx] = partialSums[lastIdx] as number; + + upsweep(localIdx); + + if (localIdx === 0 && workgroupId < scanLayout.$.sums.length) { + scanLayout.$.sums[workgroupId] = workgroupMemory.$[WORKGROUP_SIZE - 1] as number; + if (!reduceOnlySlot.$) { + workgroupMemory.$[WORKGROUP_SIZE - 1] = identitySlot.$; + } + } + + if (!reduceOnlySlot.$) { + downsweep(localIdx); + + std.workgroupBarrier(); + + const scannedSum = workgroupMemory.$[localIdx]; + + 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; + } else { + scanLayout.$.input[baseIdx + i] = operatorSlot.$( + scannedSum, + partialSums[i - 1] as number, + ); + } + } + } + } + }, + ); +} + +export function makeApplySumsKernel(schemas: ScanSchemas) { + const { operatorSlot, applySumsLayout } = schemas; + + return tgpu.computeFn({ workgroupSize: [WORKGROUP_SIZE], in: dispatchIn })( + ({ lid, wid, numWorkgroups }) => { + const workgroupId = flatWorkgroupIndex(wid, numWorkgroups); + const baseIdx = (workgroupId * WORKGROUP_SIZE + lid.x) * ELEMENTS_PER_THREAD; + const blockSum = applySumsLayout.$.sums[workgroupId]; + + for (const i of tgpu.unroll(std.range(ELEMENTS_PER_THREAD))) { + if (baseIdx + i < applySumsLayout.$.input.length) { + (applySumsLayout.$.input[baseIdx + i] as number) = operatorSlot.$( + blockSum as number, + applySumsLayout.$.input[baseIdx + i] as number, + ); + } + } + }, + ); +} diff --git a/packages/typegpu-sort/src/scan/prefixScan.ts b/packages/typegpu-sort/src/scan/prefixScan.ts index b5797de946..ce17893861 100644 --- a/packages/typegpu-sort/src/scan/prefixScan.ts +++ b/packages/typegpu-sort/src/scan/prefixScan.ts @@ -1,340 +1,304 @@ import { + d, type StorageFlag, type TgpuBuffer, type TgpuComputePipeline, - type TgpuQuerySet, type TgpuRoot, - d, } from 'typegpu'; +import { decomposeWorkgroups } from '../dispatch.ts'; +import { beginRunPass, bindPass } from '../runPass.ts'; +import type { RunOptions } from '../types.ts'; +import { makeApplySumsKernel, makeScanKernel } from './kernels.ts'; +import { BLOCK_SIZE, makeScanSchemas, type ScanElementType } from './schemas.ts'; import type { BinaryOp } from './types.ts'; -import { - identitySlot, - onlyGreatestElementSlot, - operatorSlot, - scanLayout, - uniformOpLayout, - WORKGROUP_SIZE, -} from './schemas.ts'; -import { computeBlock } from './compute/scan.ts'; -import { uniformOp } from './compute/applySums.ts'; - -const cache = new WeakMap< - TgpuRoot, - WeakMap> ->(); - -export class PrefixScanComputer { - #scanPipeline?: TgpuComputePipeline; - #reducePipeline?: TgpuComputePipeline; - #opPipeline?: TgpuComputePipeline; - #root: TgpuRoot; - #operation: BinaryOp['operation']; - #identityElement: BinaryOp['identityElement']; - - constructor( - root: TgpuRoot, - operation: BinaryOp['operation'], - identityElement: BinaryOp['identityElement'], - ) { - this.#root = root; - this.#operation = operation; - this.#identityElement = identityElement; - } - private getScanPipeline(onlyGreatestElement: boolean): TgpuComputePipeline { - const cached = onlyGreatestElement ? this.#reducePipeline : this.#scanPipeline; +export type ScanBuffer = TgpuBuffer< + d.WgslArray +> & + StorageFlag; - if (cached) { - return cached; - } +type AnyScanBuffer = TgpuBuffer> & StorageFlag; - const pipeline = this.#root - .with(operatorSlot, this.#operation) - .with(identitySlot, this.#identityElement) - .with(onlyGreatestElementSlot, onlyGreatestElement) - .createComputePipeline({ - compute: computeBlock, - }); +interface PlanStep { + pipeline: TgpuComputePipeline; + workgroups: [number, number, number]; +} - if (onlyGreatestElement) { - this.#reducePipeline = pipeline; - } else { - this.#scanPipeline = pipeline; - } +/** + * A reusable execution plan for scanning a specific buffer. All scratch buffers and + * bind groups are created once at `prepare` time, so `run` only records dispatches. + */ +export interface PrefixScanPlan { + /** + * The buffer holding the result after `run`. For a full prefix scan this is the scanned + * buffer itself, for a reduction it is a single-element buffer owned by the plan and + * reused across runs. + */ + readonly resultBuffer: ScanBuffer; + /** Dispatches the scan. Can be called repeatedly */ + run(options?: RunOptions): void; + /** Destroys the scratch buffers owned by this plan */ + destroy(): void; +} + +export interface PrefixScanComputer { + /** Creates a reusable execution plan for scanning `buffer` */ + prepare( + buffer: ScanBuffer, + options?: { reduceOnly?: boolean }, + ): PrefixScanPlan; + /** Scans `buffer` in place through a plan cached per buffer */ + scan(buffer: ScanBuffer, options?: RunOptions): ScanBuffer; + /** + * Reduces `buffer` through a plan cached per buffer. The returned single-element buffer + * belongs to that plan, so it is shared between calls on the same input buffer. + */ + reduce(buffer: ScanBuffer, options?: RunOptions): ScanBuffer; +} + +function makeComputer( + root: TgpuRoot, + operation: BinaryOp['operation'], + identityElement: BinaryOp['identityElement'], + elementType: TElement, +): PrefixScanComputer { + const schemas = makeScanSchemas(elementType); + const scanKernel = makeScanKernel(schemas, identityElement); + const applySumsKernel = makeApplySumsKernel(schemas); + const withOperation = root.with(schemas.operatorSlot, operation); + + const plans = new WeakMap< + ScanBuffer, + { scan?: PrefixScanPlan; reduce?: PrefixScanPlan } + >(); - return pipeline; + let scanPipeline: TgpuComputePipeline | undefined; + let reducePipeline: TgpuComputePipeline | undefined; + let applySumsPipeline: TgpuComputePipeline | undefined; + + function createScanPipeline(reduceOnly: boolean): TgpuComputePipeline { + return withOperation + .with(schemas.identitySlot, identityElement) + .with(schemas.reduceOnlySlot, reduceOnly) + .createComputePipeline({ compute: scanKernel }); } - private get opPipeline(): TgpuComputePipeline { - this.#opPipeline ??= this.#root.with(operatorSlot, this.#operation).createComputePipeline({ - compute: uniformOp, - }); - return this.#opPipeline; + function scanPipelineFor(reduceOnly: boolean): TgpuComputePipeline { + if (reduceOnly) { + reducePipeline ??= createScanPipeline(true); + return reducePipeline; + } + scanPipeline ??= createScanPipeline(false); + return scanPipeline; } - private getScratchBuffer(size: number): TgpuBuffer> & StorageFlag { - return this.#root.createBuffer(d.arrayOf(d.f32, size)).$usage('storage'); + function applySums(): TgpuComputePipeline { + applySumsPipeline ??= withOperation.createComputePipeline({ compute: applySumsKernel }); + return applySumsPipeline; } - private recursiveScan( - buffer: TgpuBuffer> & StorageFlag, - actualLength: number, - onlyGreatestElement: boolean, - querySet: TgpuQuerySet<'timestamp'> | null, - isFirstPass: boolean, - ): TgpuBuffer> & StorageFlag { - const numWorkgroups = Math.ceil(actualLength / (WORKGROUP_SIZE * 8)); - const scanPipeline = this.getScanPipeline(onlyGreatestElement); - - // Base case: single workgroup - if (numWorkgroups === 1) { - const finalSums = this.getScratchBuffer(1); - const bg = this.#root.createBindGroup(scanLayout, { - input: buffer, - sums: finalSums, + function prepare( + buffer: ScanBuffer, + options?: { reduceOnly?: boolean }, + ): PrefixScanPlan { + if (buffer.dataType.elementCount === 0) { + throw new Error('Cannot scan an empty buffer.'); + } + + const reduceOnly = options?.reduceOnly ?? false; + const pipeline = scanPipelineFor(reduceOnly); + + const steps: PlanStep[] = []; + const scratchBuffers: ScanBuffer[] = []; + const applyLevels: { + target: ScanBuffer; + sums: ScanBuffer; + numWorkgroups: number; + }[] = []; + + let currentBuffer = buffer; + let currentLength = buffer.dataType.elementCount; + let resultBuffer = buffer; + + for (;;) { + const numWorkgroups = Math.ceil(currentLength / BLOCK_SIZE); + const sums = root + .createBuffer(d.arrayOf(schemas.elementType, numWorkgroups)) + .$usage('storage') as ScanBuffer; + scratchBuffers.push(sums); + + steps.push({ + pipeline: pipeline.with( + root.createBindGroup(schemas.scanLayout, { + input: currentBuffer as AnyScanBuffer, + sums: sums as AnyScanBuffer, + }), + ), + workgroups: decomposeWorkgroups(numWorkgroups), }); - let pipeline = scanPipeline.with(bg); - if (querySet) { - pipeline = pipeline.withTimestampWrites({ - querySet, - ...(isFirstPass && { beginningOfPassWriteIndex: 0 }), - endOfPassWriteIndex: 1, - }); + + if (numWorkgroups === 1) { + if (reduceOnly) { + resultBuffer = sums; + } + break; } - pipeline.dispatchWorkgroups(1); - return onlyGreatestElement ? finalSums : buffer; + applyLevels.push({ target: currentBuffer, sums, numWorkgroups }); + currentBuffer = sums; + currentLength = numWorkgroups; } - // Recursive case: - let sumsBuffer = this.getScratchBuffer(numWorkgroups); - - const scanBg = this.#root.createBindGroup(scanLayout, { - input: buffer, - sums: sumsBuffer, - }); - let pipeline = scanPipeline.with(scanBg); - if (querySet && isFirstPass) { - pipeline = pipeline.withTimestampWrites({ - querySet, - beginningOfPassWriteIndex: 0, - }); + if (!reduceOnly) { + applyLevels.reverse(); + for (const level of applyLevels) { + steps.push({ + pipeline: applySums().with( + root.createBindGroup(schemas.applySumsLayout, { + input: level.target as AnyScanBuffer, + sums: level.sums as AnyScanBuffer, + }), + ), + workgroups: decomposeWorkgroups(level.numWorkgroups), + }); + } } - pipeline.dispatchWorkgroups(numWorkgroups); - - // Recursively scan the sums - sumsBuffer = this.recursiveScan( - sumsBuffer, - numWorkgroups, - onlyGreatestElement, - querySet, - false, - ); - - if (onlyGreatestElement) { - return sumsBuffer; + + return { + resultBuffer, + + run(options?: RunOptions): void { + const recording = beginRunPass(root.device, options); + for (const step of steps) { + bindPass(step.pipeline, recording.pass).dispatchWorkgroups(...step.workgroups); + } + recording.finish(); + }, + + destroy(): void { + for (const scratch of scratchBuffers) { + scratch.destroy(); + } + }, + }; + } + + function cachedPlan(buffer: ScanBuffer, reduceOnly: boolean): PrefixScanPlan { + let forBuffer = plans.get(buffer); + if (!forBuffer) { + forBuffer = {}; + plans.set(buffer, forBuffer); } - const opBg = this.#root.createBindGroup(uniformOpLayout, { - input: buffer, - sums: sumsBuffer, - }); - let opPipeline = this.opPipeline.with(opBg); - if (querySet) { - opPipeline = opPipeline.withTimestampWrites({ - querySet, - endOfPassWriteIndex: 1, - }); + const key = reduceOnly ? 'reduce' : 'scan'; + let plan = forBuffer[key]; + if (!plan) { + plan = prepare(buffer, { reduceOnly }); + forBuffer[key] = plan; } - opPipeline.dispatchWorkgroups(numWorkgroups); - return buffer; + return plan; } - compute( - buffer: TgpuBuffer> & StorageFlag, - onlyGreatestElement: boolean, - querySet?: TgpuQuerySet<'timestamp'>, - ): TgpuBuffer> & StorageFlag { - return this.recursiveScan( - buffer, - buffer.dataType.elementCount, - onlyGreatestElement, - querySet ?? null, - true, - ); + return { + prepare, + + scan(buffer, options) { + const plan = cachedPlan(buffer, false); + plan.run(options); + return plan.resultBuffer; + }, + + reduce(buffer, options) { + const plan = cachedPlan(buffer, true); + plan.run(options); + return plan.resultBuffer; + }, + }; +} + +interface CacheLike { + get(key: K): V | undefined; + set(key: K, value: V): unknown; +} + +function getOrCreate(cache: CacheLike, key: K, create: () => V): V { + const cached = cache.get(key); + if (cached !== undefined) { + return cached; } + + const created = create(); + cache.set(key, created); + return created; } +const computerCache = new WeakMap>>(); + /** - * Perform a GPU prefix-scan (parallel prefix scan depending on the - * provided operation) over the values in `inputBuffer`. For instance, this can be used to - * compute a prefix sum over an array of numbers. - * - * @param root - The TypeGPU root/context used to create pipelines, bind groups and buffers. - * @param options - Configuration object containing: - * - inputBuffer: A storage buffer with the input values to scan - * - outputBuffer: (optional) A storage buffer where the scanned values will be written. - * Defaults to in-place (overwrites `inputBuffer`). - * - operation: The binary operation to use for the scan (e.g., std.add) - * - identityElement: The identity element for the operation (e.g., 0 for addition) - * @param querySet - Optional timestamp query set (size >= 2) for GPU timing. - * Index 0 gets the begin timestamp, index 1 gets the end timestamp. - * @returns The output buffer instance which contains the scanned values. - * - * @example - * ```typescript - * const root = await tgpu.init(); - * const inputBuffer = root - * .createBuffer(d.arrayOf(d.f32, 4), [1, 2, 3, 4]) - * .$usage('storage'); - * - * // in-place (inputBuffer is modified) - * const result = prefixScan( - * root, - * { - * inputBuffer, - * operation: std.add, - * identityElement: 0, - * }, - * ); - * - * // with separate output buffer - * const outputBuffer = root - * .createBuffer(d.arrayOf(d.f32, 4)) - * .$usage('storage'); - * - * const result = prefixScan( - * root, - * { - * inputBuffer, - * outputBuffer, - * operation: std.add, - * identityElement: 0, - * }, - * ); - * ``` + * Creates a computer for the given operation, reusing the one cached for the same `root` and + * `binaryOp` so that repeated calls share pipelines. Set `dataType` to `d.u32` or `d.i32` to + * scan integer buffers (defaults to `d.f32`). */ -export function prefixScan( +export function createPrefixScanComputer( root: TgpuRoot, - options: { - inputBuffer: TgpuBuffer> & StorageFlag; - outputBuffer?: TgpuBuffer> & StorageFlag; - operation: BinaryOp['operation']; - identityElement: BinaryOp['identityElement']; - }, - querySet?: TgpuQuerySet<'timestamp'>, -): TgpuBuffer> & StorageFlag { - return runScan(root, options, false, querySet); + binaryOp: BinaryOp, +): PrefixScanComputer { + const elementType = (binaryOp.dataType ?? d.f32) as TElement; + const byOperation = getOrCreate(computerCache, root, () => new WeakMap()); + const byElement = getOrCreate(byOperation, binaryOp.operation, () => new Map()); + + return getOrCreate(byElement, `${binaryOp.identityElement}_${elementType.type}`, () => + makeComputer(root, binaryOp.operation, binaryOp.identityElement, elementType), + ) as PrefixScanComputer; } -/** - * Compute only the aggregated reduction result for `inputBuffer` using the provided operation. - * Returns only the top-level sums/reductions instead of the full scan. This is useful when - * you only need the final reduction - for instance, the sum of the whole array. - * - * @param root - The TypeGPU root/context used to create pipelines, bind groups and buffers. - * @param options - Configuration object containing: - * - inputBuffer: A storage buffer with the input values to reduce - * - operation: The binary operation to use for the reduction (e.g., std.add) - * - identityElement: The identity element for the operation (e.g., 0 for addition) - * @param querySet - Optional timestamp query set (size >= 2) for GPU timing. - * Index 0 gets the begin timestamp, index 1 gets the end timestamp. - * @returns A buffer containing the aggregated reduction result (single-element buffer). - * - * @example - * ```typescript - * const root = await tgpu.init(); - * const inputBuffer = root - * .createBuffer(d.arrayOf(d.f32, 4), [1, 2, 3, 4]) - * .$usage('storage'); - * - * // using an std function - * const result = scan( - * root, - * { - * inputBuffer, - * operation: std.add, - * identityElement: 0, - * }, - * ); - * - * // using a custom tgpu.fn - * const multiply = tgpu.fn([d.f32, d.f32], d.f32)((a, b) => a * b); - * - * const result = scan( - * root, - * { - * inputBuffer, - * operation: multiply, - * identityElement: 1, - * }, - * ); - * ``` - */ -export function scan( - root: TgpuRoot, - options: { - inputBuffer: TgpuBuffer> & StorageFlag; - operation: BinaryOp['operation']; - identityElement: BinaryOp['identityElement']; - }, - querySet?: TgpuQuerySet<'timestamp'>, -): TgpuBuffer> & StorageFlag { - return runScan(root, options, true, querySet); +interface ScanOptions { + inputBuffer: ScanBuffer; + operation: BinaryOp['operation']; + identityElement: BinaryOp['identityElement']; } -function runScan( +/** + * Performs an exclusive prefix scan over `inputBuffer` with the given associative operation, + * writing to `outputBuffer` if provided and in place otherwise. + */ +export function prefixScan( root: TgpuRoot, - options: { - inputBuffer: TgpuBuffer> & StorageFlag; - outputBuffer?: TgpuBuffer> & StorageFlag; - operation: BinaryOp['operation']; - identityElement: BinaryOp['identityElement']; - }, - onlyGreatestElement: boolean, - querySet?: TgpuQuerySet<'timestamp'>, -): TgpuBuffer> & StorageFlag { + options: ScanOptions & { outputBuffer?: ScanBuffer }, +): ScanBuffer { + const { inputBuffer, outputBuffer } = options; const computer = createPrefixScanComputer(root, { operation: options.operation, identityElement: options.identityElement, + dataType: inputBuffer.dataType.elementType, }); - if (onlyGreatestElement) { - return computer.compute(options.inputBuffer, true, querySet); + if (!outputBuffer || outputBuffer === inputBuffer) { + return computer.scan(inputBuffer); } - const outputBuffer = options.outputBuffer ?? options.inputBuffer; - if (options.inputBuffer !== outputBuffer) { - outputBuffer.copyFrom(options.inputBuffer); + if ( + outputBuffer.dataType.elementType.type !== inputBuffer.dataType.elementType.type || + outputBuffer.dataType.elementCount !== inputBuffer.dataType.elementCount + ) { + throw new Error('The input and output scan buffers must have the same type and length.'); } - return computer.compute(outputBuffer, false, querySet); + (outputBuffer as ScanBuffer).copyFrom(inputBuffer as ScanBuffer); + return computer.scan(outputBuffer); } /** - * Create or retrieve a cached `PrefixScanComputer` for the given `root` and `binaryOp`. - * - * @param root - The TypeGPU root/context to associate with the cached computer. - * @param binaryOp - The binary operation used by the computer. - * @returns A `PrefixScanComputer` instance associated with the provided `root` and `binaryOp`. + * Reduces `inputBuffer` with the given associative operation, returning a single-element + * buffer with the aggregate. The input is left untouched. */ -export function createPrefixScanComputer(root: TgpuRoot, binaryOp: BinaryOp): PrefixScanComputer { - let rootCache = cache.get(root); - if (!rootCache) { - rootCache = new WeakMap(); - cache.set(root, rootCache); - } - - let opCache = rootCache.get(binaryOp.operation); - if (!opCache) { - opCache = new Map(); - rootCache.set(binaryOp.operation, opCache); - } - - let computer = opCache.get(binaryOp.identityElement); - if (!computer) { - computer = new PrefixScanComputer(root, binaryOp.operation, binaryOp.identityElement); - opCache.set(binaryOp.identityElement, computer); - } - return computer; +export function reduce( + root: TgpuRoot, + options: ScanOptions, +): ScanBuffer { + return createPrefixScanComputer(root, { + operation: options.operation, + identityElement: options.identityElement, + dataType: options.inputBuffer.dataType.elementType, + }).reduce(options.inputBuffer); } diff --git a/packages/typegpu-sort/src/scan/schemas.ts b/packages/typegpu-sort/src/scan/schemas.ts index f8709b4e4c..8507fe6405 100644 --- a/packages/typegpu-sort/src/scan/schemas.ts +++ b/packages/typegpu-sort/src/scan/schemas.ts @@ -1,17 +1,27 @@ import { tgpu, d } from 'typegpu'; export const WORKGROUP_SIZE = 256; +export const ELEMENTS_PER_THREAD = 8; +export const BLOCK_SIZE = WORKGROUP_SIZE * ELEMENTS_PER_THREAD; -export const scanLayout = tgpu.bindGroupLayout({ - input: { storage: d.arrayOf(d.f32), access: 'mutable' }, - sums: { storage: d.arrayOf(d.f32), access: 'mutable' }, -}); +export type ScanElementType = d.F32 | d.U32 | d.I32; -export const uniformOpLayout = tgpu.bindGroupLayout({ - input: { storage: d.arrayOf(d.f32), access: 'mutable' }, - sums: { storage: d.arrayOf(d.f32), access: 'readonly' }, -}); +export function makeScanSchemas(elementType: ScanElementType) { + return { + elementType, + scanLayout: tgpu.bindGroupLayout({ + input: { storage: d.arrayOf(elementType), access: 'mutable' }, + sums: { storage: d.arrayOf(elementType), access: 'mutable' }, + }), + applySumsLayout: tgpu.bindGroupLayout({ + input: { storage: d.arrayOf(elementType), access: 'mutable' }, + sums: { storage: d.arrayOf(elementType), access: 'readonly' }, + }), + operatorSlot: tgpu.slot<(a: number, b: number) => number>(), + identitySlot: tgpu.accessor(elementType), + reduceOnlySlot: tgpu.slot(), + workgroupMemory: tgpu.workgroupVar(d.arrayOf(elementType, WORKGROUP_SIZE)), + }; +} -export const operatorSlot = tgpu.slot<(a: number, b: number) => number>(); -export const identitySlot = tgpu.slot(); -export const onlyGreatestElementSlot = tgpu.slot(); +export type ScanSchemas = ReturnType; diff --git a/packages/typegpu-sort/src/scan/types.ts b/packages/typegpu-sort/src/scan/types.ts index 0c3183e5e2..8e753ea197 100644 --- a/packages/typegpu-sort/src/scan/types.ts +++ b/packages/typegpu-sort/src/scan/types.ts @@ -1,4 +1,9 @@ -export interface BinaryOp { +import type { d } from 'typegpu'; +import type { ScanElementType } from './schemas.ts'; + +export interface BinaryOp { operation: (a: number, b: number) => number; identityElement: number; + /** Element type of the buffers to scan. Defaults to `d.f32` */ + dataType?: TElement; } diff --git a/packages/typegpu-sort/src/types.ts b/packages/typegpu-sort/src/types.ts new file mode 100644 index 0000000000..53ee0830da --- /dev/null +++ b/packages/typegpu-sort/src/types.ts @@ -0,0 +1,25 @@ +import type { TgpuCommandEncoder, TgpuComputePass } from 'typegpu'; + +interface EncoderOptions { + /** Records the dispatches as a single compute pass on this encoder. Nothing is submitted */ + encoder: GPUCommandEncoder | TgpuCommandEncoder; + pass?: never; +} + +interface PassOptions { + encoder?: never; + /** Records the dispatches into this pass. Nothing is submitted and the pass is not ended */ + pass: GPUComputePassEncoder | TgpuComputePass; +} + +/** Controls where a `run` call records its dispatches. Defaults to a standalone submit */ +export type RunOptions = EncoderOptions | PassOptions; + +export interface Sorter { + /** Number of elements this sorter was created for */ + readonly size: number; + /** Sorts the buffer in place. Can be called repeatedly */ + run(options?: RunOptions): void; + /** Destroys the internal buffers owned by this sorter */ + destroy(): void; +} diff --git a/packages/typegpu-sort/tests/bitonic.test.ts b/packages/typegpu-sort/tests/bitonic.test.ts new file mode 100644 index 0000000000..bb4865330b --- /dev/null +++ b/packages/typegpu-sort/tests/bitonic.test.ts @@ -0,0 +1,294 @@ +import { tgpu, d } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; +import { describe, expect, vi } from 'vitest'; +import { createBitonicSorter } from '../src/index.ts'; +import { defaultCompare } from '../src/bitonic/slots.ts'; +import { getConversionWarnings, getResolvedWgsl } from './utils.ts'; + +describe('bitonic sort', () => { + it('emits no implicit conversion warnings for any key type', ({ root }) => { + const warnSpy = vi.spyOn(console, 'warn'); + + for (const keyType of [d.u32, d.i32, d.f32] as const) { + const data = root.createBuffer(d.arrayOf(keyType, 256)).$usage('storage'); + createBitonicSorter(root, data).run(); + } + + expect(getConversionWarnings(warnSpy)).toMatchInlineSnapshot(`[]`); + warnSpy.mockRestore(); + }); + + it('specializes the comparator per key type', () => { + expect( + tgpu.resolve([ + tgpu.fn([d.u32, d.u32], d.bool)(defaultCompare), + tgpu.fn([d.i32, d.i32], d.bool)(defaultCompare), + tgpu.fn([d.f32, d.f32], d.bool)(defaultCompare), + ]), + ).toMatchInlineSnapshot(` + "fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn defaultCompare_1(a: i32, b: i32) -> bool { + return (a < b); + } + + fn defaultCompare_2(a: f32, b: f32) -> bool { + return (a < b); + }" + `); + }); + + it('should produce valid code for a composite payload', ({ root, device }) => { + const keys = root.createBuffer(d.arrayOf(d.u32, 1024)).$usage('storage'); + const values = root.createBuffer(d.arrayOf(d.vec2f, 1024)).$usage('storage'); + createBitonicSorter(root, keys, { values }).run(); + + expect(getResolvedWgsl(device)).toMatchInlineSnapshot(` + "fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(0) var data: array; + + var localKeys: array; + + var localVals: array; + + @group(1) @binding(0) var vals: array; + + fn loadShared(base: u32, tid: u32) { + localKeys[tid] = data[(base + tid)]; + localKeys[(tid + 256u)] = data[((base + tid) + 256u)]; + { + localVals[tid] = vals[(base + tid)]; + localVals[(tid + 256u)] = vals[((base + tid) + 256u)]; + } + } + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn swapLocalAt(a: u32, b: u32, left: u32, right: u32) { + localKeys[a] = right; + localKeys[b] = left; + { + let tmp = localVals[a]; + localVals[a] = localVals[b]; + localVals[b] = tmp; + } + } + + fn exchangeLocal(base: u32, iLocal: u32, stride: u32, k: u32) { + let jLocal = (iLocal + stride); + let left = localKeys[iLocal]; + let right = localKeys[jLocal]; + let ascending = (((base + iLocal) & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapLocalAt(iLocal, jLocal, left, right); + } + } + + fn mergeDown(base: u32, tid: u32, startShift: u32, k: u32) { + for (var jShift = startShift; (jShift > 0u); jShift--) { + workgroupBarrier(); + let stride = (1u << (jShift - 1u)); + let below = (tid & (stride - 1u)); + let above = (tid >> (jShift - 1u)); + exchangeLocal(base, (below + (above * (stride << 1u))), stride, k); + } + } + + fn storeShared(base: u32, tid: u32) { + data[(base + tid)] = localKeys[tid]; + data[((base + tid) + 256u)] = localKeys[(tid + 256u)]; + { + vals[(base + tid)] = localVals[tid]; + vals[((base + tid) + 256u)] = localVals[(tid + 256u)]; + } + } + + @compute @workgroup_size(256) fn localSort(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let base = (flatWorkgroupIndex(wid, numWorkgroups) * 512u); + if ((base >= arrayLength(&data))) { + return; + } + loadShared(base, lid.x); + for (var kShift = 1u; (kShift <= 9u); kShift++) { + mergeDown(base, lid.x, kShift, (1u << kShift)); + } + workgroupBarrier(); + storeShared(base, lid.x); + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + struct sortUniformsType { + k: u32, + jShift: u32, + } + + @group(0) @binding(1) var uniforms: sortUniformsType; + + @group(0) @binding(0) var data: array; + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + @group(1) @binding(0) var vals: array; + + fn swapAt(i: u32, j: u32, left: u32, right: u32) { + data[i] = right; + data[j] = left; + { + let tmp = vals[i]; + vals[i] = vals[j]; + vals[j] = tmp; + } + } + + @compute @workgroup_size(256) fn item(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let tid = ((flatWorkgroupIndex(wid, numWorkgroups) * 256u) + lid.x); + let k = uniforms.k; + let shift = uniforms.jShift; + let stride = (1u << shift); + let below = (tid & (stride - 1u)); + let above = (tid >> shift); + let i = (below + (above * (stride << 1u))); + let ixj = (i + stride); + if ((ixj >= arrayLength(&data))) { + return; + } + let left = data[i]; + let right = data[ixj]; + let ascending = ((i & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapAt(i, ixj, left, right); + } + } + + fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(0) var data: array; + + var localKeys: array; + + var localVals: array; + + @group(1) @binding(0) var vals: array; + + fn loadShared(base: u32, tid: u32) { + localKeys[tid] = data[(base + tid)]; + localKeys[(tid + 256u)] = data[((base + tid) + 256u)]; + { + localVals[tid] = vals[(base + tid)]; + localVals[(tid + 256u)] = vals[((base + tid) + 256u)]; + } + } + + struct sortUniformsType { + k: u32, + jShift: u32, + } + + @group(0) @binding(1) var uniforms: sortUniformsType; + + fn defaultCompare(a: u32, b: u32) -> bool { + return (a < b); + } + + fn swapLocalAt(a: u32, b: u32, left: u32, right: u32) { + localKeys[a] = right; + localKeys[b] = left; + { + let tmp = localVals[a]; + localVals[a] = localVals[b]; + localVals[b] = tmp; + } + } + + fn exchangeLocal(base: u32, iLocal: u32, stride: u32, k: u32) { + let jLocal = (iLocal + stride); + let left = localKeys[iLocal]; + let right = localKeys[jLocal]; + let ascending = (((base + iLocal) & k) == 0u); + if (select(defaultCompare(left, right), defaultCompare(right, left), ascending)) { + swapLocalAt(iLocal, jLocal, left, right); + } + } + + fn mergeDown(base: u32, tid: u32, startShift: u32, k: u32) { + for (var jShift = startShift; (jShift > 0u); jShift--) { + workgroupBarrier(); + let stride = (1u << (jShift - 1u)); + let below = (tid & (stride - 1u)); + let above = (tid >> (jShift - 1u)); + exchangeLocal(base, (below + (above * (stride << 1u))), stride, k); + } + } + + fn storeShared(base: u32, tid: u32) { + data[(base + tid)] = localKeys[tid]; + data[((base + tid) + 256u)] = localKeys[(tid + 256u)]; + { + vals[(base + tid)] = localVals[tid]; + vals[((base + tid) + 256u)] = localVals[(tid + 256u)]; + } + } + + @compute @workgroup_size(256) fn localMerge(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let base = (flatWorkgroupIndex(wid, numWorkgroups) * 512u); + if ((base >= arrayLength(&data))) { + return; + } + loadShared(base, lid.x); + mergeDown(base, lid.x, 9u, uniforms.k); + workgroupBarrier(); + storeShared(base, lid.x); + }" + `); + }); + + it('performs no buffer writes or allocations during run (uniforms precreated)', ({ + root, + device, + }) => { + const data = root.createBuffer(d.arrayOf(d.u32, 1024)).$usage('storage'); + const sorter = createBitonicSorter(root, data); + + sorter.run(); + const writesAfterFirst = (device.mock.queue.writeBuffer as { mock: { calls: unknown[] } }).mock + .calls.length; + const buffersAfterFirst = device.mock.createBuffer.mock.calls.length; + + sorter.run(); + + expect( + (device.mock.queue.writeBuffer as { mock: { calls: unknown[] } }).mock.calls.length, + ).toBe(writesAfterFirst); + expect(device.mock.createBuffer.mock.calls.length).toBe(buffersAfterFirst); + }); + + it('rejects payload sorting when padding would be required', ({ root }) => { + const keys = root.createBuffer(d.arrayOf(d.u32, 3)).$usage('storage'); + const values = root.createBuffer(d.arrayOf(d.u32, 3)).$usage('storage'); + + expect(() => createBitonicSorter(root, keys, { values })).toThrowErrorMatchingInlineSnapshot( + `[Error: Bitonic sorting with a values buffer requires a power-of-two element count.]`, + ); + }); + + it('rejects empty buffers', ({ root }) => { + const keys = root.createBuffer(d.arrayOf(d.u32, 0)).$usage('storage'); + expect(() => createBitonicSorter(root, keys)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a bitonic sorter for an empty buffer.]`, + ); + }); +}); diff --git a/packages/typegpu-sort/tests/radix.test.ts b/packages/typegpu-sort/tests/radix.test.ts new file mode 100644 index 0000000000..2d88154f14 --- /dev/null +++ b/packages/typegpu-sort/tests/radix.test.ts @@ -0,0 +1,107 @@ +import { tgpu, d } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; +import { describe, expect, vi } from 'vitest'; +import { createRadixSorter } from '../src/index.ts'; +import { makeDigitFn, makeRadixSchemas } from '../src/radix/schemas.ts'; +import { getConversionWarnings } from './utils.ts'; + +describe('radix sort', () => { + it('emits no implicit conversion warnings for any key type', ({ root }) => { + const warnSpy = vi.spyOn(console, 'warn'); + + for (const keyType of [d.u32, d.i32, d.f32] as const) { + const data = root.createBuffer(d.arrayOf(keyType, 512)).$usage('storage'); + createRadixSorter(root, data).run(); + } + + expect(getConversionWarnings(warnSpy)).toMatchInlineSnapshot(`[]`); + warnSpy.mockRestore(); + }); + + it('extracts i32 digits without an INT_MIN literal', () => { + const digit = tgpu.fn([d.i32, d.u32], d.u32)(makeDigitFn(d.i32, 'ascending')); + expect(tgpu.resolve([digit])).toMatchInlineSnapshot(` + "fn digitOfI32(v: i32, shift: u32) -> u32 { + let raw = u32(((v >> shift) & 255i)); + return (raw ^ select(0u, 128u, (shift == 24u))); + }" + `); + }); + + it('canonicalizes signed zero before extracting f32 digits', () => { + const digit = tgpu.fn([d.f32, d.u32], d.u32)(makeDigitFn(d.f32, 'ascending')); + expect(tgpu.resolve([digit])).toMatchInlineSnapshot(` + "fn digitOfF32(v: f32, shift: u32) -> u32 { + let bits = select(bitcast(v), 0u, (v == 0f)); + let mask = select(2147483648u, 4294967295u, ((bits >> 31u) == 1u)); + return (((bits ^ mask) >> shift) & 255u); + }" + `); + }); + + it('inverts digits for a descending sort', () => { + const digit = tgpu.fn([d.u32, d.u32], d.u32)(makeDigitFn(d.u32, 'descending')); + expect(tgpu.resolve([digit])).toMatchInlineSnapshot(` + "fn digitOfU32(v: u32, shift: u32) -> u32 { + return ((v >> shift) & 255u); + } + + fn descendingDigit(v: u32, shift: u32) -> u32 { + return (255u - digitOfU32(v, shift)); + }" + `); + }); + + it('writes keys only when no values buffer is provided', () => { + const { writeOutput } = makeRadixSchemas(d.u32, 'ascending'); + expect(tgpu.resolve([tgpu.fn([d.u32, d.u32, d.u32])(writeOutput)])).toMatchInlineSnapshot(` + "@group(0) @binding(1) var dst: array; + + fn writeOutput(key: u32, srcIdx: u32, dstIdx: u32) { + dst[dstIdx] = key; + }" + `); + }); + + it('reorders the payload alongside the keys', () => { + const { writeOutput } = makeRadixSchemas(d.u32, 'ascending', d.vec4f); + expect(tgpu.resolve([tgpu.fn([d.u32, d.u32, d.u32])(writeOutput)])).toMatchInlineSnapshot(` + "@group(0) @binding(1) var dst: array; + + @group(1) @binding(1) var dstVals: array; + + @group(1) @binding(0) var srcVals: array; + + fn writeOutput(key: u32, srcIdx: u32, dstIdx: u32) { + dst[dstIdx] = key; + { + dstVals[dstIdx] = srcVals[srcIdx]; + } + }" + `); + }); + + it('allocates no new GPU resources on repeated runs', ({ root, device }) => { + const data = root.createBuffer(d.arrayOf(d.u32, 4096)).$usage('storage'); + const sorter = createRadixSorter(root, data); + + sorter.run(); + const buffersAfterFirst = device.mock.createBuffer.mock.calls.length; + const bindGroupsAfterFirst = device.mock.createBindGroup.mock.calls.length; + const pipelinesAfterFirst = device.mock.createComputePipeline.mock.calls.length; + + sorter.run(); + sorter.run(); + + expect(device.mock.createBuffer.mock.calls.length).toBe(buffersAfterFirst); + expect(device.mock.createBindGroup.mock.calls.length).toBe(bindGroupsAfterFirst); + expect(device.mock.createComputePipeline.mock.calls.length).toBe(pipelinesAfterFirst); + }); + + it('rejects empty buffers', ({ root }) => { + const keys = root.createBuffer(d.arrayOf(d.u32, 0)).$usage('storage'); + expect(() => createRadixSorter(root, keys)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a radix sorter for an empty buffer.]`, + ); + }); +}); diff --git a/packages/typegpu-sort/tests/scan.test.ts b/packages/typegpu-sort/tests/scan.test.ts new file mode 100644 index 0000000000..0605ac437f --- /dev/null +++ b/packages/typegpu-sort/tests/scan.test.ts @@ -0,0 +1,195 @@ +import { tgpu, d, std } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; +import { describe, expect, vi } from 'vitest'; +import { createPrefixScanComputer, prefixScan } from '../src/index.ts'; +import { makeScanSchemas } from '../src/scan/schemas.ts'; +import { getConversionWarnings, getResolvedWgsl } from './utils.ts'; + +describe('prefix scan', () => { + it('emits no implicit conversion warnings for any element type', ({ root }) => { + const warnSpy = vi.spyOn(console, 'warn'); + + for (const dataType of [d.u32, d.i32, d.f32] as const) { + const buffer = root.createBuffer(d.arrayOf(dataType, 4096)).$usage('storage'); + prefixScan(root, { inputBuffer: buffer, operation: std.add, identityElement: 0 }); + } + + expect(getConversionWarnings(warnSpy)).toMatchInlineSnapshot(`[]`); + warnSpy.mockRestore(); + }); + + it('types the workgroup memory and layouts after the element type', () => { + const { workgroupMemory, scanLayout, applySumsLayout } = makeScanSchemas(d.i32); + expect(tgpu.resolve([workgroupMemory, scanLayout, applySumsLayout])).toMatchInlineSnapshot(` + "var workgroupMemory: array; + + @group(0) @binding(0) var input: array; + + @group(0) @binding(1) var sums: array; + + @group(1) @binding(0) var input_1: array; + + @group(1) @binding(1) var sums_1: array;" + `); + }); + + it('reuses scratch buffers and bind groups across repeated scans', ({ root, device }) => { + const computer = createPrefixScanComputer(root, { + operation: std.add, + identityElement: 0, + dataType: d.u32, + }); + const buffer = root.createBuffer(d.arrayOf(d.u32, 4096)).$usage('storage'); + + computer.scan(buffer); + const buffersAfterFirst = device.mock.createBuffer.mock.calls.length; + const bindGroupsAfterFirst = device.mock.createBindGroup.mock.calls.length; + const modulesAfterFirst = device.mock.createShaderModule.mock.calls.length; + + computer.scan(buffer); + computer.scan(buffer); + + expect(device.mock.createBuffer.mock.calls.length).toBe(buffersAfterFirst); + expect(device.mock.createBindGroup.mock.calls.length).toBe(bindGroupsAfterFirst); + expect(device.mock.createShaderModule.mock.calls.length).toBe(modulesAfterFirst); + }); + + it('should produce valid code for a reduction', ({ root, device }) => { + const computer = createPrefixScanComputer(root, { + operation: std.add, + identityElement: 0, + dataType: d.u32, + }); + const buffer = root.createBuffer(d.arrayOf(d.u32, 4096)).$usage('storage'); + + computer.reduce(buffer); + + expect(getResolvedWgsl(device)).toMatchInlineSnapshot(` + "fn flatWorkgroupIndex(wid: vec3u, numWorkgroups: vec3u) -> u32 { + return ((wid.x + (wid.y * numWorkgroups.x)) + ((wid.z * numWorkgroups.x) * numWorkgroups.y)); + } + + @group(0) @binding(0) var input: array; + + var workgroupMemory: array; + + fn upsweep(localIdx: u32) { + var offset = 1u; + for (var span = 128u; (span > 0u); span >>= 1u) { + workgroupBarrier(); + if ((localIdx < span)) { + let ai = ((offset * ((2u * localIdx) + 1u)) - 1u); + let bi = ((offset * ((2u * localIdx) + 2u)) - 1u); + workgroupMemory[bi] = (workgroupMemory[ai] + workgroupMemory[bi]); + } + offset <<= 1u; + } + } + + @group(0) @binding(1) var sums: array; + + @compute @workgroup_size(256) fn item(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u, @builtin(num_workgroups) numWorkgroups: vec3u) { + let workgroupId = flatWorkgroupIndex(wid, numWorkgroups); + let localIdx = lid.x; + let baseIdx = (((workgroupId * 256u) + localIdx) * 8u); + var partialSums = array(0u, 0u, 0u, 0u, 0u, 0u, 0u, 0u); + var prev = 0u; + var lastIdx = 0u; + // unrolled iteration #0 + { + if (((baseIdx + 0u) < arrayLength(&input))) { + partialSums[0i] = (prev + input[(baseIdx + 0u)]); + prev = partialSums[0i]; + lastIdx = 0u; + } + } + // unrolled iteration #1 + { + if (((baseIdx + 1u) < arrayLength(&input))) { + partialSums[1i] = (prev + input[(baseIdx + 1u)]); + prev = partialSums[1i]; + lastIdx = 1u; + } + } + // unrolled iteration #2 + { + if (((baseIdx + 2u) < arrayLength(&input))) { + partialSums[2i] = (prev + input[(baseIdx + 2u)]); + prev = partialSums[2i]; + lastIdx = 2u; + } + } + // unrolled iteration #3 + { + if (((baseIdx + 3u) < arrayLength(&input))) { + partialSums[3i] = (prev + input[(baseIdx + 3u)]); + prev = partialSums[3i]; + lastIdx = 3u; + } + } + // unrolled iteration #4 + { + if (((baseIdx + 4u) < arrayLength(&input))) { + partialSums[4i] = (prev + input[(baseIdx + 4u)]); + prev = partialSums[4i]; + lastIdx = 4u; + } + } + // unrolled iteration #5 + { + if (((baseIdx + 5u) < arrayLength(&input))) { + partialSums[5i] = (prev + input[(baseIdx + 5u)]); + prev = partialSums[5i]; + lastIdx = 5u; + } + } + // unrolled iteration #6 + { + if (((baseIdx + 6u) < arrayLength(&input))) { + partialSums[6i] = (prev + input[(baseIdx + 6u)]); + prev = partialSums[6i]; + lastIdx = 6u; + } + } + // unrolled iteration #7 + { + if (((baseIdx + 7u) < arrayLength(&input))) { + partialSums[7i] = (prev + input[(baseIdx + 7u)]); + prev = partialSums[7i]; + lastIdx = 7u; + } + } + workgroupMemory[localIdx] = partialSums[lastIdx]; + upsweep(localIdx); + if (((localIdx == 0u) && (workgroupId < arrayLength(&sums)))) { + sums[workgroupId] = workgroupMemory[255i]; + } + }" + `); + expect(device.mock.createComputePipeline.mock.calls.length).toMatchInlineSnapshot(`1`); + }); + + it('rejects empty and mismatched buffers', ({ root }) => { + const computer = createPrefixScanComputer(root, { + operation: std.add, + identityElement: 0, + }); + const empty = root.createBuffer(d.arrayOf(d.f32, 0)).$usage('storage'); + expect(() => computer.prepare(empty)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot scan an empty buffer.]`, + ); + + const input = root.createBuffer(d.arrayOf(d.u32, 4)).$usage('storage'); + const output = root.createBuffer(d.arrayOf(d.u32, 3)).$usage('storage'); + expect(() => + prefixScan(root, { + inputBuffer: input, + outputBuffer: output, + operation: std.add, + identityElement: 0, + }), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: The input and output scan buffers must have the same type and length.]`, + ); + }); +}); diff --git a/packages/typegpu-sort/tests/utils.ts b/packages/typegpu-sort/tests/utils.ts new file mode 100644 index 0000000000..9fc66620d8 --- /dev/null +++ b/packages/typegpu-sort/tests/utils.ts @@ -0,0 +1,11 @@ +export function getResolvedWgsl(device: { + mock: { createShaderModule: { mock: { calls: unknown[][] } } }; +}): string { + return device.mock.createShaderModule.mock.calls + .map((call) => (call[0] as { code: string }).code) + .join('\n\n'); +} + +export function getConversionWarnings(warnSpy: { mock: { calls: unknown[][] } }): unknown[][] { + return warnSpy.mock.calls.filter((call) => String(call[0]).includes('Implicit conversions')); +} diff --git a/packages/typegpu-sort/vitest.config.mts b/packages/typegpu-sort/vitest.config.mts new file mode 100644 index 0000000000..5cc441af79 --- /dev/null +++ b/packages/typegpu-sort/vitest.config.mts @@ -0,0 +1,14 @@ +import { createJiti } from 'jiti'; +import { typegpuBuiltAliases } from 'typegpu-testing-utility/config'; +import type TypeGPUPlugin from 'unplugin-typegpu/vite'; +import { defineConfig } from 'vitest/config'; + +const jiti = createJiti(import.meta.url); +const typegpu = await jiti.import('unplugin-typegpu/vite', { default: true }); + +export default defineConfig({ + plugins: [typegpu({ forceTgpuAlias: 'tgpu', earlyPruning: false })], + resolve: { + alias: typegpuBuiltAliases(), + }, +}); diff --git a/packages/typegpu-testing-utility/src/extendedIt.ts b/packages/typegpu-testing-utility/src/extendedIt.ts index ebf60d52a9..d39918f9e1 100644 --- a/packages/typegpu-testing-utility/src/extendedIt.ts +++ b/packages/typegpu-testing-utility/src/extendedIt.ts @@ -157,6 +157,7 @@ export const it = base limits: { maxUniformBuffersPerShaderStage: 12, maxStorageBuffersPerShaderStage: 8, + maxComputeWorkgroupStorageSize: 16384, }, destroy: vi.fn(), }; diff --git a/packages/typegpu/tests/accessor.test.ts b/packages/typegpu/tests/accessor.test.ts index 284d7a07c5..c0400af72b 100644 --- a/packages/typegpu/tests/accessor.test.ts +++ b/packages/typegpu/tests/accessor.test.ts @@ -540,6 +540,15 @@ describe('tgpu.accessor', () => { ); }); + it('throws when $ is accessed in simulation mode', () => { + const valueAccess = tgpu.accessor(d.f32, 1); + expect(() => + tgpu['~unstable'].simulate(() => valueAccess.$), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: \`tgpu.accessor\` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call. Use \`tgpu.slot\` for non-WGSL values instead.]`, + ); + }); + it('allows for arbitrarily nested access functions', ({ root }) => { const counterMutable = root.createMutable(d.u32); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 26295c0a2e..56a659697a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -937,6 +937,9 @@ importers: typegpu: specifier: workspace:* version: link:../typegpu + typegpu-testing-utility: + specifier: workspace:* + version: link:../typegpu-testing-utility typescript: specifier: npm:tsover@^6.0.2 version: tsover@6.0.2