feat: Add radix to sort package and improve existing algos - #2756
feat: Add radix to sort package and improve existing algos#2756reczkok wants to merge 15 commits into
Conversation
|
pkg.pr.new packages benchmark commit |
Bundle size comparison (
|
| 🟢 Decreased | ➖ Unchanged | 🔴 Increased | ❔ Unknown |
|---|---|---|---|
| 0 | 322 | 0 | 0 |
import { ... } in PR vs import * as ... in PR (is the library tree-Shakeable?):
| Test | tsdown |
|---|---|
| tgpu_init.ts | 267.87 kB ( |
| tgpu_initFromDevice.ts | 267.34 kB ( |
| tgpu_resolve.ts | 168.44 kB ( |
| tgpu_resolveWithContext.ts | 168.38 kB ( |
| tgpu_bindGroupLayout.ts | 73.79 kB ( |
| tgpu_mutableAccessor.ts | 68.51 kB ( |
| tgpu_accessor.ts | 68.51 kB ( |
| tgpu_privateVar.ts | 67.20 kB ( |
| tgpu_workgroupVar.ts | 67.20 kB ( |
| tgpu_const.ts | 66.62 kB ( |
| tgpu_lazy.ts | 66.42 kB ( |
| tgpu_fragmentFn.ts | 38.92 kB ( |
| tgpu_fn.ts | 38.86 kB ( |
| tgpu_vertexFn.ts | 38.73 kB ( |
| tgpu_computeFn.ts | 38.44 kB ( |
| tgpu_vertexLayout.ts | 27.57 kB ( |
| tgpu_comptime.ts | 15.17 kB ( |
| tgpu_unroll.ts | 1.75 kB ( |
| tgpu_slot.ts | 1.70 kB ( |
If you wish to run a comparison for other, slower bundlers, run the 'Tree-shake test' from the GitHub Actions menu.
Resolution Time Benchmark---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Random Branching (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.91, 1.78, 3.95, 6.02, 6.87, 11.50, 20.74, 24.37]
line [0.95, 1.90, 3.96, 6.13, 7.11, 10.61, 21.72, 26.10]
line [0.96, 1.84, 4.34, 6.69, 7.32, 11.34, 21.69, 24.56]
---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Linear Recursion (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.29, 0.47, 0.66, 0.78, 1.06, 1.13, 1.33, 1.48]
line [0.30, 0.50, 0.66, 0.81, 1.28, 1.20, 1.33, 1.50]
line [0.34, 0.54, 0.63, 0.89, 1.10, 1.16, 1.43, 1.56]
---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Full Tree (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.85, 2.14, 4.18, 6.99, 12.70, 24.87, 53.02, 106.38]
line [0.82, 1.90, 3.41, 6.64, 12.62, 25.59, 52.27, 106.72]
line [0.95, 2.07, 4.58, 6.66, 13.19, 27.70, 55.10, 112.12]
|
d6b4ea3 to
2adb4ae
Compare
There was a problem hiding this comment.
Missing correctness tests for the sort itself
There was a problem hiding this comment.
Well no good way to run the shaders without introducing a lot of testing scaffolding, the sort example is a correctness test
2adb4ae to
dad0228
Compare
There was a problem hiding this comment.
Pull request overview
Adds a new @typegpu/sort radix sorter and refactors the existing scan/bitonic implementations to better reuse GPU resources (pipelines/bind groups/scratch buffers) and to compose with externally managed command encoders / compute passes.
Changes:
- Introduces stable LSD radix sort (with optional payload reordering) plus test coverage.
- Refactors prefix scan into reusable cached “plans” and shared dispatch/run-pass utilities.
- Updates docs/examples/tests to reflect the new APIs and expanded algorithm set.
Reviewed changes
Copilot reviewed 41 out of 45 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pnpm-lock.yaml | Adds typegpu-testing-utility workspace dependency for the repo toolchain. |
| packages/typegpu/tests/accessor.test.ts | Adds regression test for accessor $ access during simulation mode. |
| packages/typegpu-testing-utility/src/extendedIt.ts | Extends mocked device limits (workgroup storage size) to cover new kernels. |
| packages/typegpu-sort/vitest.config.mts | Adds Vitest config for the new typegpu-sort package tests. |
| packages/typegpu-sort/tests/utils.ts | Adds test helpers for extracting generated WGSL and conversion warnings. |
| packages/typegpu-sort/tests/scan.test.ts | Adds scan tests for warnings, schemas, caching behavior, and validation errors. |
| packages/typegpu-sort/tests/radix.test.ts | Adds radix sort tests for warnings, digit extraction behavior, and caching. |
| packages/typegpu-sort/tests/bitonic.test.ts | Adds/updates bitonic tests incl. payload behavior and “no allocations/writes” expectations. |
| packages/typegpu-sort/src/types.ts | Introduces shared Sorter + RunOptions abstraction for composing into encoders/passes. |
| packages/typegpu-sort/src/scan/types.ts | Extends scan BinaryOp typing to carry optional element type (d.u32/d.i32/d.f32). |
| packages/typegpu-sort/src/scan/schemas.ts | Reworks scan schemas to be element-type aware; defines constants and per-type layouts/slots. |
| packages/typegpu-sort/src/scan/prefixScan.ts | Major refactor: prepares reusable plans, caches per-buffer plans, adds reduce, and supports RunOptions. |
| packages/typegpu-sort/src/scan/kernels.ts | New scan/apply-sums kernels using shared dispatch helpers and unrolled per-thread processing. |
| packages/typegpu-sort/src/scan/index.ts | Updates scan exports: adds reduce, plan types, buffer types, and element-type exports. |
| packages/typegpu-sort/src/scan/compute/shared.ts | Removes old scan shared compute helpers (replaced by kernels.ts). |
| packages/typegpu-sort/src/scan/compute/scan.ts | Removes old scan kernel implementation (replaced by kernels.ts). |
| packages/typegpu-sort/src/scan/compute/applySums.ts | Removes old apply-sums kernel implementation (replaced by kernels.ts). |
| packages/typegpu-sort/src/runPass.ts | Adds shared helpers to run work either standalone or recorded into an existing encoder/pass. |
| packages/typegpu-sort/src/radix/types.ts | Adds radix sorter options (direction + optional payload buffer). |
| packages/typegpu-sort/src/radix/schemas.ts | Adds radix constants, layouts, workgroup memory, and digit extraction logic per key type/direction. |
| packages/typegpu-sort/src/radix/scatter.ts | Adds radix scatter kernel (tile-based, uses workgroup bitsets + ranks). |
| packages/typegpu-sort/src/radix/radixSort.ts | Implements radix sorter orchestration, resource ownership, and multi-pass execution. |
| packages/typegpu-sort/src/radix/index.ts | Exposes radix sorter public API. |
| packages/typegpu-sort/src/radix/count.ts | Adds radix histogram count kernel. |
| packages/typegpu-sort/src/index.ts | Updates package exports to include radix and new scan APIs/types. |
| packages/typegpu-sort/src/dispatch.ts | Adds shared dispatch decomposition + flat workgroup indexing helper for kernels. |
| packages/typegpu-sort/src/bitonic/utils.ts | Removes old bitonic dispatch utilities (replaced by shared dispatch.ts). |
| packages/typegpu-sort/src/bitonic/types.ts | Aligns bitonic types with shared Sorter/RunOptions and adds optional payload support. |
| packages/typegpu-sort/src/bitonic/slots.ts | Refactors comparator/defaults; introduces per-key-type default padding values. |
| packages/typegpu-sort/src/bitonic/index.ts | Simplifies bitonic exports and aligns with new public API surface. |
| packages/typegpu-sort/src/bitonic/bitonicSort.ts | Major refactor: pre-creates resources, adds local/shared-memory kernels, and supports payload buffers. |
| packages/typegpu-sort/README.md | Updates docs: adds radix sort section, updates scan API (reduce), and documents encoder/pass composition. |
| packages/typegpu-sort/package.json | Adds typegpu-testing-utility dependency for tests/config. |
| apps/typegpu-docs/tests/individual-example-tests/sort.test.ts | Adds/updates docs example test snapshots for the new consolidated sort example. |
| apps/typegpu-docs/tests/individual-example-tests/bitonic-sort.test.ts | Removes old dedicated bitonic sort example test. |
| apps/typegpu-docs/src/examples/tests/prefix-scan/index.ts | Updates example test to use reduce instead of scan. |
| apps/typegpu-docs/src/examples/algorithms/sort/meta.json | Renames example from “Bitonic Sort” to “Sort”. |
| apps/typegpu-docs/src/examples/algorithms/sort/index.ts | Updates example to include radix option + benchmarking and revised timing integration. |
| apps/typegpu-docs/src/examples/algorithms/sort/index.html | Adds overlay/spinner UI for sort/benchmark actions. |
| apps/typegpu-docs/src/examples/algorithms/concurrent-chart/meta.json | Removes concurrent-chart example metadata. |
| apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.ts | Removes concurrent-chart example implementation. |
| apps/typegpu-docs/src/examples/algorithms/concurrent-chart/index.html | Removes concurrent-chart example markup/styles. |
| apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts | Removes concurrent-chart example compute/timing helper. |
Files not reviewed (1)
- pnpm-lock.yaml: Generated file
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ({ lid, wid, numWorkgroups }) => { | ||
| const workgroupId = flatWorkgroupIndex(wid, numWorkgroups); | ||
| const baseIdx = (workgroupId * WORKGROUP_SIZE + lid.x) * ELEMENTS_PER_THREAD; | ||
| const blockSum = applySumsLayout.$.sums[workgroupId]; | ||
|
|
![pullfrog[bot]](https://avatars.githubusercontent.com/in/1768019?s=60&v=4){kind=small}
There was a problem hiding this comment.
(Note from @reczkok)
Mr. Frog pulled the alarm sirens but there is no cause for panic, the Radix sorter is tested via the sort example - we can add some automated testing via browser or deno in the future. Also you can't break an API that isn't public so yeah, false alarm
Important
The headline feature — the new radix sorter — plus the reworked bitonic sorter and prefix scan ship with no automated correctness test. Every new test asserts resolved-WGSL snapshots, the absence of conversion warnings, or GPU-resource reuse counts; none runs a sort/scan on real data and checks the output. A regression in the scatter rank/prefix logic, the digit transforms, or the scan's exclusive-prefix semantics would pass CI silently. Please add at least one data-correctness test (see the section below) before merging.
Reviewed changes
- New stable LSD radix sorter (
createRadixSorter) foru32/i32/f32keys with optional payload reordering and ascending/descending directions, plus itscount/scatterkernels and schema layer. - Bitonic sorter reworked: typed key support, a
valuespayload path, block-local workgroup-shared kernels with a shared-memory-size-gated global fallback, and padding moved to per-type defaults. - Prefix scan rewritten around a reusable
createPrefixScanComputer+prepare()plan API, integer element types, anencoder/passcomposition surface, and areducehelper replacing the oldscan. - Shared
dispatch.ts/runPass.tsinfrastructure and newRunOptions/Sorter/PrefixScanPlantypes. - Docs:
sortexample rewritten to showcase both sorters (with a benchmark harness), newsort.test.ts, and deletion of the oldbitonic-sort.test.tsandconcurrent-chartexample. - New package tests,
vitest.config.mts, and atypegpu-testing-utilitymock-limit addition.
I traced the radix algorithm end-to-end (count → flat exclusive scan over the digit-major hist layout → bitset-ranked scatter) and it is correct, including the partial-final-tile bounds handling, i32/f32 digit transforms, stable within-tile placement, and descending direction.
⚠️ No data-correctness test for the radix sorter (and the rewritten bitonic/scan)
All tests in packages/typegpu-sort/tests/* and the docs sort.test.ts are codegen snapshots plus device-mock resource-count assertions. Nothing verifies that sorter.run() leaves the buffer actually sorted, or that the scan leaves the buffer actually exclusive-prefix-scanned. radix.test.ts in particular (the only genuinely new algorithm) covers digit functions and resource reuse but never checks a real ordering result — a bug in the scatter rank/prefix or an accidental switch of the flat scan from exclusive to inclusive would go completely unnoticed by CI.
Technical details
# Add a radix data-correctness test
## Affected sites
- packages/typegpu-sort/src/radix/{count,scatter,radixSort}.ts — the new algorithm, entirely untested on data
- packages/typegpu-sort/tests/radix.test.ts — today only WGSL snapshots and resource-reuse counts
- packages/typegpu-sort/tests/{bitonic,scan}.test.ts — same codegen-only pattern
## Required outcome
- A test that computes a reference result on CPU (e.g. `[...arr].sort(...)` for keys, and an index payload that must land exactly on its key's sorted position), runs the sorter, reads the GPU buffer back, and asserts equality — covering at least: u32 ascending and descending, a non-power-of-2 length (partial final tile + bounds checks), and a payload buffer reordered alongside the keys.
- Analogous check (even a single fixed small input) that `prefixScan` yields the exact exclusive prefix vs a JS reference, since the entire radix scheme depends on that flat scan being exclusive.
## Notes
- The mock `device` from `typegpu-testing-utility` is shader/codegen-oriented, so a real correctness test likely needs a run-and-compare harness (like the existing `apps/typegpu-docs/src/examples/tests/prefix-scan/`) or a test on a real device — worth aligning with however the repo validates numeric results elsewhere.ℹ️ Breaking public-API changes to @typegpu/sort
This PR removes the scan export (renamed reduce) and drops compareSlot, defaultCompare, decomposeWorkgroups, and BitonicSorterRunOptions from the package's public index (compareSlot/defaultCompare are now internal to the bitonic module). Monorepo callers were updated and the now-obsolete examples deleted, so nothing internal breaks — but this is a breaking release for external consumers. Worth confirming the removal of compareSlot/defaultCompare from the public surface is intended (custom comparators still work via BitonicSorterOptions.compare).
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
| } | ||
|
|
||
| export function bindPass(pipeline: TgpuComputePipeline, pass: RunPass): TgpuComputePipeline { | ||
| if ('resourceType' in pass) { |
There was a problem hiding this comment.
Both branches of this if ('resourceType' in pass) are identical (return pipeline.with(pass)), so the guard is dead code and suggests two paths are handled differently when they aren't. Since TgpuComputePipeline.with() already accepts both TgpuComputePass and a raw GPUComputePassEncoder, this can collapse to a single return pipeline.with(pass);.
dad0228 to
cd07350
Compare
No description provided.