From a245c7e44cafaad0329273624d96cf617f7f426b Mon Sep 17 00:00:00 2001 From: Ib Green Date: Tue, 4 Aug 2026 22:11:52 -0400 Subject: [PATCH] fix(experimental): scale GPU sort dispatches --- .../src/gpu-primitives/gpu-sort.ts | 176 +++++++++----- .../test/gpu-primitives/gpu-sort.node.spec.ts | 216 ++++++++++++++++++ .../test/gpu-primitives/gpu-sort.spec.ts | 94 +++++++- 3 files changed, 431 insertions(+), 55 deletions(-) create mode 100644 modules/experimental/test/gpu-primitives/gpu-sort.node.spec.ts diff --git a/modules/experimental/src/gpu-primitives/gpu-sort.ts b/modules/experimental/src/gpu-primitives/gpu-sort.ts index fa3e128f1a..f61a3ec0f2 100644 --- a/modules/experimental/src/gpu-primitives/gpu-sort.ts +++ b/modules/experimental/src/gpu-primitives/gpu-sort.ts @@ -5,7 +5,12 @@ import {type Binding} from '@luma.gl/core'; import {Computation} from '@luma.gl/engine'; import {GPUCommandGraph, type GraphBufferUse, type GraphDataView} from './gpu-command-graph'; -import {GPUScan} from './gpu-scan'; +import { + getBoundedDispatchLayout, + getBoundedInvocationIndexSource, + type GPUBoundedDispatchLayout +} from './gpu-dispatch-utils'; +import {addGPUScanToGraphWithDispatchLimit, GPUScan} from './gpu-scan'; import { createTransientView, getViewBinding, @@ -138,24 +143,46 @@ export class GPUSort { * encode, submit, or read back commands. */ addToGraph(graph: GPUCommandGraph): void { - for (const view of [this.keys, this.values, this.outputKeys, this.outputValues]) { - if (view.buffer.graph !== graph) { - throw new Error(`${this.id} views must belong to the target graph`); - } - } - if (this.keys.length === 0) { - return; - } - if (this.keys.length === 1) { - addCopyPairPass(graph, this); - return; - } - if (this.resolvedAlgorithm === 'bitonic') { - addBitonicSort(graph, this); - } else { - addRadixSort(graph, this); + addGPUSortToGraphWithDispatchLimit( + this, + graph, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + } +} + +/** Adds one stable sort while propagating an explicit bounded dispatch limit. @internal */ +export function addGPUSortToGraphWithDispatchLimit( + sort: GPUSort, + graph: GPUCommandGraph, + maxComputeWorkgroupsPerDimension: number +): void { + for (const view of [sort.keys, sort.values, sort.outputKeys, sort.outputValues]) { + if (view.buffer.graph !== graph) { + throw new Error(`${sort.id} views must belong to the target graph`); } } + + if (sort.keys.length === 0) { + return; + } + if (sort.keys.length === 1) { + addCopyPairPass(graph, sort); + return; + } + + const dispatchLayout = getBoundedDispatchLayout( + 'GPUSort', + sort.keys.length, + RADIX_WORKGROUP_SIZE, + maxComputeWorkgroupsPerDimension + ); + + if (sort.resolvedAlgorithm === 'bitonic') { + addBitonicSort(graph, sort, dispatchLayout, maxComputeWorkgroupsPerDimension); + } else { + addRadixSort(graph, sort, dispatchLayout, maxComputeWorkgroupsPerDimension); + } } /** Enforces out-of-place writes and distinct writable destinations. */ @@ -177,7 +204,8 @@ function addCopyPairPass( sort: GPUSort, inputKeys: GraphDataView<'uint32'> = sort.keys, inputValues: GraphDataView<'uint32'> = sort.values, - identifier = 'copy-pair' + identifier = 'copy-pair', + dispatchLayout: GPUBoundedDispatchLayout = {x: 1, y: 1, z: 1} ): void { const source = /* wgsl */ ` const ELEMENT_COUNT: u32 = ${sort.keys.length}u; @@ -191,9 +219,10 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u; @group(0) @binding(3) var outputValues: array; @compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, RADIX_WORKGROUP_SIZE)} if (index >= ELEMENT_COUNT) { return; } outputKeys[OUTPUT_KEYS_OFFSET + index] = keys[KEYS_OFFSET + index]; outputValues[OUTPUT_VALUES_OFFSET + index] = values[VALUES_OFFSET + index]; @@ -213,13 +242,24 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u; outputKeys: sort.outputKeys, outputValues: sort.outputValues }, - dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE) + dispatchLayout }); } /** Adds padded-index initialization, every bitonic stage, and the final stable gather. */ -function addBitonicSort(graph: GPUCommandGraph, sort: GPUSort): void { +function addBitonicSort( + graph: GPUCommandGraph, + sort: GPUSort, + dispatchLayout: GPUBoundedDispatchLayout, + maxComputeWorkgroupsPerDimension: number +): void { const paddedLength = getNextPowerOfTwo(sort.keys.length); + const paddedDispatchLayout = getBoundedDispatchLayout( + 'GPUSort bitonic', + paddedLength, + BITONIC_WORKGROUP_SIZE, + maxComputeWorkgroupsPerDimension + ); const indicesA = createTransientView( graph, `${sort.id}-bitonic-indices-a`, @@ -232,15 +272,23 @@ function addBitonicSort(graph: GPUCommandGraph, sort: GP 'uint32', paddedLength ); - addBitonicInitializePass(graph, sort, indicesA, paddedLength); + addBitonicInitializePass(graph, sort, indicesA, paddedLength, paddedDispatchLayout); let currentIndices = indicesA; let nextIndices = indicesB; for (const stage of getBitonicStages(paddedLength)) { - addBitonicStagePass(graph, sort, currentIndices, nextIndices, paddedLength, stage); + addBitonicStagePass( + graph, + sort, + currentIndices, + nextIndices, + paddedLength, + stage, + paddedDispatchLayout + ); [currentIndices, nextIndices] = [nextIndices, currentIndices]; } - addBitonicGatherPass(graph, sort, currentIndices); + addBitonicGatherPass(graph, sort, currentIndices, dispatchLayout); } /** Initializes logical indices and invalid padding for a power-of-two bitonic network. */ @@ -248,7 +296,8 @@ function addBitonicInitializePass( graph: GPUCommandGraph, sort: GPUSort, indices: GraphDataView<'uint32'>, - paddedLength: number + paddedLength: number, + dispatchLayout: GPUBoundedDispatchLayout ): void { const source = /* wgsl */ ` const INVALID_INDEX: u32 = ${INVALID_INDEX}u; @@ -258,9 +307,10 @@ const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u; @group(0) @binding(0) var indices: array; @compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, BITONIC_WORKGROUP_SIZE)} if (index < PADDED_LENGTH) { indices[INDICES_OFFSET + index] = select(INVALID_INDEX, index, index < LOGICAL_LENGTH); } @@ -270,7 +320,7 @@ const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u; source, resources: [{buffer: indices, usage: 'storage-write'}], bindings: {indices}, - dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE) + dispatchLayout }); } @@ -281,7 +331,8 @@ function addBitonicStagePass( indicesIn: GraphDataView<'uint32'>, indicesOut: GraphDataView<'uint32'>, paddedLength: number, - stage: BitonicStage + stage: BitonicStage, + dispatchLayout: GPUBoundedDispatchLayout ): void { const descending = sort.direction === 'descending'; const source = /* wgsl */ ` @@ -313,9 +364,10 @@ fn comes_before(leftIndex: u32, rightIndex: u32) -> bool { } @compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, BITONIC_WORKGROUP_SIZE)} if (index >= PADDED_LENGTH) { return; } let partnerIndex = index ^ COMPARE_STRIDE; if (partnerIndex <= index) { return; } @@ -339,7 +391,7 @@ fn comes_before(leftIndex: u32, rightIndex: u32) -> bool { {buffer: indicesOut, usage: 'storage-write'} ], bindings: {keys: sort.keys, indicesIn, indicesOut}, - dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE) + dispatchLayout }); } @@ -347,7 +399,8 @@ fn comes_before(leftIndex: u32, rightIndex: u32) -> bool { function addBitonicGatherPass( graph: GPUCommandGraph, sort: GPUSort, - indices: GraphDataView<'uint32'> + indices: GraphDataView<'uint32'>, + dispatchLayout: GPUBoundedDispatchLayout ): void { const source = /* wgsl */ ` const LOGICAL_LENGTH: u32 = ${sort.keys.length}u; @@ -363,9 +416,10 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u; @group(0) @binding(4) var outputValues: array; @compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, BITONIC_WORKGROUP_SIZE)} if (index >= LOGICAL_LENGTH) { return; } let sourceIndex = indices[INDICES_OFFSET + index]; outputKeys[OUTPUT_KEYS_OFFSET + index] = keys[KEYS_OFFSET + sourceIndex]; @@ -388,12 +442,17 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u; outputKeys: sort.outputKeys, outputValues: sort.outputValues }, - dispatchCount: Math.ceil(sort.keys.length / BITONIC_WORKGROUP_SIZE) + dispatchLayout }); } /** Adds stable least-significant-bit radix partitions and the final output copy if needed. */ -function addRadixSort(graph: GPUCommandGraph, sort: GPUSort): void { +function addRadixSort( + graph: GPUCommandGraph, + sort: GPUSort, + dispatchLayout: GPUBoundedDispatchLayout, + maxComputeWorkgroupsPerDimension: number +): void { const scratchKeys = createTransientView( graph, `${sort.id}-radix-scratch-keys`, @@ -424,12 +483,13 @@ function addRadixSort(graph: GPUCommandGraph, sort: GPUS ); const nextKeys = bit % 2 === 0 ? scratchKeys : sort.outputKeys; const nextValues = bit % 2 === 0 ? scratchValues : sort.outputValues; - addRadixClassifyPass(graph, sort, currentKeys, flags, bit); - new GPUScan({ + addRadixClassifyPass(graph, sort, currentKeys, flags, bit, dispatchLayout); + const scan = new GPUScan({ id: `${sort.id}-radix-bit-${bit}-scan`, input: flags, output: offsets - }).addToGraph(graph); + }); + addGPUScanToGraphWithDispatchLimit(scan, graph, maxComputeWorkgroupsPerDimension); addRadixScatterPass( graph, sort, @@ -439,14 +499,15 @@ function addRadixSort(graph: GPUCommandGraph, sort: GPUS offsets, nextKeys, nextValues, - bit + bit, + dispatchLayout ); currentKeys = nextKeys; currentValues = nextValues; } if (sort.keyBits % 2 !== 0) { - addCopyPairPass(graph, sort, currentKeys, currentValues, 'radix-final-copy'); + addCopyPairPass(graph, sort, currentKeys, currentValues, 'radix-final-copy', dispatchLayout); } } @@ -456,7 +517,8 @@ function addRadixClassifyPass( sort: GPUSort, keys: GraphDataView<'uint32'>, flags: GraphDataView<'uint32'>, - bit: number + bit: number, + dispatchLayout: GPUBoundedDispatchLayout ): void { const firstBit = sort.direction === 'ascending' ? 0 : 1; const source = /* wgsl */ ` @@ -469,9 +531,10 @@ const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u; @group(0) @binding(1) var flags: array; @compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, RADIX_WORKGROUP_SIZE)} if (index >= ELEMENT_COUNT) { return; } let bitValue = (keys[KEYS_OFFSET + index] >> BIT_INDEX) & 1u; flags[FLAGS_OFFSET + index] = select(0u, 1u, bitValue == FIRST_BIT); @@ -484,7 +547,7 @@ const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u; {buffer: flags, usage: 'storage-write'} ], bindings: {keys, flags}, - dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE) + dispatchLayout }); } @@ -498,7 +561,8 @@ function addRadixScatterPass( offsets: GraphDataView<'uint32'>, outputKeys: GraphDataView<'uint32'>, outputValues: GraphDataView<'uint32'>, - bit: number + bit: number, + dispatchLayout: GPUBoundedDispatchLayout ): void { const lastIndex = sort.keys.length - 1; const source = /* wgsl */ ` @@ -518,9 +582,10 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(outputValues)}u; @group(0) @binding(5) var outputValues: array; @compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main( - @builtin(global_invocation_id) globalId: vec3 + @builtin(local_invocation_index) localInvocationIndex: u32, + @builtin(workgroup_id) workgroupId: vec3 ) { - let index = globalId.x; + ${getBoundedInvocationIndexSource(dispatchLayout, RADIX_WORKGROUP_SIZE)} if (index >= ELEMENT_COUNT) { return; } let firstOffset = offsets[OFFSETS_OFFSET + index]; let firstCount = offsets[OFFSETS_OFFSET + LAST_INDEX] + flags[FLAGS_OFFSET + LAST_INDEX]; @@ -541,7 +606,7 @@ const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(outputValues)}u; {buffer: outputValues, usage: 'storage-write'} ], bindings: {keys, values, flags, offsets, outputKeys, outputValues}, - dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE) + dispatchLayout }); } @@ -573,7 +638,7 @@ function addComputationPass( source: string; resources: GraphBufferUse[]; bindings: Record; - dispatchCount: number; + dispatchLayout: GPUBoundedDispatchLayout; } ): void { graph.addComputePass({ @@ -599,7 +664,12 @@ function addComputationPass( bindings[name] = getViewBinding(view, getBuffer); } computation.setBindings(bindings); - computation.dispatch(computePass, props.dispatchCount); + computation.dispatch( + computePass, + props.dispatchLayout.x, + props.dispatchLayout.y, + props.dispatchLayout.z + ); }, destroy: () => computation.destroy() }; diff --git a/modules/experimental/test/gpu-primitives/gpu-sort.node.spec.ts b/modules/experimental/test/gpu-primitives/gpu-sort.node.spec.ts new file mode 100644 index 0000000000..24a36d6867 --- /dev/null +++ b/modules/experimental/test/gpu-primitives/gpu-sort.node.spec.ts @@ -0,0 +1,216 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import { + GPUBatchSort, + GPUCommandGraph, + GraphVectorView, + GPUSort, + type GPUSortAlgorithm +} from '@luma.gl/experimental'; +import {NullDevice} from '@luma.gl/test-utils'; +import {describe, expect, test, vi} from 'vitest'; +import { + getBoundedDispatchLayout, + getBoundedInvocationIndexSource +} from '../../src/gpu-primitives/gpu-dispatch-utils'; +import {addGPUSortToGraphWithDispatchLimit} from '../../src/gpu-primitives/gpu-sort'; + +const WORKGROUP_SIZE = 256; + +describe('bounded GPU sort dispatch', () => { + test('preserves one-dimensional boundaries and safely uses all three dimensions', () => { + const maximum = 65_535; + const oneDimensionalCapacity = maximum * WORKGROUP_SIZE; + + expect( + getBoundedDispatchLayout('GPUSort', oneDimensionalCapacity, WORKGROUP_SIZE, maximum) + ).toEqual({x: maximum, y: 1, z: 1}); + expect( + getBoundedDispatchLayout('GPUSort', oneDimensionalCapacity + 1, WORKGROUP_SIZE, maximum) + ).toEqual({x: maximum, y: 2, z: 1}); + expect(getBoundedDispatchLayout('GPUSort', 4 * WORKGROUP_SIZE + 1, WORKGROUP_SIZE, 2)).toEqual({ + x: 2, + y: 2, + z: 2 + }); + expect(() => + getBoundedDispatchLayout('GPUSort', 8 * WORKGROUP_SIZE + 1, WORKGROUP_SIZE, 2) + ).toThrow(/exceeding the 3D dispatch limit/i); + expect(getBoundedDispatchLayout('GPUSort', 0x80000000, WORKGROUP_SIZE, maximum)).toEqual({ + x: maximum, + y: 129, + z: 1 + }); + + const source = getBoundedInvocationIndexSource({x: 2, y: 2, z: 2}, WORKGROUP_SIZE); + expect(source).toContain('workgroupId.z * 2u + workgroupId.y'); + expect(source).toContain('* 2u + workgroupId.x'); + expect(source.indexOf('workgroupIndex >= 16777216u')).toBeLessThan( + source.indexOf('workgroupIndex * 256u + localInvocationIndex') + ); + }); + + test('propagates one synthetic limit through every bitonic and radix scan stage', () => { + for (const algorithm of ['bitonic', 'radix'] as const) { + const fixture = createSortGraphFixture(4 * WORKGROUP_SIZE + 1, algorithm); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + const createBuffer = vi.spyOn(fixture.device, 'createBuffer'); + + try { + addGPUSortToGraphWithDispatchLimit(fixture.sort, fixture.graph, 2); + + const identifiers = addComputePass.mock.calls.map(([pass]) => pass.id); + if (algorithm === 'bitonic') { + expect(identifiers).toContain('node-sort-bitonic-initialize'); + expect(identifiers).toContain('node-sort-bitonic-2048-1'); + expect(identifiers).toContain('node-sort-bitonic-gather'); + } else { + expect(identifiers).toContain('node-sort-radix-bit-0-classify'); + expect(identifiers).toContain('node-sort-radix-bit-0-scan-level-0-scan'); + expect(identifiers).toContain('node-sort-radix-bit-0-scan-level-0-add-offsets'); + expect(identifiers).toContain('node-sort-radix-bit-0-scatter'); + expect(identifiers).toContain('node-sort-radix-bit-31-classify'); + expect(identifiers).toContain('node-sort-radix-bit-31-scatter'); + } + expect(createBuffer).not.toHaveBeenCalled(); + } finally { + addComputePass.mockRestore(); + createBuffer.mockRestore(); + fixture.device.destroy(); + } + } + }); + + test('rejects padded bitonic overflow before mutating the command graph', () => { + const fixture = createSortGraphFixture(4_097, 'bitonic'); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + const createTransientBuffer = vi.spyOn(fixture.graph, 'createTransientBuffer'); + + try { + expect(() => addGPUSortToGraphWithDispatchLimit(fixture.sort, fixture.graph, 3)).toThrow( + /GPUSort bitonic.*exceeding the 3D dispatch limit/i + ); + expect(addComputePass).not.toHaveBeenCalled(); + expect(createTransientBuffer).not.toHaveBeenCalled(); + } finally { + addComputePass.mockRestore(); + createTransientBuffer.mockRestore(); + fixture.device.destroy(); + } + }); + + test('rejects radix ranges beyond the complete bounded three-dimensional capacity', () => { + const fixture = createSortGraphFixture(8 * WORKGROUP_SIZE + 1, 'radix'); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + + try { + expect(() => addGPUSortToGraphWithDispatchLimit(fixture.sort, fixture.graph, 2)).toThrow( + /GPUSort.*exceeding the 3D dispatch limit/i + ); + expect(addComputePass).not.toHaveBeenCalled(); + } finally { + addComputePass.mockRestore(); + fixture.device.destroy(); + } + }); + + test('preserves empty and single-row fast paths without consulting dispatch limits', () => { + for (const length of [0, 1]) { + const fixture = createSortGraphFixture(length, 'auto'); + const addComputePass = vi.spyOn(fixture.graph, 'addComputePass'); + + try { + addGPUSortToGraphWithDispatchLimit(fixture.sort, fixture.graph, 0); + expect(addComputePass.mock.calls.map(([pass]) => pass.id)).toEqual( + length === 0 ? [] : ['node-sort-copy-pair'] + ); + } finally { + addComputePass.mockRestore(); + fixture.device.destroy(); + } + } + }); + + test('propagates the real device limit through independent batch sort chunks', () => { + const device = new NullDevice({id: 'bounded-batch-sort-node-device'}); + Object.defineProperty(device, 'type', {value: 'webgpu'}); + device.limits.maxComputeWorkgroupsPerDimension = 2; + const graph = new GPUCommandGraph(device, {id: 'bounded-batch-sort-node-graph'}); + const lengths = [1_025, 0, 513]; + const sort = new GPUBatchSort({ + id: 'bounded-batch-sort', + keys: createSortVector(graph, 'keys', lengths), + values: createSortVector(graph, 'values', lengths), + outputKeys: createSortVector(graph, 'output-keys', lengths), + outputValues: createSortVector(graph, 'output-values', lengths) + }); + const addComputePass = vi.spyOn(graph, 'addComputePass'); + + try { + sort.addToGraph(graph); + + const identifiers = addComputePass.mock.calls.map(([pass]) => pass.id); + expect(identifiers).toContain('bounded-batch-sort-chunk-0-bitonic-initialize'); + expect(identifiers).toContain('bounded-batch-sort-chunk-0-bitonic-2048-1'); + expect(identifiers).toContain('bounded-batch-sort-chunk-0-bitonic-gather'); + expect(identifiers).toContain('bounded-batch-sort-chunk-2-bitonic-initialize'); + expect(identifiers).toContain('bounded-batch-sort-chunk-2-bitonic-1024-1'); + expect(identifiers).toContain('bounded-batch-sort-chunk-2-bitonic-gather'); + expect(identifiers.some(identifier => identifier?.includes('chunk-1'))).toBe(false); + } finally { + addComputePass.mockRestore(); + device.destroy(); + } + }); +}); + +function createSortGraphFixture( + length: number, + algorithm: GPUSortAlgorithm +): {device: NullDevice; graph: GPUCommandGraph; sort: GPUSort} { + const device = new NullDevice({id: 'bounded-sort-node-device'}); + Object.defineProperty(device, 'type', {value: 'webgpu'}); + const graph = new GPUCommandGraph(device, {id: 'bounded-sort-node-graph'}); + const sort = new GPUSort({ + id: 'node-sort', + keys: createSortView(graph, 'keys', length), + values: createSortView(graph, 'values', length), + outputKeys: createSortView(graph, 'output-keys', length), + outputValues: createSortView(graph, 'output-values', length), + algorithm + }); + return {device, graph, sort}; +} + +function createSortView(graph: GPUCommandGraph, id: string, length: number) { + const buffer = graph.createTransientBuffer({ + id, + byteLength: Math.max(length, 1) * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.STORAGE + }); + return graph.createDataView(buffer, {format: 'uint32', length}); +} + +function createSortVector( + graph: GPUCommandGraph, + id: string, + lengths: readonly number[] +): GraphVectorView<'uint32'> { + const length = lengths.reduce((total, chunkLength) => total + chunkLength, 0); + return new GraphVectorView({ + id, + name: id, + format: 'uint32', + length, + valueLength: length, + stride: 1, + byteStride: Uint32Array.BYTES_PER_ELEMENT, + rowByteLength: Uint32Array.BYTES_PER_ELEMENT, + data: lengths.map((chunkLength, chunkIndex) => + createSortView(graph, `${id}-chunk-${chunkIndex}`, chunkLength) + ) + }); +} diff --git a/modules/experimental/test/gpu-primitives/gpu-sort.spec.ts b/modules/experimental/test/gpu-primitives/gpu-sort.spec.ts index 809b988289..3575b21176 100644 --- a/modules/experimental/test/gpu-primitives/gpu-sort.spec.ts +++ b/modules/experimental/test/gpu-primitives/gpu-sort.spec.ts @@ -4,6 +4,7 @@ import test from 'test/utils/vitest-tape'; import {Buffer, type Device} from '@luma.gl/core'; +import {Computation} from '@luma.gl/engine'; import { GPUBatchSort, GPUCommandGraph, @@ -14,6 +15,8 @@ import { } from '@luma.gl/experimental'; import {GPUData, GPUVector} from '@luma.gl/tables'; import {getWebGPUTestDevice} from '@luma.gl/test-utils'; +import {vi} from 'vitest'; +import {addGPUSortToGraphWithDispatchLimit} from '../../src/gpu-primitives/gpu-sort'; test('GPUSort bitonic stably sorts paired uint32 values in both directions', async t => { const device = await getWebGPUTestDevice(); @@ -88,6 +91,88 @@ test('GPUSort radix processes only the requested significant key bits', async t 'odd key widths copy their final scratch result into the caller-owned outputs' ); } + + t.end(); +}); + +test('GPUSort bounds bitonic and radix stages across all three dispatch dimensions', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const rowCount = 1_025; + const keys = Uint32Array.from({length: rowCount}, (_, index) => + index % 23 === 0 ? 7 : (Math.imul(index, 1_664_525) + 1_013_904_223) >>> 0 + ); + const values = Uint32Array.from({length: rowCount}, (_, index) => 2_000 + index); + + for (const [algorithm, direction] of [ + ['bitonic', 'ascending'], + ['radix', 'descending'] + ] as const) { + const dispatchSpy = vi.spyOn(Computation.prototype, 'dispatch'); + + try { + const result = await runSort(device, keys, values, algorithm, direction, undefined, 2); + const expected = getStableSortedPairs(keys, values, direction); + t.deepEqual(result.keys, expected.keys, `${algorithm} sorts every multidimensional key`); + t.deepEqual( + result.values, + expected.values, + `${algorithm} preserves duplicate-key payload order across workgroups` + ); + + const dispatches = dispatchSpy.mock.instances.map((computation, index) => ({ + id: (computation as Computation).id, + dimensions: dispatchSpy.mock.calls[index].slice(1) + })); + const expectedPasses = + algorithm === 'bitonic' + ? ['sort-bitonic-initialize', 'sort-bitonic-2048-1', 'sort-bitonic-gather'] + : [ + 'sort-radix-bit-0-classify', + 'sort-radix-bit-0-scan-level-0-scan', + 'sort-radix-bit-0-scan-level-0-add-offsets', + 'sort-radix-bit-0-scatter', + 'sort-radix-bit-31-classify', + 'sort-radix-bit-31-scatter' + ]; + + for (const identifier of expectedPasses) { + t.deepEqual( + dispatches.find(dispatch => dispatch.id === identifier)?.dimensions, + [2, 2, 2], + `${identifier} respects the synthetic per-dimension dispatch limit` + ); + } + } finally { + dispatchSpy.mockRestore(); + } + } + + const limitedKeys = Uint32Array.from(keys, key => key & 0x7fff); + const dispatchSpy = vi.spyOn(Computation.prototype, 'dispatch'); + try { + const result = await runSort(device, limitedKeys, values, 'radix', 'ascending', 15, 2); + const expected = getStableSortedPairs(limitedKeys, values, 'ascending'); + t.deepEqual(result.keys, expected.keys, 'odd-width multidimensional radix keys match'); + t.deepEqual(result.values, expected.values, 'odd-width multidimensional radix remains stable'); + + const finalCopyDispatchIndex = dispatchSpy.mock.instances.findIndex( + computation => (computation as Computation).id === 'sort-radix-final-copy' + ); + t.deepEqual( + dispatchSpy.mock.calls[finalCopyDispatchIndex]?.slice(1), + [2, 2, 2], + 'odd-width radix final copy respects the synthetic per-dimension dispatch limit' + ); + } finally { + dispatchSpy.mockRestore(); + } + t.end(); }); @@ -375,7 +460,8 @@ async function runSort( values: Uint32Array, algorithm: GPUSortAlgorithm, direction: GPUSortDirection, - keyBits?: number + keyBits?: number, + maxComputeWorkgroupsPerDimension?: number ): Promise { const byteLength = Math.max(keys.length, 1) * Uint32Array.BYTES_PER_ELEMENT; const keysBuffer = device.createBuffer({ @@ -413,7 +499,11 @@ async function runSort( direction, keyBits }); - sort.addToGraph(graph); + if (maxComputeWorkgroupsPerDimension === undefined) { + sort.addToGraph(graph); + } else { + addGPUSortToGraphWithDispatchLimit(sort, graph, maxComputeWorkgroupsPerDimension); + } const compiled = graph.compile(); const commandEncoder = device.createCommandEncoder({id: 'sort-test-encoder'}); compiled.encode(commandEncoder, {parameters: undefined});