Repository navigation
Clean up names, comments, file layout, tests and docs across the package - #179
Merged
Merged
Conversation
prototype/ calls APIs that no longer exist (merge_sort!, block_size and temp keywords, OhMyThreads schedulers), and its parallel sample sort became src/sort/cpu_sample_sort.jl. The Pluto notebook in benchmark/ is unreferenced and mostly an embedded manifest. mapreduce_1d.jl benchmarked accumulate instead of reduce, and the AMDGPU branch of runbenchmarks.jl never defined ArrayType. CUDA.jl sorts permutations now, so sortperm.jl no longer needs excluding there. The plotting packages were only used by commented-out code.
The package read README.md into a constant at load time only so the docs could transclude its sections; that helper now lives in docs/, which drops the Markdown dependency. TypeWrap and the u8..i64 literals are unused, and Base.ispow2 replaces AK's own. The CUDA weakdep stays: its compat entry keeps CUDA 5, which has no CUDACore and so no extension, from being installed alongside AK (#122).
It was not a single-pass scan: the lookback ran as a second kernel after every block had finished, so it moved more data than ScanPrefixes, and since how it grouped elements depended on which blocks had published, floating-point results varied between runs. `Auto` never chose it (no tuning set `prefer_lookback`), and its docs claimed a forward-progress requirement that a kernel launched after all blocks finished does not have. With it go `prefer_lookback`, `_supports_lookback`, the device fences in the extensions, the block flags, and the kernels only it used. The AMDGPU extension held nothing else and is removed, as is the UnsafeAtomics dependency.
Nothing calls `unrolled_map_index`, the `_si0` binary searches, or the `comp` methods of `_searchsortedfirst`/`_searchsortedlast` (every caller passes an Ordering). The merge sort host loops doubled a `size_group` they never read, the radix scatter kernel had a second name for its block size, and `mapreduce_nd!` picked a block size for the threaded path, which returns before using it. `IntLinSpace` had one use; its `i == length` branch compared against `Base.length` and never ran, so the general formula is inlined there. The sample sort's commented-out single-threaded copy goes too. `@inbounds @inline function` annotated the definition, not the body, so the merge sort's binary searches kept their bounds checks; they now propagate the kernels' `inbounds=true`.
`reduce_group!`, `block_exclusive_scan!` with its sub-group width helper, and the vector loads are used by several families: the ordered reduction scans with the block scan from accumulate/, which was included after it. They now live in one file included before the families.
The accumulator type rule, seeds, lanes and `_finish` are used by sort, reduce and accumulate alike, so they move out of reduce/ into a file included before the families. The rest of reduce/utilities.jl goes to its users: `_BoolValued` to predicates.jl, `_mapreduce_nd_single!` to the nd reduction, and `_unval` to utils.jl.
mapreduce_nd.jl held some 1250 lines of four kinds of code. The source layout (dims canonicalization, strides, offsets, vector widths) goes to layout.jl, the shape constants and selection to shapes.jl, and the kernels with their index helpers to kernels_nd.jl; mapreduce_nd.jl keeps the entry point, the launchers and the threaded CPU path.
arithmetics.jl held shorthands for two unrelated families. The reductions move to reduce/shorthands.jl and the scans to the end of accumulate/accumulate.jl, next to the operations they call.
Reduce and accumulate define their algorithms in tuning.jl, next to the tuning that fills and selects them; sort now does too, and includes tuning.jl first like the other families.
After the dead searches went, sort/utils.jl held only the two zero-based searches the merge sort kernels use. They join AK's other binary searches in searchsorted.jl, which is now included before the sorts.
Internal names need no leading underscore: the module namespace hides them. The buffer requirement becomes `Scratch` (made by `scratch`), and allocating it `allocate_scratch`, as `allocate` is KernelAbstractions' export.
The capability predicates, the field checks and fills shared by the algorithm families, and backend resolution lose their underscores. `_backend_free` becomes `backendless`, and `_similar`/`_copy` become `similar_on`/`copy_on`, after the backend they allocate on when the array has none.
Each family resolves with `resolve_<family>`, selects with `select_<family>` and checks with `check_<family>`; sorting's `_check` becomes `check_sort` to match. The shared steps become `check_fields` (the fields set explicitly) and `fill_defaults` (`fill` would shadow Base's).
Every operation that plans scratch has `<op>_setup`, returning its plan with what the run needs, and `<op>_run`; a call of one operation inside another goes through `<op>_nested`. Scans use `scan_setup` and `scan_run!`, to pair with `mapreduce_setup`.
A maybe-empty partial result is an `Optional` (set or not, `isset`), the map and operator wrapped for a seed come from `seeded_ops`, and the marker of inputs that are already partial results is `PartialInputs`. The accumulator rule is `accumulator_type`, and `resolve_acctype` applies it or the caller's `acctype`. The `init` markers of the writing kernels become `FoldInto` and `Unfinished`. The internal arguments that carry the seed are called `seed`.
The launch wrapper behind `foreachindex` and every operation's loops becomes `foreachindex_on` (its kernel `foreachindex_kernel!`), and the offset-axes helpers `one_based`, `has_offset_parent`, `copyto_on!` and `fill_on!`. Of the block primitives, the sub-group width of the scans becomes `scan_subgroup_width`, and the vector load and fold `vload` and `fold_vector`.
`workspace_size(AK.map!, ...)` threw a `MethodError` naming the internal `plan`. A fallback method now says that the operation takes no workspace, so the docstring need not list those operations.
Both derived the settings of their element-wise loops from the resolved algorithm with identical methods.
`similar_on`, `copy_on` and the offset-axes helpers make and copy arrays; they have nothing to do with resolving algorithms.
`backend_vote` repeated the type lists of `backendless`. It now asks `backendless` first and leaves the forwarding of Base's wrappers to `get_backend`, which already votes like the wrapped array.
Inline the single use of `_describe`, let `show` of a buffer requirement reuse `size_pairs`, and keep the methods of `scratch_buffers` together.
The renames left the arguments on continuation lines aligned with the old, longer or shorter, names.
`foreachindex_on` was the only caller of its threaded and kernel halves, and `map!` of its two element getters. The kernel half's `block_size` check repeated `check_launch`.
Resolution checked an algorithm's explicit fields before filling it and all of its fields again after, so that filling never computed with bad values; but filling only copies values, and the one derived default (the scans' items per thread) guards itself. The drivers of the sorts and of the whole-array GPU reduction re-checked the same fields with `@argcheck`, though only resolved algorithms reach them. Reductions still check the fields that only apply to whole arrays before filling, since they must see which ones were set.
Five algorithms checked for a power-of-two block size up to 1024, and four for tiles within `typemax(Int32)` elements, each with its own copy. `check_block_size` and `check_tile` do it once.
`task_partition` copied `itask_partition` down to its validation. `task_count` gives the number of tasks of a partitioner without building one, which `findall` did only to size its counts, and which the threaded reduction and scan now ask before building one for their single-task paths. `findall`'s own single-task path did what the partitioner does for one task, so it goes.
The host backend's kernels and OpenCL's CPU devices both run on the CPU, and both defined the same reduction and scan tunings with their own copy of the reasoning. `cpu_reduce_tuning` and `cpu_scan_tuning` hold them, with the measurement, once.
Every other `<op>_setup` returns a `NamedTuple` with the plan, which the run takes its backend and algorithm from.
They named the `max_tasks` argument `num_tasks`, and the `itask_partition` example called `task_partition`. Also drop the comments that restate the constructor's code.
`sample_sort_barrier!` only isolated the ordering built from the keywords, which `sample_sortperm!` now receives as an argument. The index comparison stays a plain closure: `Base.Order.Perm`, which the sort along `dims` uses, breaks ties with a second comparison and made `sortperm` of 10^7 elements 15-40% slower on 4 threads.
The two differed only in the binary search they run for each query.
`sort!` with MergeSort maps `by` over the input into its key copy and sorts that, but `sortperm!` copied the raw values and evaluated `by` in every comparison, O(log² n) times per element. It now fills its key copy with `by.(v)` too, in a buffer of `by`'s result type, and sorts it under the ordering `by` wraps; the low-memory path, which has no key copy, is unchanged. On an RTX 5080, sortperm! of 10^7 Float32 with a `by` of 16 sin/cos steps goes from 194 to 6.9 ms (10^6: 14.4 to 0.38 ms; 1000x1000 along dims=1: 3.7 to 0.15 ms); with `by=abs` it stays at 6.9-7.0 ms.
The RadixSort docstring states its contract and its settings before its layout; the sort! docstring and the docs no longer suggest that Auto picks BitonicSort or RadixSort anywhere, since no backend's tuning sets their thresholds and GPUs sort with MergeSort. The internal radix driver gets a comment instead of a docstring, the zero-based searches say what they return, and the duplicated, restating and bold-label comments go.
The file tests searchsortedfirst! and searchsortedlast!, which live in src/searchsorted.jl.
The helpers every generic test file uses were an expression quoted in runtests.jl; as a file they read like the test files that use them.
The six backends' setups differed only in the package to load, the backend, its largest block and whether AK's kernels run, each repeating the same globals and their comments.
The reduction and scan tests each defined the same 2x2 matrix type with its product and random generator, and findall's tests their own offset vector next to the setup's OffsetAxes. The lazy array moves along, as the other array types the tests build sources from.
Resolution only looks at types and sizes, but its tests sat in five places under test/generic, which runs them again for every backend, each with its own fake backend and tuning switch. They now share one of each in test/algorithms.jl, which runs once. What does depend on the backend, that every tuning its device returns resolves, moves to generic/tunings.jl and now covers every family, not only sorting.
The resolution tests pinned the default tunings' values, so retuning a default meant changing tests. They now set the fake backend's tuning and check that Auto picks and fills from it, and that the host's algorithms take the host's tuning. The accumulator-type and seed rules were tested on the internal functions; the reduction contract tests now cover the cases they had that it did not, on real inputs.
They checked that keywords and arguments of earlier versions are now MethodErrors, which any unknown keyword is; the operations that had no unknown-keyword test (foreachindex, findall) get one.
On GPUs, Auto() picks the merge sort for every permutation and, with the default tuning, for every sort, and on the host the sample sort, so the fuzz loops of merge_sort and merge_sortperm on GPUs and of sample_sort and sample_sortperm on the host repeated those of sort and sortperm. Each algorithm now has one fuzz loop, which compares with Base's result rather than checking that it is sorted; sortperm's also checks for a permutation. sortperm_extended's element types, lengths and inputs join sortperm's tests, as exact comparisons with Base's stable permutation, which replace its stability and validity checks. sort_by_transform becomes one loop over the `by` keywords, also for the sample sort. The rejections of every algorithm and the radix and bitonic sorts' own are gathered in one testset; the radix sort's stability test of plain integers, where it cannot be observed, goes. Testsets that only run AK's kernels are skipped from outside.
Its loop over random arrays along dims called sum.
Which kernel a reduction along dims launches was asserted next to the tests of each shape's results, four times with its own helper, the vector-load width with hand-built strides, and the one-pass columns kernel was reached by calling its 22-argument launcher. One testset now checks that the inputs of those tests reach each shape and load width, computing both from the arrays as the reduction does; the one-pass columns kernel is reached by an input of the strided-slices tests, which also fold into their destination as the launcher call did.
Each configuration has one whole-array scan algorithm since the decoupled look-back scan went, so the testsets that looped over the algorithms, and scan_alg's argument choosing one, no longer select anything.
The workspace tests built Workspace objects field by field in three places to alias or swap a buffer or nested algorithm, and compared the planned sizes through an internal conversion. One helper now forges the workspaces, and the sizes are compared with those of the allocated buffers.
The testsets were named in two styles, `reduce_1d` or `merge_sort` next to `sort: offset axes`; they now all take the latter.
Comments that continued the previous one with "..." or a parenthesis become sentences of their own, those that named the kernel paths an input takes describe the input, and the convention for the algorithm of tests that choose none is stated once, in test/setup.jl, with a function for it. Testsets that run only AK's kernels are skipped from outside, and top-level testsets are two blank lines apart.
Each family's test file defined its default algorithm and a function building one with settings, and the sorting tests a constant for one set of settings. They follow one convention and the split test files of a family need the same ones, so the setup defines them all, sorting with a builder like the others'. The point type the reductions map moves along.
sort.jl was the longest test file on every backend, and ParallelTestRunner runs files in parallel; the permutation and key/value tests and those along `dims` move to files of their own.
reduce.jl was the longest test file on CUDA and OpenCL; the reductions along `dims`, their launch shapes and the non-commutative reductions, whose cases are mostly along `dims`, move to reduce_dims.jl.
Nothing is exported, so that the functions can share Base's names, but until now nothing marked which names are the API either. Declare the functions, algorithm types and helpers the manual documents `public`, plus the re-exported `get_backend` and `synchronize`. The tuning hooks and structs stay internal, as the tuning page says. The keyword does not parse on Julia 1.10, hence the `eval`.
The canonical URL pointed at KernelAbstractions' docs, and `warnonly=true` turned every missing docstring and broken `@ref` into a warning nobody reads. Check all docstrings, so that internal helpers documented with a docstring must also appear in the manual, and fail the build otherwise.
The test file was renamed to searchsorted.jl already; the page follows.
Only foreachindex, foraxes, map!, reverse! and the binary searches take a block size; the other operations take their settings through `alg`. The Performance Tips page repeated that claim and the workspace page, so drop it. Say on the overview what the package is: CUB-style primitives with their own contract, which GPUArrays.jl and power users build on. List OpenCL among the backends, document the re-exported `get_backend` and `synchronize`, keep the acknowledgements on the References page only, document `SortAlgorithm` with the other sorting types, and write the "Other names" labels as sentences.
Say what the package is for (CUB-style primitives that GPUArrays.jl and power users build on), add OpenCL to the backends and the CI table, point the Buildkite badges at the pipeline's step labels (Julia 1.10 and 1.13; the old names matched no step), drop the allocating binary searches that do not exist from the function table, and fix the package name in the threaded `Pkg.test` command. The Status and Roadmap sections showed keywords and scratch arrays that no longer exist.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A cleanup of the whole package, without changing what it computes. Much of AK 0.5 (the host API rework and the kernel PRs since) was written with LLM agents, and it shows: about 300 internal names with leading underscores, comments that narrate earlier versions of the code ("the original, non-striding behavior", "historical settings"), the same checks run two or three times on the way to a kernel, files whose names don't match their contents, and tests that pin internals. This PR makes the code read like one maintainer wrote it, so that the structural follow-ups (merging the duplicated kernels) are reviewable.
What changes, by kind (one commit per logical step, renames and moves never mixed with edits):
prototype/(21 of its 31 scripts called APIs that no longer exist), the Pluto benchmark notebook,TypeWrapand theu8...i64literals (unused, not what AK is for), AK's ownispow2(wrong fortypemin), dead code, andDecoupledLookback. The latter was not a single-pass scan: its lookback ran as a second kernel after all blocks had finished, so it moved more data thanScanPrefixes, its floating-point results varied between runs,Autonever chose it, and its docs claimed a forward-progress requirement it didn't have. No package uses it. With it go the device fences in the extensions, the AMDGPU extension (which held nothing else) and the UnsafeAtomics dependency.DocHelpers, which read README.md at load time for the docs, moves to docs/, dropping the Markdown dependency.arithmetics.jl's shorthands join their families; every family keeps its algorithm structs next to its tuning.plan/Workspace/Scratch,resolve_X/check_X/fill_defaults,X_setup/X_run,*_kernel!), and CUB's terms in kernels (tiles, items per thread, digits, aggregates). The internal partial-result seed isseed, not the user-facingneutral.check_block_size/check_tile;task_partitionisitask_partitionwithout the index; one destination check for scans andreverse!; dispatch instead ofalg isachains; the sorting drivers receive the resolved ordering only.mapreducedim!; README and docs pages drop stale claims (block sizes on every function, removed keywords in the roadmap,Autochoosing radix/bitonic sorts, which no backend tunes yet) and say what AK is: CUB-style primitives with their own contract, which GPUArrays and power users build on. The documented API is declaredpublic(behind a version check for 1.10), and the docs build fails on missing docstrings and broken references.test/setup.jl; long files are split so they run in parallel. The suites take 30-50% less time (CUDA 3m56 → 2m24) with the same coverage.Behaviour that does change:
sortperm!withMergeSortevaluatesbyonce per element instead of in every comparison (10^7Float32s with an expensiveby: 194 ms → 6.9 ms on an RTX 5080; cheapbys are unchanged). Thelowmemvariant still evaluates it per comparison.reverse!(dst, src)rejects overlapping arrays (reverse!(x, x)used to corruptx) and throws aDimensionMismatchfor mismatched sizes, asaccumulate!does.workspace_sizeandworkspaceof operations that take no workspace throw anArgumentErrorinstead of aMethodErrornaming an internal function.Tested with the full suites on CPU threads, PoCL (also running AK's kernels on the host backend), CUDA (RTX 5080, RTX 6000 Ada), oneAPI (Arc A770, oneAPI.jl main) and Metal (M1), on Julia 1.12 and 1.10. On AMDGPU (an integrated gfx1036, AMDGPU.jl main) the only failures also occur without this PR:
@localmemof isbits unions (JuliaGPU/AMDGPU.jl#1184) and that iGPU's sporadically lost kernel results. The docs build cleanly with the stricter checks.