28617818631

Committed 02 Jul 2026 07:57PM UTC coverage: 71.941% (+0.4%) from 71.566%

Build # 28617818631

Build Type

push

github

Committed by

web-flow

Commit Message

[codex] Document GPU picking and command-graph roadmap (#2711)

Coverage Stats

10762 of 16860 branches covered (63.83%)

Branch coverage included in aggregate %.

21074 of 27393 relevant lines covered (76.93%)

5474.44 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.68

/modules/experimental/src/gpu-primitives/gpu-sort.ts

// luma.gl
// SPDX-License-Identifier: MIT
// Copyright (c) vis.gl contributors

import {type Binding} from '@luma.gl/core';
import {Computation} from '@luma.gl/engine';
import {GPUCommandGraph, type GraphBufferUse, type GraphBufferView} from './gpu-command-graph';
import {GPUScan} from './gpu-scan';
import {
  createTransientView,
  getViewBinding,
  getViewElementOffset,
  validatePackedUint32View
} from './graph-buffer-view-utils';

const BITONIC_WORKGROUP_SIZE = 256;
const RADIX_WORKGROUP_SIZE = 256;
const INVALID_INDEX = 0xffffffff;
const MAXIMUM_LOGICAL_LENGTH = 0x80000000;
const AUTO_BITONIC_MAXIMUM_LENGTH = 65_536;

/** Sort implementation requested by {@link GPUSort}. */
export type GPUSortAlgorithm = 'auto' | 'bitonic' | 'radix';

/** Final key ordering requested by {@link GPUSort}. */
export type GPUSortDirection = 'ascending' | 'descending';

/** Properties for one graph-native stable uint32 key/value sort. */
export type GPUSortProps = {
  id?: string;
  keys: GraphBufferView<'uint32'>;
  values: GraphBufferView<'uint32'>;
  outputKeys: GraphBufferView<'uint32'>;
  outputValues: GraphBufferView<'uint32'>;
  algorithm?: GPUSortAlgorithm;
  direction?: GPUSortDirection;
};

type BitonicStage = {
  blockWidth: number;
  compareStride: number;
};

/**
 * Stable graph-native sort for paired packed uint32 keys and values.
 *
 * @remarks
 * The operation is out-of-place. Inputs and outputs are caller-owned graph views, while all
 * implementation scratch is graph-owned. `addToGraph()` only records work; the caller retains
 * control of graph compilation, command encoding, submission, and optional readback.
 */
export class GPUSort {
  readonly id: string;
  readonly keys: GraphBufferView<'uint32'>;
  readonly values: GraphBufferView<'uint32'>;
  readonly outputKeys: GraphBufferView<'uint32'>;
  readonly outputValues: GraphBufferView<'uint32'>;
  readonly algorithm: GPUSortAlgorithm;
  readonly direction: GPUSortDirection;
  readonly resolvedAlgorithm: Exclude<GPUSortAlgorithm, 'auto'>;

  constructor(props: GPUSortProps) {
    this.id = props.id ?? 'gpu-sort';
    this.keys = props.keys;
    this.values = props.values;
    this.outputKeys = props.outputKeys;
    this.outputValues = props.outputValues;
    this.algorithm = props.algorithm ?? 'auto';
    this.direction = props.direction ?? 'ascending';

    for (const [name, view] of [
      ['keys', this.keys],
      ['values', this.values],
      ['outputKeys', this.outputKeys],
      ['outputValues', this.outputValues]
    ] as const) {
      validatePackedUint32View(view, `${this.id} ${name}`);
    }
    if (!['auto', 'bitonic', 'radix'].includes(this.algorithm)) {
      throw new Error(`${this.id} algorithm must be auto, bitonic, or radix`);
    }
    if (!['ascending', 'descending'].includes(this.direction)) {
      throw new Error(`${this.id} direction must be ascending or descending`);
    }
    if (
      this.values.length !== this.keys.length ||
      this.outputKeys.length !== this.keys.length ||
      this.outputValues.length !== this.keys.length
    ) {
      throw new Error(`${this.id} key, value, and output lengths must match`);
    }
    if (this.keys.length > MAXIMUM_LOGICAL_LENGTH) {
      throw new Error(`${this.id} supports at most ${MAXIMUM_LOGICAL_LENGTH} rows`);
    }
    validateSeparateWritableBuffers(this);

    this.resolvedAlgorithm =
      this.algorithm === 'auto'
        ? this.keys.length <= AUTO_BITONIC_MAXIMUM_LENGTH
          ? 'bitonic'
          : 'radix'
        : this.algorithm;
  }

  /** Adds the selected sort implementation and its scratch buffers to a command graph. */
  addToGraph<Parameters>(graph: GPUCommandGraph<Parameters>): 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);
    }
  }
}

function validateSeparateWritableBuffers(sort: GPUSort): void {
  if (
    sort.outputKeys.buffer === sort.outputValues.buffer ||
    sort.outputKeys.buffer === sort.keys.buffer ||
    sort.outputKeys.buffer === sort.values.buffer ||
    sort.outputValues.buffer === sort.keys.buffer ||
    sort.outputValues.buffer === sort.values.buffer
  ) {
    throw new Error(`${sort.id} outputs must use separate buffers from inputs and each other`);
  }
}

function addCopyPairPass<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
  const source = /* wgsl */ `
const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
const VALUES_OFFSET: u32 = ${getViewElementOffset(sort.values)}u;
const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(sort.outputKeys)}u;
const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u;
@group(0) @binding(0) var<storage, read> keys: array<u32>;
@group(0) @binding(1) var<storage, read> values: array<u32>;
@group(0) @binding(2) var<storage, read_write> outputKeys: array<u32>;
@group(0) @binding(3) var<storage, read_write> outputValues: array<u32>;

@compute @workgroup_size(1) fn main() {
  outputKeys[OUTPUT_KEYS_OFFSET] = keys[KEYS_OFFSET];
  outputValues[OUTPUT_VALUES_OFFSET] = values[VALUES_OFFSET];
}`;
  addComputationPass(graph, {
    id: `${sort.id}-copy-pair`,
    source,
    resources: [
      {buffer: sort.keys, usage: 'storage-read'},
      {buffer: sort.values, usage: 'storage-read'},
      {buffer: sort.outputKeys, usage: 'storage-write'},
      {buffer: sort.outputValues, usage: 'storage-write'}
    ],
    bindings: {
      keys: sort.keys,
      values: sort.values,
      outputKeys: sort.outputKeys,
      outputValues: sort.outputValues
    },
    dispatchCount: 1
  });
}

function addBitonicSort<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
  const paddedLength = getNextPowerOfTwo(sort.keys.length);
  const indicesA = createTransientView(
    graph,
    `${sort.id}-bitonic-indices-a`,
    'uint32',
    paddedLength
  );
  const indicesB = createTransientView(
    graph,
    `${sort.id}-bitonic-indices-b`,
    'uint32',
    paddedLength
  );
  addBitonicInitializePass(graph, sort, indicesA, paddedLength);

  let currentIndices = indicesA;
  let nextIndices = indicesB;
  for (const stage of getBitonicStages(paddedLength)) {
    addBitonicStagePass(graph, sort, currentIndices, nextIndices, paddedLength, stage);
    [currentIndices, nextIndices] = [nextIndices, currentIndices];
  }
  addBitonicGatherPass(graph, sort, currentIndices);
}

function addBitonicInitializePass<Parameters>(
  graph: GPUCommandGraph<Parameters>,
  sort: GPUSort,
  indices: GraphBufferView<'uint32'>,
  paddedLength: number
): void {
  const source = /* wgsl */ `
const INVALID_INDEX: u32 = ${INVALID_INDEX}u;
const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
const PADDED_LENGTH: u32 = ${paddedLength}u;
const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u;
@group(0) @binding(0) var<storage, read_write> indices: array<u32>;

@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
  @builtin(global_invocation_id) globalId: vec3<u32>
) {
  let index = globalId.x;
  if (index < PADDED_LENGTH) {
    indices[INDICES_OFFSET + index] = select(INVALID_INDEX, index, index < LOGICAL_LENGTH);
  }
}`;
  addComputationPass(graph, {
    id: `${sort.id}-bitonic-initialize`,
    source,
    resources: [{buffer: indices, usage: 'storage-write'}],
    bindings: {indices},
    dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE)
  });
}

function addBitonicStagePass<Parameters>(
  graph: GPUCommandGraph<Parameters>,
  sort: GPUSort,
  indicesIn: GraphBufferView<'uint32'>,
  indicesOut: GraphBufferView<'uint32'>,
  paddedLength: number,
  stage: BitonicStage
): void {
  const descending = sort.direction === 'descending';
  const source = /* wgsl */ `
const INVALID_INDEX: u32 = ${INVALID_INDEX}u;
const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
const PADDED_LENGTH: u32 = ${paddedLength}u;
const BLOCK_WIDTH: u32 = ${stage.blockWidth}u;
const COMPARE_STRIDE: u32 = ${stage.compareStride}u;
const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
const INDICES_IN_OFFSET: u32 = ${getViewElementOffset(indicesIn)}u;
const INDICES_OUT_OFFSET: u32 = ${getViewElementOffset(indicesOut)}u;
@group(0) @binding(0) var<storage, read> keys: array<u32>;
@group(0) @binding(1) var<storage, read> indicesIn: array<u32>;
@group(0) @binding(2) var<storage, read_write> indicesOut: array<u32>;

fn is_valid(index: u32) -> bool {
  return index != INVALID_INDEX && index < LOGICAL_LENGTH;
}

fn comes_before(leftIndex: u32, rightIndex: u32) -> bool {
  let leftValid = is_valid(leftIndex);
  let rightValid = is_valid(rightIndex);
  if (leftValid != rightValid) { return leftValid; }
  if (!leftValid) { return false; }
  let leftKey = keys[KEYS_OFFSET + leftIndex];
  let rightKey = keys[KEYS_OFFSET + rightIndex];
  if (leftKey == rightKey) { return leftIndex < rightIndex; }
  return ${descending ? 'leftKey > rightKey' : 'leftKey < rightKey'};
}

@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
  @builtin(global_invocation_id) globalId: vec3<u32>
) {
  let index = globalId.x;
  if (index >= PADDED_LENGTH) { return; }
  let partnerIndex = index ^ COMPARE_STRIDE;
  if (partnerIndex <= index) { return; }
  let leftIndex = indicesIn[INDICES_IN_OFFSET + index];
  let rightIndex = indicesIn[INDICES_IN_OFFSET + partnerIndex];
  let ascending = (index & BLOCK_WIDTH) == 0u;
  let shouldSwap = select(
    comes_before(leftIndex, rightIndex),
    comes_before(rightIndex, leftIndex),
    ascending
  );
  indicesOut[INDICES_OUT_OFFSET + index] = select(leftIndex, rightIndex, shouldSwap);
  indicesOut[INDICES_OUT_OFFSET + partnerIndex] = select(rightIndex, leftIndex, shouldSwap);
}`;
  addComputationPass(graph, {
    id: `${sort.id}-bitonic-${stage.blockWidth}-${stage.compareStride}`,
    source,
    resources: [
      {buffer: sort.keys, usage: 'storage-read'},
      {buffer: indicesIn, usage: 'storage-read'},
      {buffer: indicesOut, usage: 'storage-write'}
    ],
    bindings: {keys: sort.keys, indicesIn, indicesOut},
    dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE)
  });
}

function addBitonicGatherPass<Parameters>(
  graph: GPUCommandGraph<Parameters>,
  sort: GPUSort,
  indices: GraphBufferView<'uint32'>
): void {
  const source = /* wgsl */ `
const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
const VALUES_OFFSET: u32 = ${getViewElementOffset(sort.values)}u;
const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u;
const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(sort.outputKeys)}u;
const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u;
@group(0) @binding(0) var<storage, read> keys: array<u32>;
@group(0) @binding(1) var<storage, read> values: array<u32>;
@group(0) @binding(2) var<storage, read> indices: array<u32>;
@group(0) @binding(3) var<storage, read_write> outputKeys: array<u32>;
@group(0) @binding(4) var<storage, read_write> outputValues: array<u32>;

@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
  @builtin(global_invocation_id) globalId: vec3<u32>
) {
  let index = globalId.x;
  if (index >= LOGICAL_LENGTH) { return; }
  let sourceIndex = indices[INDICES_OFFSET + index];
  outputKeys[OUTPUT_KEYS_OFFSET + index] = keys[KEYS_OFFSET + sourceIndex];
  outputValues[OUTPUT_VALUES_OFFSET + index] = values[VALUES_OFFSET + sourceIndex];
}`;
  addComputationPass(graph, {
    id: `${sort.id}-bitonic-gather`,
    source,
    resources: [
      {buffer: sort.keys, usage: 'storage-read'},
      {buffer: sort.values, usage: 'storage-read'},
      {buffer: indices, usage: 'storage-read'},
      {buffer: sort.outputKeys, usage: 'storage-write'},
      {buffer: sort.outputValues, usage: 'storage-write'}
    ],
    bindings: {
      keys: sort.keys,
      values: sort.values,
      indices,
      outputKeys: sort.outputKeys,
      outputValues: sort.outputValues
    },
    dispatchCount: Math.ceil(sort.keys.length / BITONIC_WORKGROUP_SIZE)
  });
}

function addRadixSort<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
  const scratchKeys = createTransientView(
    graph,
    `${sort.id}-radix-scratch-keys`,
    'uint32',
    sort.keys.length
  );
  const scratchValues = createTransientView(
    graph,
    `${sort.id}-radix-scratch-values`,
    'uint32',
    sort.keys.length
  );
  let currentKeys = sort.keys;
  let currentValues = sort.values;

  for (let bit = 0; bit < 32; bit++) {
    const flags = createTransientView(
      graph,
      `${sort.id}-radix-bit-${bit}-flags`,
      'uint32',
      sort.keys.length
    );
    const offsets = createTransientView(
      graph,
      `${sort.id}-radix-bit-${bit}-offsets`,
      'uint32',
      sort.keys.length
    );
    const nextKeys = bit % 2 === 0 ? scratchKeys : sort.outputKeys;
    const nextValues = bit % 2 === 0 ? scratchValues : sort.outputValues;
    addRadixClassifyPass(graph, sort, currentKeys, flags, bit);
    new GPUScan({
      id: `${sort.id}-radix-bit-${bit}-scan`,
      input: flags,
      output: offsets
    }).addToGraph(graph);
    addRadixScatterPass(
      graph,
      sort,
      currentKeys,
      currentValues,
      flags,
      offsets,
      nextKeys,
      nextValues,
      bit
    );
    currentKeys = nextKeys;
    currentValues = nextValues;
  }
}

function addRadixClassifyPass<Parameters>(
  graph: GPUCommandGraph<Parameters>,
  sort: GPUSort,
  keys: GraphBufferView<'uint32'>,
  flags: GraphBufferView<'uint32'>,
  bit: number
): void {
  const firstBit = sort.direction === 'ascending' ? 0 : 1;
  const source = /* wgsl */ `
const ELEMENT_COUNT: u32 = ${sort.keys.length}u;
const BIT_INDEX: u32 = ${bit}u;
const FIRST_BIT: u32 = ${firstBit}u;
const KEYS_OFFSET: u32 = ${getViewElementOffset(keys)}u;
const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u;
@group(0) @binding(0) var<storage, read> keys: array<u32>;
@group(0) @binding(1) var<storage, read_write> flags: array<u32>;

@compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main(
  @builtin(global_invocation_id) globalId: vec3<u32>
) {
  let index = globalId.x;
  if (index >= ELEMENT_COUNT) { return; }
  let bitValue = (keys[KEYS_OFFSET + index] >> BIT_INDEX) & 1u;
  flags[FLAGS_OFFSET + index] = select(0u, 1u, bitValue == FIRST_BIT);
}`;
  addComputationPass(graph, {
    id: `${sort.id}-radix-bit-${bit}-classify`,
    source,
    resources: [
      {buffer: keys, usage: 'storage-read'},
      {buffer: flags, usage: 'storage-write'}
    ],
    bindings: {keys, flags},
    dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE)
  });
}

function addRadixScatterPass<Parameters>(
  graph: GPUCommandGraph<Parameters>,
  sort: GPUSort,
  keys: GraphBufferView<'uint32'>,
  values: GraphBufferView<'uint32'>,
  flags: GraphBufferView<'uint32'>,
  offsets: GraphBufferView<'uint32'>,
  outputKeys: GraphBufferView<'uint32'>,
  outputValues: GraphBufferView<'uint32'>,
  bit: number
): void {
  const lastIndex = sort.keys.length - 1;
  const source = /* wgsl */ `
const ELEMENT_COUNT: u32 = ${sort.keys.length}u;
const LAST_INDEX: u32 = ${lastIndex}u;
const KEYS_OFFSET: u32 = ${getViewElementOffset(keys)}u;
const VALUES_OFFSET: u32 = ${getViewElementOffset(values)}u;
const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u;
const OFFSETS_OFFSET: u32 = ${getViewElementOffset(offsets)}u;
const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(outputKeys)}u;
const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(outputValues)}u;
@group(0) @binding(0) var<storage, read> keys: array<u32>;
@group(0) @binding(1) var<storage, read> values: array<u32>;
@group(0) @binding(2) var<storage, read> flags: array<u32>;
@group(0) @binding(3) var<storage, read> offsets: array<u32>;
@group(0) @binding(4) var<storage, read_write> outputKeys: array<u32>;
@group(0) @binding(5) var<storage, read_write> outputValues: array<u32>;

@compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main(
  @builtin(global_invocation_id) globalId: vec3<u32>
) {
  let index = globalId.x;
  if (index >= ELEMENT_COUNT) { return; }
  let firstOffset = offsets[OFFSETS_OFFSET + index];
  let firstCount = offsets[OFFSETS_OFFSET + LAST_INDEX] + flags[FLAGS_OFFSET + LAST_INDEX];
  let isFirst = flags[FLAGS_OFFSET + index] != 0u;
  let outputIndex = select(firstCount + index - firstOffset, firstOffset, isFirst);
  outputKeys[OUTPUT_KEYS_OFFSET + outputIndex] = keys[KEYS_OFFSET + index];
  outputValues[OUTPUT_VALUES_OFFSET + outputIndex] = values[VALUES_OFFSET + index];
}`;
  addComputationPass(graph, {
    id: `${sort.id}-radix-bit-${bit}-scatter`,
    source,
    resources: [
      {buffer: keys, usage: 'storage-read'},
      {buffer: values, usage: 'storage-read'},
      {buffer: flags, usage: 'storage-read'},
      {buffer: offsets, usage: 'storage-read'},
      {buffer: outputKeys, usage: 'storage-write'},
      {buffer: outputValues, usage: 'storage-write'}
    ],
    bindings: {keys, values, flags, offsets, outputKeys, outputValues},
    dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE)
  });
}

function getNextPowerOfTwo(length: number): number {
  let paddedLength = 1;
  while (paddedLength < length) {
    paddedLength *= 2;
  }
  return paddedLength;
}

function getBitonicStages(paddedLength: number): BitonicStage[] {
  const stages: BitonicStage[] = [];
  for (let blockWidth = 2; blockWidth <= paddedLength; blockWidth *= 2) {
    for (let compareStride = blockWidth / 2; compareStride >= 1; compareStride /= 2) {
      stages.push({blockWidth, compareStride});
    }
  }
  return stages;
}

function addComputationPass<GraphParameters>(
  graph: GPUCommandGraph<GraphParameters>,
  props: {
    id: string;
    source: string;
    resources: GraphBufferUse[];
    bindings: Record<string, GraphBufferView>;
    dispatchCount: number;
  }
): void {
  graph.addComputePass({
    id: props.id,
    resources: props.resources,
    compile: ({device}) => {
      const computation = new Computation(device, {
        id: props.id,
        source: props.source,
        shaderLayout: {
          bindings: Object.keys(props.bindings).map((name, location) => ({
            name,
            type: 'storage' as const,
            group: 0,
            location
          }))
        }
      });
      return {
        encode: ({computePass, getBuffer}) => {
          const bindings: Record<string, Binding> = {};
          for (const [name, view] of Object.entries(props.bindings)) {
            bindings[name] = getViewBinding(view, getBuffer);
          }
          computation.setBindings(bindings);
          computation.dispatch(computePass, props.dispatchCount);
        },
        destroy: () => computation.destroy()
      };
    }
  });
}

1	// luma.gl
2	// SPDX-License-Identifier: MIT
3	// Copyright (c) vis.gl contributors
4
5	import {type Binding} from '@luma.gl/core';
6	import {Computation} from '@luma.gl/engine';
7	import {GPUCommandGraph, type GraphBufferUse, type GraphBufferView} from './gpu-command-graph';
8	import {GPUScan} from './gpu-scan';
9	import {
10	createTransientView,
11	getViewBinding,
12	getViewElementOffset,
13	validatePackedUint32View
14	} from './graph-buffer-view-utils';
15
16	const BITONIC_WORKGROUP_SIZE = 256;	11✔
17	const RADIX_WORKGROUP_SIZE = 256;	11✔
18	const INVALID_INDEX = 0xffffffff;	11✔
19	const MAXIMUM_LOGICAL_LENGTH = 0x80000000;	11✔
20	const AUTO_BITONIC_MAXIMUM_LENGTH = 65_536;	11✔
21
22	/** Sort implementation requested by {@link GPUSort}. */
23	export type GPUSortAlgorithm = 'auto' \| 'bitonic' \| 'radix';
24
25	/** Final key ordering requested by {@link GPUSort}. */
26	export type GPUSortDirection = 'ascending' \| 'descending';
27
28	/** Properties for one graph-native stable uint32 key/value sort. */
29	export type GPUSortProps = {
30	id?: string;
31	keys: GraphBufferView<'uint32'>;
32	values: GraphBufferView<'uint32'>;
33	outputKeys: GraphBufferView<'uint32'>;
34	outputValues: GraphBufferView<'uint32'>;
35	algorithm?: GPUSortAlgorithm;
36	direction?: GPUSortDirection;
37	};
38
39	type BitonicStage = {
40	blockWidth: number;
41	compareStride: number;
42	};
43
44	/**
45	* Stable graph-native sort for paired packed uint32 keys and values.
46	*
47	* @remarks
48	* The operation is out-of-place. Inputs and outputs are caller-owned graph views, while all
49	* implementation scratch is graph-owned. `addToGraph()` only records work; the caller retains
50	* control of graph compilation, command encoding, submission, and optional readback.
51	*/
52	export class GPUSort {
53	readonly id: string;
54	readonly keys: GraphBufferView<'uint32'>;
55	readonly values: GraphBufferView<'uint32'>;
56	readonly outputKeys: GraphBufferView<'uint32'>;
57	readonly outputValues: GraphBufferView<'uint32'>;
58	readonly algorithm: GPUSortAlgorithm;
59	readonly direction: GPUSortDirection;
60	readonly resolvedAlgorithm: Exclude<GPUSortAlgorithm, 'auto'>;
61
62	constructor(props: GPUSortProps) {
63	this.id = props.id ?? 'gpu-sort';	13✔
64	this.keys = props.keys;	13✔
65	this.values = props.values;	13✔
66	this.outputKeys = props.outputKeys;	13✔
67	this.outputValues = props.outputValues;	13✔
68	this.algorithm = props.algorithm ?? 'auto';	13✔
69	this.direction = props.direction ?? 'ascending';	13✔
70
71	for (const [name, view] of [	13✔
72	['keys', this.keys],
73	['values', this.values],
74	['outputKeys', this.outputKeys],
75	['outputValues', this.outputValues]
76	] as const) {
77	validatePackedUint32View(view, `${this.id} ${name}`);	49✔
78	}
79	if (!['auto', 'bitonic', 'radix'].includes(this.algorithm)) {	12!
80	throw new Error(`${this.id} algorithm must be auto, bitonic, or radix`);	×
81	}
82	if (!['ascending', 'descending'].includes(this.direction)) {	12!
83	throw new Error(`${this.id} direction must be ascending or descending`);	×
84	}
85	if (	12✔
86	this.values.length !== this.keys.length \|\|	34✔
87	this.outputKeys.length !== this.keys.length \|\|
88	this.outputValues.length !== this.keys.length
89	) {
90	throw new Error(`${this.id} key, value, and output lengths must match`);	1✔
91	}
92	if (this.keys.length > MAXIMUM_LOGICAL_LENGTH) {	11!
93	throw new Error(`${this.id} supports at most ${MAXIMUM_LOGICAL_LENGTH} rows`);	×
94	}
95	validateSeparateWritableBuffers(this);	11✔
96
97	this.resolvedAlgorithm =	10✔
98	this.algorithm === 'auto'	10✔
99	? this.keys.length <= AUTO_BITONIC_MAXIMUM_LENGTH	5✔
100	? 'bitonic'
101	: 'radix'
102	: this.algorithm;
103	}
104
105	/** Adds the selected sort implementation and its scratch buffers to a command graph. */
106	addToGraph<Parameters>(graph: GPUCommandGraph<Parameters>): void {
107	for (const view of [this.keys, this.values, this.outputKeys, this.outputValues]) {	7✔
108	if (view.buffer.graph !== graph) {	25✔
109	throw new Error(`${this.id} views must belong to the target graph`);	1✔
110	}
111	}
112	if (this.keys.length === 0) {	6✔
113	return;	1✔
114	}
115	if (this.keys.length === 1) {	5✔
116	addCopyPairPass(graph, this);	1✔
117	return;	1✔
118	}
119	if (this.resolvedAlgorithm === 'bitonic') {	4✔
120	addBitonicSort(graph, this);	2✔
121	} else {
122	addRadixSort(graph, this);	2✔
123	}
124	}
125	}
126
127	function validateSeparateWritableBuffers(sort: GPUSort): void {
128	if (	11✔
129	sort.outputKeys.buffer === sort.outputValues.buffer \|\|	52✔
130	sort.outputKeys.buffer === sort.keys.buffer \|\|
131	sort.outputKeys.buffer === sort.values.buffer \|\|
132	sort.outputValues.buffer === sort.keys.buffer \|\|
133	sort.outputValues.buffer === sort.values.buffer
134	) {
135	throw new Error(`${sort.id} outputs must use separate buffers from inputs and each other`);	1✔
136	}
137	}
138
139	function addCopyPairPass<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
140	const source = /* wgsl */ `	1✔
141	const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
142	const VALUES_OFFSET: u32 = ${getViewElementOffset(sort.values)}u;
143	const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(sort.outputKeys)}u;
144	const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u;
145	@group(0) @binding(0) var<storage, read> keys: array<u32>;
146	@group(0) @binding(1) var<storage, read> values: array<u32>;
147	@group(0) @binding(2) var<storage, read_write> outputKeys: array<u32>;
148	@group(0) @binding(3) var<storage, read_write> outputValues: array<u32>;
149
150	@compute @workgroup_size(1) fn main() {
151	outputKeys[OUTPUT_KEYS_OFFSET] = keys[KEYS_OFFSET];
152	outputValues[OUTPUT_VALUES_OFFSET] = values[VALUES_OFFSET];
153	}`;
154	addComputationPass(graph, {	1✔
155	id: `${sort.id}-copy-pair`,
156	source,
157	resources: [
158	{buffer: sort.keys, usage: 'storage-read'},
159	{buffer: sort.values, usage: 'storage-read'},
160	{buffer: sort.outputKeys, usage: 'storage-write'},
161	{buffer: sort.outputValues, usage: 'storage-write'}
162	],
163	bindings: {
164	keys: sort.keys,
165	values: sort.values,
166	outputKeys: sort.outputKeys,
167	outputValues: sort.outputValues
168	},
169	dispatchCount: 1
170	});
171	}
172
173	function addBitonicSort<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
174	const paddedLength = getNextPowerOfTwo(sort.keys.length);	2✔
175	const indicesA = createTransientView(	2✔
176	graph,
177	`${sort.id}-bitonic-indices-a`,
178	'uint32',
179	paddedLength
180	);
181	const indicesB = createTransientView(	2✔
182	graph,
183	`${sort.id}-bitonic-indices-b`,
184	'uint32',
185	paddedLength
186	);
187	addBitonicInitializePass(graph, sort, indicesA, paddedLength);	2✔
188
189	let currentIndices = indicesA;	2✔
190	let nextIndices = indicesB;	2✔
191	for (const stage of getBitonicStages(paddedLength)) {	2✔
192	addBitonicStagePass(graph, sort, currentIndices, nextIndices, paddedLength, stage);	20✔
193	[currentIndices, nextIndices] = [nextIndices, currentIndices];	20✔
194	}
195	addBitonicGatherPass(graph, sort, currentIndices);	2✔
196	}
197
198	function addBitonicInitializePass<Parameters>(
199	graph: GPUCommandGraph<Parameters>,
200	sort: GPUSort,
201	indices: GraphBufferView<'uint32'>,
202	paddedLength: number
203	): void {
204	const source = /* wgsl */ `	2✔
205	const INVALID_INDEX: u32 = ${INVALID_INDEX}u;
206	const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
207	const PADDED_LENGTH: u32 = ${paddedLength}u;
208	const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u;
209	@group(0) @binding(0) var<storage, read_write> indices: array<u32>;
210
211	@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
212	@builtin(global_invocation_id) globalId: vec3<u32>
213	) {
214	let index = globalId.x;
215	if (index < PADDED_LENGTH) {
216	indices[INDICES_OFFSET + index] = select(INVALID_INDEX, index, index < LOGICAL_LENGTH);
217	}
218	}`;
219	addComputationPass(graph, {	2✔
220	id: `${sort.id}-bitonic-initialize`,
221	source,
222	resources: [{buffer: indices, usage: 'storage-write'}],
223	bindings: {indices},
224	dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE)
225	});
226	}
227
228	function addBitonicStagePass<Parameters>(
229	graph: GPUCommandGraph<Parameters>,
230	sort: GPUSort,
231	indicesIn: GraphBufferView<'uint32'>,
232	indicesOut: GraphBufferView<'uint32'>,
233	paddedLength: number,
234	stage: BitonicStage
235	): void {
236	const descending = sort.direction === 'descending';	20✔
237	const source = /* wgsl */ `	20✔
238	const INVALID_INDEX: u32 = ${INVALID_INDEX}u;
239	const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
240	const PADDED_LENGTH: u32 = ${paddedLength}u;
241	const BLOCK_WIDTH: u32 = ${stage.blockWidth}u;
242	const COMPARE_STRIDE: u32 = ${stage.compareStride}u;
243	const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
244	const INDICES_IN_OFFSET: u32 = ${getViewElementOffset(indicesIn)}u;
245	const INDICES_OUT_OFFSET: u32 = ${getViewElementOffset(indicesOut)}u;
246	@group(0) @binding(0) var<storage, read> keys: array<u32>;
247	@group(0) @binding(1) var<storage, read> indicesIn: array<u32>;
248	@group(0) @binding(2) var<storage, read_write> indicesOut: array<u32>;
249
250	fn is_valid(index: u32) -> bool {
251	return index != INVALID_INDEX && index < LOGICAL_LENGTH;
252	}
253
254	fn comes_before(leftIndex: u32, rightIndex: u32) -> bool {
255	let leftValid = is_valid(leftIndex);
256	let rightValid = is_valid(rightIndex);
257	if (leftValid != rightValid) { return leftValid; }
258	if (!leftValid) { return false; }
259	let leftKey = keys[KEYS_OFFSET + leftIndex];
260	let rightKey = keys[KEYS_OFFSET + rightIndex];
261	if (leftKey == rightKey) { return leftIndex < rightIndex; }
262	return ${descending ? 'leftKey > rightKey' : 'leftKey < rightKey'};	20✔
263	}
264
265	@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
266	@builtin(global_invocation_id) globalId: vec3<u32>
267	) {
268	let index = globalId.x;
269	if (index >= PADDED_LENGTH) { return; }
270	let partnerIndex = index ^ COMPARE_STRIDE;
271	if (partnerIndex <= index) { return; }
272	let leftIndex = indicesIn[INDICES_IN_OFFSET + index];
273	let rightIndex = indicesIn[INDICES_IN_OFFSET + partnerIndex];
274	let ascending = (index & BLOCK_WIDTH) == 0u;
275	let shouldSwap = select(
276	comes_before(leftIndex, rightIndex),
277	comes_before(rightIndex, leftIndex),
278	ascending
279	);
280	indicesOut[INDICES_OUT_OFFSET + index] = select(leftIndex, rightIndex, shouldSwap);
281	indicesOut[INDICES_OUT_OFFSET + partnerIndex] = select(rightIndex, leftIndex, shouldSwap);
282	}`;
283	addComputationPass(graph, {	20✔
284	id: `${sort.id}-bitonic-${stage.blockWidth}-${stage.compareStride}`,
285	source,
286	resources: [
287	{buffer: sort.keys, usage: 'storage-read'},
288	{buffer: indicesIn, usage: 'storage-read'},
289	{buffer: indicesOut, usage: 'storage-write'}
290	],
291	bindings: {keys: sort.keys, indicesIn, indicesOut},
292	dispatchCount: Math.ceil(paddedLength / BITONIC_WORKGROUP_SIZE)
293	});
294	}
295
296	function addBitonicGatherPass<Parameters>(
297	graph: GPUCommandGraph<Parameters>,
298	sort: GPUSort,
299	indices: GraphBufferView<'uint32'>
300	): void {
301	const source = /* wgsl */ `	2✔
302	const LOGICAL_LENGTH: u32 = ${sort.keys.length}u;
303	const KEYS_OFFSET: u32 = ${getViewElementOffset(sort.keys)}u;
304	const VALUES_OFFSET: u32 = ${getViewElementOffset(sort.values)}u;
305	const INDICES_OFFSET: u32 = ${getViewElementOffset(indices)}u;
306	const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(sort.outputKeys)}u;
307	const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(sort.outputValues)}u;
308	@group(0) @binding(0) var<storage, read> keys: array<u32>;
309	@group(0) @binding(1) var<storage, read> values: array<u32>;
310	@group(0) @binding(2) var<storage, read> indices: array<u32>;
311	@group(0) @binding(3) var<storage, read_write> outputKeys: array<u32>;
312	@group(0) @binding(4) var<storage, read_write> outputValues: array<u32>;
313
314	@compute @workgroup_size(${BITONIC_WORKGROUP_SIZE}) fn main(
315	@builtin(global_invocation_id) globalId: vec3<u32>
316	) {
317	let index = globalId.x;
318	if (index >= LOGICAL_LENGTH) { return; }
319	let sourceIndex = indices[INDICES_OFFSET + index];
320	outputKeys[OUTPUT_KEYS_OFFSET + index] = keys[KEYS_OFFSET + sourceIndex];
321	outputValues[OUTPUT_VALUES_OFFSET + index] = values[VALUES_OFFSET + sourceIndex];
322	}`;
323	addComputationPass(graph, {	2✔
324	id: `${sort.id}-bitonic-gather`,
325	source,
326	resources: [
327	{buffer: sort.keys, usage: 'storage-read'},
328	{buffer: sort.values, usage: 'storage-read'},
329	{buffer: indices, usage: 'storage-read'},
330	{buffer: sort.outputKeys, usage: 'storage-write'},
331	{buffer: sort.outputValues, usage: 'storage-write'}
332	],
333	bindings: {
334	keys: sort.keys,
335	values: sort.values,
336	indices,
337	outputKeys: sort.outputKeys,
338	outputValues: sort.outputValues
339	},
340	dispatchCount: Math.ceil(sort.keys.length / BITONIC_WORKGROUP_SIZE)
341	});
342	}
343
344	function addRadixSort<Parameters>(graph: GPUCommandGraph<Parameters>, sort: GPUSort): void {
345	const scratchKeys = createTransientView(	2✔
346	graph,
347	`${sort.id}-radix-scratch-keys`,
348	'uint32',
349	sort.keys.length
350	);
351	const scratchValues = createTransientView(	2✔
352	graph,
353	`${sort.id}-radix-scratch-values`,
354	'uint32',
355	sort.keys.length
356	);
357	let currentKeys = sort.keys;	2✔
358	let currentValues = sort.values;	2✔
359
360	for (let bit = 0; bit < 32; bit++) {	2✔
361	const flags = createTransientView(	64✔
362	graph,
363	`${sort.id}-radix-bit-${bit}-flags`,
364	'uint32',
365	sort.keys.length
366	);
367	const offsets = createTransientView(	64✔
368	graph,
369	`${sort.id}-radix-bit-${bit}-offsets`,
370	'uint32',
371	sort.keys.length
372	);
373	const nextKeys = bit % 2 === 0 ? scratchKeys : sort.outputKeys;	64✔
374	const nextValues = bit % 2 === 0 ? scratchValues : sort.outputValues;	64✔
375	addRadixClassifyPass(graph, sort, currentKeys, flags, bit);	64✔
376	new GPUScan({	64✔
377	id: `${sort.id}-radix-bit-${bit}-scan`,
378	input: flags,
379	output: offsets
380	}).addToGraph(graph);
381	addRadixScatterPass(	64✔
382	graph,
383	sort,
384	currentKeys,
385	currentValues,
386	flags,
387	offsets,
388	nextKeys,
389	nextValues,
390	bit
391	);
392	currentKeys = nextKeys;	64✔
393	currentValues = nextValues;	64✔
394	}
395	}
396
397	function addRadixClassifyPass<Parameters>(
398	graph: GPUCommandGraph<Parameters>,
399	sort: GPUSort,
400	keys: GraphBufferView<'uint32'>,
401	flags: GraphBufferView<'uint32'>,
402	bit: number
403	): void {
404	const firstBit = sort.direction === 'ascending' ? 0 : 1;	64✔
405	const source = /* wgsl */ `	64✔
406	const ELEMENT_COUNT: u32 = ${sort.keys.length}u;
407	const BIT_INDEX: u32 = ${bit}u;
408	const FIRST_BIT: u32 = ${firstBit}u;
409	const KEYS_OFFSET: u32 = ${getViewElementOffset(keys)}u;
410	const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u;
411	@group(0) @binding(0) var<storage, read> keys: array<u32>;
412	@group(0) @binding(1) var<storage, read_write> flags: array<u32>;
413
414	@compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main(
415	@builtin(global_invocation_id) globalId: vec3<u32>
416	) {
417	let index = globalId.x;
418	if (index >= ELEMENT_COUNT) { return; }
419	let bitValue = (keys[KEYS_OFFSET + index] >> BIT_INDEX) & 1u;
420	flags[FLAGS_OFFSET + index] = select(0u, 1u, bitValue == FIRST_BIT);
421	}`;
422	addComputationPass(graph, {	64✔
423	id: `${sort.id}-radix-bit-${bit}-classify`,
424	source,
425	resources: [
426	{buffer: keys, usage: 'storage-read'},
427	{buffer: flags, usage: 'storage-write'}
428	],
429	bindings: {keys, flags},
430	dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE)
431	});
432	}
433
434	function addRadixScatterPass<Parameters>(
435	graph: GPUCommandGraph<Parameters>,
436	sort: GPUSort,
437	keys: GraphBufferView<'uint32'>,
438	values: GraphBufferView<'uint32'>,
439	flags: GraphBufferView<'uint32'>,
440	offsets: GraphBufferView<'uint32'>,
441	outputKeys: GraphBufferView<'uint32'>,
442	outputValues: GraphBufferView<'uint32'>,
443	bit: number
444	): void {
445	const lastIndex = sort.keys.length - 1;	64✔
446	const source = /* wgsl */ `	64✔
447	const ELEMENT_COUNT: u32 = ${sort.keys.length}u;
448	const LAST_INDEX: u32 = ${lastIndex}u;
449	const KEYS_OFFSET: u32 = ${getViewElementOffset(keys)}u;
450	const VALUES_OFFSET: u32 = ${getViewElementOffset(values)}u;
451	const FLAGS_OFFSET: u32 = ${getViewElementOffset(flags)}u;
452	const OFFSETS_OFFSET: u32 = ${getViewElementOffset(offsets)}u;
453	const OUTPUT_KEYS_OFFSET: u32 = ${getViewElementOffset(outputKeys)}u;
454	const OUTPUT_VALUES_OFFSET: u32 = ${getViewElementOffset(outputValues)}u;
455	@group(0) @binding(0) var<storage, read> keys: array<u32>;
456	@group(0) @binding(1) var<storage, read> values: array<u32>;
457	@group(0) @binding(2) var<storage, read> flags: array<u32>;
458	@group(0) @binding(3) var<storage, read> offsets: array<u32>;
459	@group(0) @binding(4) var<storage, read_write> outputKeys: array<u32>;
460	@group(0) @binding(5) var<storage, read_write> outputValues: array<u32>;
461
462	@compute @workgroup_size(${RADIX_WORKGROUP_SIZE}) fn main(
463	@builtin(global_invocation_id) globalId: vec3<u32>
464	) {
465	let index = globalId.x;
466	if (index >= ELEMENT_COUNT) { return; }
467	let firstOffset = offsets[OFFSETS_OFFSET + index];
468	let firstCount = offsets[OFFSETS_OFFSET + LAST_INDEX] + flags[FLAGS_OFFSET + LAST_INDEX];
469	let isFirst = flags[FLAGS_OFFSET + index] != 0u;
470	let outputIndex = select(firstCount + index - firstOffset, firstOffset, isFirst);
471	outputKeys[OUTPUT_KEYS_OFFSET + outputIndex] = keys[KEYS_OFFSET + index];
472	outputValues[OUTPUT_VALUES_OFFSET + outputIndex] = values[VALUES_OFFSET + index];
473	}`;
474	addComputationPass(graph, {	64✔
475	id: `${sort.id}-radix-bit-${bit}-scatter`,
476	source,
477	resources: [
478	{buffer: keys, usage: 'storage-read'},
479	{buffer: values, usage: 'storage-read'},
480	{buffer: flags, usage: 'storage-read'},
481	{buffer: offsets, usage: 'storage-read'},
482	{buffer: outputKeys, usage: 'storage-write'},
483	{buffer: outputValues, usage: 'storage-write'}
484	],
485	bindings: {keys, values, flags, offsets, outputKeys, outputValues},
486	dispatchCount: Math.ceil(sort.keys.length / RADIX_WORKGROUP_SIZE)
487	});
488	}
489
490	function getNextPowerOfTwo(length: number): number {
491	let paddedLength = 1;	2✔
492	while (paddedLength < length) {	2✔
493	paddedLength *= 2;	8✔
494	}
495	return paddedLength;	2✔
496	}
497
498	function getBitonicStages(paddedLength: number): BitonicStage[] {
499	const stages: BitonicStage[] = [];	2✔
500	for (let blockWidth = 2; blockWidth <= paddedLength; blockWidth *= 2) {	2✔
501	for (let compareStride = blockWidth / 2; compareStride >= 1; compareStride /= 2) {	8✔
502	stages.push({blockWidth, compareStride});	20✔
503	}
504	}
505	return stages;	2✔
506	}
507
508	function addComputationPass<GraphParameters>(
509	graph: GPUCommandGraph<GraphParameters>,
510	props: {
511	id: string;
512	source: string;
513	resources: GraphBufferUse[];
514	bindings: Record<string, GraphBufferView>;
515	dispatchCount: number;
516	}
517	): void {
518	graph.addComputePass({	153✔
519	id: props.id,
520	resources: props.resources,
521	compile: ({device}) => {
522	const computation = new Computation(device, {	153✔
523	id: props.id,
524	source: props.source,
525	shaderLayout: {
526	bindings: Object.keys(props.bindings).map((name, location) => ({	588✔
527	name,
528	type: 'storage' as const,
529	group: 0,
530	location
531	}))
532	}
533	});
534	return {	153✔
535	encode: ({computePass, getBuffer}) => {
536	const bindings: Record<string, Binding> = {};	153✔
537	for (const [name, view] of Object.entries(props.bindings)) {	153✔
538	bindings[name] = getViewBinding(view, getBuffer);	588✔
539	}
540	computation.setBindings(bindings);	153✔
541	computation.dispatch(computePass, props.dispatchCount);	153✔
542	},
543	destroy: () => computation.destroy()	153✔
544	};
545	}
546	});
547	}

visgl / luma.gl / 28617818631

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous