feat: Add @typegpu/sort scaffolding with simple bitonic sort implementation (#2142)

Author: reczkok

apps/typegpu-docs/package.json

@@ -29,10 +29,10 @@
     "@stackblitz/sdk": "^1.11.0",
     "@tailwindcss/vite": "^4.1.18",
     "@typegpu/color": "workspace:*",
-    "@typegpu/concurrent-scan": "workspace:*",
     "@typegpu/geometry": "workspace:*",
     "@typegpu/noise": "workspace:*",
     "@typegpu/sdf": "workspace:*",
+    "@typegpu/sort": "workspace:*",
     "@typegpu/three": "workspace:*",
     "@types/react": "^19.1.8",
     "@types/react-dom": "^19.1.6",

apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.html

@@ -0,0 +1,48 @@
+<canvas></canvas>
+<div id="sort-overlay" hidden>
+  <div id="sort-spinner"></div>
+  <span id="sort-status"></span>
+</div>
+<style>
+  #sort-overlay {
+    position: absolute;
+    inset: 0;
+    flex-direction: column;
+    align-items: center;
+    justify-content: center;
+    gap: 12px;
+    backdrop-filter: blur(8px);
+    background: rgba(23, 23, 36, 0.5);
+    opacity: 0;
+    transition: opacity 0.3s ease-in-out;
+  }
+
+  #sort-overlay:not([hidden]) {
+    display: flex;
+  }
+
+  #sort-overlay.visible {
+    opacity: 1;
+  }
+
+  #sort-spinner {
+    width: 40px;
+    height: 40px;
+    border: 3px solid #2e295f;
+    border-top-color: #6453d2;
+    border-radius: 50%;
+    animation: spin 0.8s linear infinite;
+  }
+
+  #sort-status {
+    color: #c3c4f1;
+    font-size: 0.875rem;
+    font-family: 'Aeonik', sans-serif;
+  }
+
+  @keyframes spin {
+    to {
+      transform: rotate(360deg);
+    }
+  }
+</style>

apps/typegpu-docs/src/examples/algorithms/bitonic-sort/index.ts

@@ -0,0 +1,281 @@
+import tgpu, { d, std } from 'typegpu';
+import {
+  type BitonicSorter,
+  type BitonicSorterOptions,
+  createBitonicSorter,
+  decomposeWorkgroups,
+} from '@typegpu/sort';
+import { randf } from '@typegpu/noise';
+import { fullScreenTriangle } from 'typegpu/common';
+import { defineControls } from '../../common/defineControls.ts';
+
+const maxBufferSize = await navigator.gpu.requestAdapter().then((adapter) => {
+  if (!adapter) {
+    throw new Error('No GPU adapter found');
+  }
+  const limits = adapter.limits;
+  return Math.min(limits.maxStorageBufferBindingSize, limits.maxBufferSize);
+});
+
+const root = await tgpu.init({
+  device: {
+    optionalFeatures: ['timestamp-query'],
+    requiredLimits: {
+      maxStorageBufferBindingSize: maxBufferSize,
+      maxBufferSize: maxBufferSize,
+    },
+  },
+});
+const hasTimestampQuery = root.enabledFeatures.has('timestamp-query');
+const querySet = hasTimestampQuery ? root.createQuerySet('timestamp', 2) : null;
+
+const canvas = document.querySelector('canvas') as HTMLCanvasElement;
+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 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 SortOrderKey = 'ascending' | 'descending' | 'bit-reversed' | 'xor-scatter';
+
+const sortOrders: Record<SortOrderKey, BitonicSorterOptions> = {
+  ascending: {},
+  descending: {
+    compare: (a, b) => {
+      'use gpu';
+      return a > b;
+    },
+    paddingValue: 0,
+  },
+  'bit-reversed': {
+    compare: (a, b) => {
+      'use gpu';
+      return std.reverseBits(a) < std.reverseBits(b);
+    },
+  },
+  'xor-scatter': {
+    compare: (a, b) => {
+      'use gpu';
+      return (a ^ 0xaa) < (b ^ 0xaa);
+    },
+  },
+};
+
+const state = {
+  arraySize: arraySizeOptions[2],
+  sortOrder: 'ascending' as SortOrderKey,
+};
+
+const WORKGROUP_SIZE = 256;
+
+const renderLayout = tgpu.bindGroupLayout({
+  data: {
+    storage: d.arrayOf(d.u32),
+    access: 'readonly',
+  },
+});
+
+const initLayout = tgpu.bindGroupLayout({
+  data: {
+    storage: d.arrayOf(d.u32),
+    access: 'mutable',
+  },
+});
+
+const initSeed = root.createUniform(d.f32, 0);
+
+const fragmentFn = tgpu.fragmentFn({
+  in: { uv: d.vec2f },
+  out: d.vec4f,
+})((input) => {
+  const data = renderLayout.$.data;
+  const arrayLength = data.length;
+
+  const cols = d.u32(std.round(std.sqrt(d.f32(arrayLength))));
+  const rows = d.u32(std.round(arrayLength / cols));
+
+  const col = d.u32(std.floor(input.uv.x * d.f32(cols)));
+  const row = d.u32(std.floor(input.uv.y * d.f32(rows)));
+  const idx = row * cols + col;
+
+  if (idx >= arrayLength) {
+    return d.vec4f(0.1, 0.1, 0.1, 1);
+  }
+
+  const value = data[idx];
+  const normalized = value / 255;
+
+  return d.vec4f(normalized, normalized, normalized, 1);
+});
+
+const initKernel = 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;
+  const idx = input.gid.x + input.gid.y * spanX + input.gid.z * spanY;
+
+  if (idx >= initLayout.$.data.length) {
+    return;
+  }
+
+  randf.seed3(d.vec3f(d.f32(idx & 0xffff), d.f32(idx >> 16), initSeed.$));
+  const n = randf.sample();
+  initLayout.$.data[idx] = d.u32(std.floor(n * 256.0));
+});
+
+const renderPipeline = root.createRenderPipeline({
+  vertex: fullScreenTriangle,
+  fragment: fragmentFn,
+  targets: { format: presentationFormat },
+});
+
+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,
+});
+
+function createSorters(buf: typeof buffer) {
+  return Object.fromEntries(
+    Object.entries(sortOrders).map(([key, opts]) => [key, createBitonicSorter(root, buf, opts)]),
+  ) as Record<SortOrderKey, BitonicSorter>;
+}
+
+let sorters = createSorters(buffer);
+
+function recreateBuffer() {
+  for (const s of Object.values(sorters)) {
+    s.destroy();
+  }
+  buffer.destroy();
+
+  buffer = root.createBuffer(d.arrayOf(d.u32, state.arraySize)).$usage('storage');
+
+  bindGroup = root.createBindGroup(renderLayout, {
+    data: buffer,
+  });
+
+  initBindGroup = root.createBindGroup(initLayout, {
+    data: buffer,
+  });
+
+  sorters = createSorters(buffer);
+}
+
+function generateRandomArray() {
+  const workgroupsTotal = Math.ceil(state.arraySize / WORKGROUP_SIZE);
+  const [workgroupsX, workgroupsY, workgroupsZ] = decomposeWorkgroups(workgroupsTotal);
+
+  initSeed.write(Math.random());
+
+  initPipeline.with(initBindGroup).dispatchWorkgroups(workgroupsX, workgroupsY, workgroupsZ);
+
+  render();
+}
+
+function render() {
+  renderPipeline.withColorAttachment({ view: context }).with(bindGroup).draw(3);
+}
+
+const overlay = document.getElementById('sort-overlay') as HTMLDivElement;
+const spinnerEl = document.getElementById('sort-spinner') as HTMLDivElement;
+const statusEl = document.getElementById('sort-status') as HTMLSpanElement;
+canvas.parentElement?.appendChild(overlay);
+
+let hideTimeoutId: ReturnType<typeof setTimeout> | null = null;
+
+function showOverlay(text: string, showSpinner = true) {
+  if (hideTimeoutId !== null) {
+    clearTimeout(hideTimeoutId);
+    hideTimeoutId = null;
+  }
+  spinnerEl.hidden = !showSpinner;
+  statusEl.textContent = text;
+  overlay.hidden = false;
+  overlay.classList.add('visible');
+}
+
+function hideOverlay(delayMs = 1500) {
+  hideTimeoutId = setTimeout(() => {
+    hideTimeoutId = null;
+    overlay.classList.remove('visible');
+    overlay.addEventListener('transitionend', () => (overlay.hidden = true), {
+      once: true,
+    });
+  }, delayMs);
+}
+
+async function sort() {
+  const sorter = sorters[state.sortOrder];
+
+  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;
+  }
+
+  render();
+
+  const timeStr =
+    gpuTimeMs !== null
+      ? ` in ${
+          gpuTimeMs >= 1000 ? `${(gpuTimeMs / 1000).toFixed(2)}s` : `${gpuTimeMs.toFixed(2)}ms`
+        }`
+      : '';
+  showOverlay(`\u2714 Sorted${timeStr}`, false);
+  hideOverlay();
+}
+
+// #region Example controls & Cleanup
+
+const sortOrderKeys = Object.keys(sortOrders) as SortOrderKey[];
+
+export const controls = defineControls({
+  'Array Size': {
+    initial: arraySizeOptions[2],
+    options: arraySizeOptions,
+    onSelectChange: (value) => {
+      state.arraySize = isNaN(value) ? 64 : value;
+      recreateBuffer();
+      generateRandomArray();
+    },
+  },
+  'Sort Order': {
+    initial: 'ascending',
+    options: sortOrderKeys,
+    onSelectChange: (value) => {
+      state.sortOrder = value;
+    },
+  },
+  Reshuffle: { onButtonClick: generateRandomArray },
+  Sort: { onButtonClick: sort },
+});
+
+export function onCleanup() {
+  for (const s of Object.values(sorters)) {
+    s.destroy();
+  }
+  root.destroy();
+}
+
+// #endregion

apps/typegpu-docs/src/examples/algorithms/bitonic-sort/meta.json

@@ -0,0 +1,5 @@
+{
+  "title": "Bitonic Sort",
+  "category": "algorithms",
+  "tags": ["experimental", "compute"]
+}

apps/typegpu-docs/src/examples/algorithms/bitonic-sort/thumbnail.png

@@ -1,4 +1,4 @@
-import { initCache, prefixScan } from '@typegpu/concurrent-scan';
+import { createPrefixScanComputer, prefixScan } from '@typegpu/sort';
 import type { TgpuRoot } from 'typegpu';
 import { d, std } from 'typegpu';
 
@@ -45,7 +45,7 @@ export async function performCalculationsWithTime(
   const jsTime = performance.now() - jsStartTime;
 
   // GPU version
-  initCache(root, { operation: std.add, identityElement: 0 });
+  createPrefixScanComputer(root, { operation: std.add, identityElement: 0 });
   const querySet = root.createQuerySet('timestamp', 2);
   const gpuStartTime = performance.now();
   const calcResult = prefixScan(

apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts

@@ -1,7 +1,7 @@
 import tgpu from 'typegpu';
 import * as d from 'typegpu/data';
 import * as std from 'typegpu/std';
-import type { BinaryOp } from '@typegpu/concurrent-scan';
+import type { BinaryOp } from '@typegpu/sort';
 
 // tgpu functions
 

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

@@ -1,6 +1,6 @@
 import tgpu from 'typegpu';
 import * as d from 'typegpu/data';
-import { type BinaryOp, prefixScan, scan } from '@typegpu/concurrent-scan';
+import { type BinaryOp, prefixScan, scan } from '@typegpu/sort';
 import * as std from 'typegpu/std';
 import { addFn, concat10, isArrayEqual, mulFn, prefixScanJS, scanJS } from './functions.ts';
 

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

@@ -137,11 +137,11 @@ export const SANDBOX_MODULES: Record<string, SandboxModuleDefinition> = {
     import: { reroute: 'typegpu-color/src/index.ts' },
     typeDef: { reroute: 'typegpu-color/src/index.ts' },
   },
-  '@typegpu/concurrent-scan': {
-    import: { reroute: 'typegpu-concurrent-scan/src/index.ts' },
-    typeDef: { reroute: 'typegpu-concurrent-scan/src/index.ts' },
-  },
   '@typegpu/three': {
     typeDef: { reroute: 'typegpu-three/src/index.ts' },
   },
+  '@typegpu/sort': {
+    import: { reroute: 'typegpu-sort/src/index.ts' },
+    typeDef: { reroute: 'typegpu-sort/src/index.ts' },
+  },
 };