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 20 commits into
Conversation
|
pkg.pr.new packages benchmark commit |
Bundle size comparison (
|
| 🟢 Decreased | ➖ Unchanged | 🔴 Increased | ❔ Unknown |
|---|---|---|---|
| 0 | 324 | 0 | 0 |
import { ... } in PR vs import * as ... in PR (is the library tree-Shakeable?):
| Test | tsdown |
|---|---|
| tgpu_init.ts | 274.51 kB ( |
| tgpu_initFromDevice.ts | 273.97 kB ( |
| tgpu_resolve.ts | 174.80 kB ( |
| tgpu_resolveWithContext.ts | 174.74 kB ( |
| tgpu_bindGroupLayout.ts | 74.47 kB ( |
| tgpu_mutableAccessor.ts | 69.19 kB ( |
| tgpu_accessor.ts | 69.19 kB ( |
| tgpu_privateVar.ts | 67.88 kB ( |
| tgpu_workgroupVar.ts | 67.88 kB ( |
| tgpu_const.ts | 67.30 kB ( |
| tgpu_lazy.ts | 67.10 kB ( |
| tgpu_fragmentFn.ts | 39.65 kB ( |
| tgpu_fn.ts | 39.60 kB ( |
| tgpu_vertexFn.ts | 39.47 kB ( |
| tgpu_computeFn.ts | 39.17 kB ( |
| tgpu_vertexLayout.ts | 28.30 kB ( |
| tgpu_comptime.ts | 15.91 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.86, 1.67, 3.78, 5.95, 6.38, 8.92, 20.32, 22.55]
line [0.81, 1.66, 3.53, 5.61, 6.57, 9.23, 19.94, 22.59]
line [0.87, 1.70, 3.53, 5.41, 6.55, 11.56, 20.96, 22.46]
---
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.27, 0.47, 0.59, 0.69, 0.94, 1.00, 1.18, 1.33]
line [0.29, 0.44, 0.61, 0.73, 0.98, 1.03, 1.19, 1.32]
line [0.33, 0.50, 0.60, 0.71, 0.97, 1.04, 1.26, 1.36]
---
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.82, 1.84, 3.40, 5.96, 10.61, 22.70, 49.46, 99.62]
line [0.74, 1.95, 3.40, 6.18, 10.71, 23.26, 48.72, 99.10]
line [0.73, 1.89, 3.39, 6.02, 10.88, 22.92, 48.80, 98.11]
|
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.
![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) ｜ 𝕏
dad0228 to
cd07350
Compare
cd07350 to
7746f74
Compare
There was a problem hiding this comment.
ℹ️ No new issues in the new commits — the shift-operator cleanup is correct and behavior-identical. Not marking this as an approval only because the previously-flagged
bindPassdead branch inrunPass.tsis still open and unchanged by these commits (the data-correctness gap was intentionally deferred in the prior review note).
Reviewed changes
Incremental re-review of the delta since the prior Pullfrog review (head dad02286). The branch was rebased onto a newer main, so most of cd07350..7746f74 is upstream content (docs/blog/examples, 0.12 posts); scoping to the PR's own files, the net change is a small shift-operator pass plus test updates:
- Converted
u32>>/>>=to>>>/>>>=in the GPU code acrossradix/schemas.ts(digit extraction),radix/scatter.ts(bitset-word index),scan/kernels.ts(scan span/offset halving), andbitonic/bitonicSort.ts(lane indexing). - Strengthened the package warning tests (
bitonic/radix/scan): each now assertsconsole.warnis never called instead of over-filtering for "Implicit conversions", and dropped the now-unusedgetConversionWarningshelper fromtests/utils.ts. - Refreshed WGSL snapshot assertions to match current codegen output (statement brace wrapping, unroll markers).
I verified the shift change is behavior-identical: in the TypeGPU DSL both >> and >>> on a u32 operand lower to the same WGSL >> token; the only difference is that >> on u32 emits a deprecation warning (wgslGenerator.ts:583) that >>> avoids. The i32 path correctly keeps >> (arithmetic shift, not deprecated), and digitOfI32 is unchanged. Deprecation warnings route to console.warn (tgpuLogger.ts:73), so the strengthened no-warning assertions would genuinely catch a regression back to >> — they have teeth.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
There was a problem hiding this comment.
ℹ️ No new issues in commit
b0ad345a— it cleanly addresses the previously-flagged dead branch and adds a necessary bounds guard. Not marking this as an approval because the prior review's data-correctness-testing gap remains unaddressed (intentionally deferred by the author).
Reviewed changes
Re-review of the single new commit b0ad345a ("widen api and fix stupid") landed after the last review (head 7746f74). Three small changes:
- Removed the dead
if ('resourceType' in pass)branch inbindPass(runPass.ts) — now a singlepipeline.with(pass), which is valid sincewith()handles bothTgpuComputePassand a rawGPUComputePassEncoder. This addresses the prior review's open thread. - Added an early-return guard
if (workgroupId >= applySumsLayout.$.sums.length) return;at the top ofmakeApplySumsKernel(scan/kernels.ts). - Consolidated the
with(pass: TgpuComputePass)andwith(pass: GPUComputePassEncoder)overloads into a singlewith(pass: TgpuComputePass | GPUComputePassEncoder)(computePipeline.ts) — pure simplification, no behavior change.
I verified the applySums guard is not just dead defensiveness: decomposeWorkgroups (dispatch.ts:10-20) can produce a grid larger than the requested count for large buffers (e.g. [65535,2,1] for total > 65535), so flatWorkgroupIndex can exceed sums.length, which previously read sums[workgroupId] out of bounds (silently returning 0 in WGSL). The guard closes that path and is consistent with the equivalent guard already present in makeScanKernel.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
There was a problem hiding this comment.
ℹ️ No new critical issues in commit
5d44384("fix stupid") — one minor type-assertion nitpick inline below. Not marked as an approval because the prior review's data-correctness-testing gap for radix/bitonic/scan remains unaddressed (intentionally deferred by the author), and this delta is behavior-neutral.
Reviewed changes
Incremental re-review of the single new commit 5d44384 ("fix stupid") landed after the last review (head b0ad345a). Three small changes:
- Reverted the
with(pass: TgpuComputePass | GPUComputePassEncoder)overload consolidation incomputePipeline.tsback to the two separate overloads (with(pass: TgpuComputePass)andwith(pass: GPUComputePassEncoder)), restoring the pre-PRwith()signature. Net diff of this file versusmainis empty. - Added
pass as TgpuComputePassinbindPass(runPass.ts) so the union-typedRunPasssatisfies the now-splitwith()overloads. - Removed an accessor "throws when
$is accessed in simulation mode" test (accessor.test.ts) that was added earlier in the PR; the file now equals itsmainbaseline (net diff empty).
I verified the cast is functionally safe: TgpuComputePipelineImpl.with (computePipeline.ts:271-291) dispatches on isTgpuComputePass / isGPUComputePassEncoder, so a real GPUComputePassEncoder is still adopted correctly at runtime regardless of the assertion. The removed accessor test asserted behavior that still holds (makeDereferenceable.ts:58-63 falls through to normalMode.get, which throws), so its removal is a coverage-only loss, not a behavior change.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
| } | ||
|
|
||
| export function bindPass(pipeline: TgpuComputePipeline, pass: RunPass): TgpuComputePipeline { | ||
| return pipeline.with(pass as TgpuComputePass); |
There was a problem hiding this comment.
This as TgpuComputePass assertion is unsound: RunPass = GPUComputePassEncoder | TgpuComputePass, so when a raw GPUComputePassEncoder is passed the assertion lies about its type. It only works by accident of the runtime type-guards in TgpuComputePipelineImpl.with (computePipeline.ts:271-291). This cast exists solely because the prior commit's single union overload (with(pass: TgpuComputePass | GPUComputePassEncoder)), which accepted RunPass with no cast, was reverted. Consider a union-aware typed helper so the compiler enforces the real contract instead of an unchecked assertion.
There was a problem hiding this comment.
It's just TS not understanding overloads, we don't care
There was a problem hiding this comment.
ℹ️ No new issues in the new commits — the eager
initSync/initAsyncinitialization API is correctly implemented and well-tested. Not marking this as an approval because the previously-flagged data-correctness-testing gap remains open (author-deferred), and therunPass.ts:12cast nitpick is still open and unchanged by this delta.
Reviewed changes
Incremental re-review of the delta since the prior Pullfrog review (head 5d44384). Single commit 6347db51 ("remove useless feature and add initSync/initAsync to sort packaga APIs") across 9 files:
- Added
initSync()/initAsync()to theSorterinterface (types.ts) andPrefixScanPlan(prefixScan.ts), implemented in all three sorters/plans (bitonic, radix, scan). Each iterates its recorded steps and delegates to the existingTgpuComputePipeline.initSync/initAsync— I verified the underlying API exists (computePipeline.ts:449-509), that.with(bindGroup)wrappers share the pipeline core (so initializing the base covers all steps), and thatinitSyncwhile aninitAsyncis in flight throws (pre-existing, documented pipeline behavior). - Scan
prepare(): hoistedapplySums()out of the level loop and added theapplyLevels.length > 0guard. Behavior-identical to the old code (the old loop body only ran when levels were non-empty, soapplySums()was never called for the single-workgroup case); the guard preserves exactly that. - Bitonic: hoisted
padPipeline/unpadPipelinecreation into named locals — pure cleanup, no behavior change. - New tests assert exact eager-init pipeline counts (bitonic
3, radix3, scan1for reduce then3total). I traced the numbers against the mock device and they are consistent — the radix test also confirmsrun()afterinitAsync()makes zero synchronouscreateComputePipelinecalls, which gives the assertions teeth. prefix-scandocs example: droppedrequiredFeatures: ['timestamp-query']fromtgpu.init— the example uses no query sets (verified), so this is a safe cleanup.
All 22 typegpu-sort package tests pass and tsc --noEmit is clean.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
|
Run failed. View the logs →
|
aleksanderkatan
left a comment
aleksanderkatan
left a comment
There was a problem hiding this comment.
What I've ween looks good! 👍
Please, split such PRs into smaller ones. It is unreasonable to expect a proper review of one PR of this size and we keep doing this over and over again. It's exhausting, and probably introduces many hidden errors to our codebase.
| * For `f32` keys sorted ascending, NaNs with a cleared sign bit sort after +Infinity and | ||
| * NaNs with a set sign bit sort before -Infinity. -0 and +0 compare equal. |
There was a problem hiding this comment.
I'd skip this note. We don't throw on writing these values to buffers (only when casting with a schema), but the spec still says "Implementations may assume that overflow, infinities, and NaNs are not present during shader execution."
| * For `f32` keys sorted ascending, NaNs with a cleared sign bit sort after +Infinity and | |
| * NaNs with a set sign bit sort before -Infinity. -0 and +0 compare equal. | |
| * For `f32` keys sorted ascending, -0 and +0 compare equal. |
| export function defaultCompare(a: number, b: number): boolean { | ||
| 'use gpu'; | ||
| return a < b; | ||
| } |
There was a problem hiding this comment.
Do we plan on updating & allowing std functions here?
There was a problem hiding this comment.
Another PR! This needs changes in std, and I'd like to talk through the signatures of std functions (we had a soft convention that when JS supports it, std shouldn't, this would break it)
| } | ||
|
|
||
| export function bindPass(pipeline: TgpuComputePipeline, pass: RunPass): TgpuComputePipeline { | ||
| return pipeline.with(pass as TgpuComputePass); |
There was a problem hiding this comment.
It's just TS not understanding overloads, we don't care
I actually don't quite get why separating this one would be any easier on reviewers. The 2500 lines of additions are pretty neatly split into: examples + docs, logic, and tests and even the files are in this order. I can apply a stacked PR pattern but the dependancies between the changes are significant (so the core logic would land as one PR unless I artificially split which is more work). What that would probably do (and have already done in the soul + worklet stuff) is double maintenance and extend the review time. I will apply this pattern from now on (since that has become a repo pattern) but I see no point in shaming a PR that happened before stacked PRs were even a thing |
Maybe we just are built different. I spent 2.5 hours reviewing the refactor, I am 29 out of 44 files through, and I already am so tired that I don't want to do reviews for the rest of the day. And I didn't even read the radix sort code yet. Unless we don't care about reviews and a quick skim is good enough, but then why review at all? |
Nice! |
No description provided.