Repository navigation
Conversation
Merged
1 of 3 tasks
This was referenced Sep 24, 2026
Closed
Closed
maleadt
changed the title
maleadt
force-pushed
the
tb/api
branch
2 times, most recently
from
September 29, 2026 21:09
e7b4e09 to
5d78328
Compare
The foundation of the host-API restructure, and the first family to use it.
Algorithms are values that carry their settings: `alg=Auto()` lets AK choose,
an explicit algorithm is used as given or rejected with an `ArgumentError`.
`Auto(; stable=true)` carries requirements; with `stable=true` it chooses an
unstable algorithm only where equal elements are bitwise identical.
Choices can differ per device. `sort_tuning(backend, T)` returns a
`SortTuning` of thresholds and settings, and `_resolve_sort` selects,
fills unset fields and checks the result against the backend's capabilities
(`_runs_kernels`, `_runs_threads`). The default tuning selects the same
algorithms with the same settings as before; only merge sort's `by` keys are
now computed by AK's `map!` (see below).
`backend` becomes a keyword, derived from all array arguments, which must
agree; ranges and other non-arrays do not count, and without any array the
host backend is used.
The threaded CPU sort becomes `CPUThreads.SampleSort`, chosen by `Auto` on
the host backend. On KernelAbstractions 0.10 the host backend also runs AK's
kernels, so explicit kernel algorithms work on host arrays there; the
`--cpu-ka` test configuration tests that, in a new CI job.
Sorting entry points:
- `sort!`, `sort`, `sortperm!`, `sortperm` take `backend` and `alg` keywords;
`prefer_threads`, `max_tasks`, `min_elems` and the loose `block_size` are
gone (they are algorithm fields now), as is `alg=nothing`.
- `sort_by_key!` replaces `merge_sort_by_key!`/`merge_sort_by_key`, with
`alg` and a threaded host implementation.
- `merge_sort!`, `merge_sort`, `merge_sortperm!`, `merge_sortperm`,
`merge_sortperm_lowmem!`, `merge_sortperm_lowmem`, `bitonic_sort!`,
`sample_sort!`, `sample_sortperm!`, `bitonic_defaults` and `SampleSort`
are removed; the algorithms are reached through `alg`.
- `MergeSort` gains a `block_size` field.
- `RadixSort` settings are bounded (`block_size <= 1024`,
`items_per_thread <= 64`), so its local-memory checks cannot overflow.
- Merge sort computes `by` keys with `map!` into an array of the keys' type
instead of broadcasting, which made a `BitArray` of `Bool` keys that
kernels cannot take.
- Nested operations (radix sort's key-range reduction, the index and key
initialisation of merge sort) receive the resolved backend.
- The threaded sort and sort permutation of an N-d array without `dims` no
longer fail for small inputs (`Base.sort!` needs `dims` for matrices).
- Merge sort is tested on reshaped views and on `Union{Missing, Int32}`
arrays, which `main` failed to compile on CUDA. The views need
KernelAbstractions 0.9.43, which rebuilds them in `@Const`
(JuliaGPU/KernelAbstractions.jl#794), so the compat bound rises; the bits
unions need CUDA.jl's cached loads of them (JuliaGPU/CUDA.jl#3297), so the
CUDACore bound rises to 6.4.1. On
KernelAbstractions 0.10's POCL backend (`--cpu-ka`), they need
SPIRVIntrinsics 1.1.4, which allocates local memory of bits unions
(JuliaGPU/OpenCL.jl#516), and SPIRV_LLVM_Backend_jll 23.1.1+2, whose fix of
pointers extracted from aggregates sorting by key needs
(llvm/llvm-project#227599).
`foreachindex`, `foraxes`, `map!`, `map`, `reverse!`, `reverse`,
`searchsortedfirst!` and `searchsortedlast!` take `backend` as a keyword,
derived from their arrays (the destination first); a loop over a range with
no backend runs on the host. On the host backend they always run on Julia
threads, and `prefer_threads` is gone. Their launch settings, `block_size`,
`max_tasks` and `min_elems`, are explicit keywords, checked on every backend.
- `map!` and the batched searches return their destination (the searches
returned `nothing`).
- The batched searches take `rev::Union{Nothing,Bool}=nothing` and `order`,
like every other ordering-taking operation.
- The allocating `searchsortedfirst(v, xs)` and `searchsortedlast(v, xs)` are
removed: Base has these names with different semantics (a vector `x` is
one value), so `AK.searchsortedfirst` is now Base's.
Internal callers pass their backend to the internal `_foreachindex`.
`reduce` and `mapreduce` take `backend` and `alg` keywords, like sorting. The GPU tree reduction becomes `BlockReduce(block_size, items_per_thread, switch_below)` and the threaded host reduction `CPUThreads.Partitioned( max_tasks, min_elems)`, which scans, `findall` and the predicates will share. `Auto` picks `Partitioned` on the host backend and `BlockReduce` elsewhere. `reduce_tuning(backend, T)` returns a `ReduceTuning` whose values fill unset fields; its defaults are the settings used so far, so every call launches the same kernels as before. The number of blocks a reduction along `dims` aims for (`TARGET_BLOCKS`) is a tuning value too, `target_blocks`. The reduction contract is unchanged here: `init` is still required, and results take its type. The next commit changes that. - The positional `backend` and the loose `block_size`, `items_per_thread`, `switch_below`, `max_tasks`, `min_elems` and `prefer_threads` keywords are gone from `reduce`, `mapreduce`, `sum`, `prod`, `maximum`, `minimum` and `count`; the convenience reductions forward `alg`. - `BlockReduce` with an explicit `items_per_thread` or `switch_below` along `dims` is an `ArgumentError`; the settings were silently ignored there. - Resolution rejects whole-array `BlockReduce` tiles (`block_size * items_per_thread`) of one element, which never finish, and of more than `typemax(Int32)` elements. A tuning's `target_blocks` must be positive. - Base's views, reshapes and permutations of ranges and other backend-free values do not determine the backend either. Otherwise `get_backend` would throw for a `Broadcasted` over reshaped ranges, which the host reduction has accepted so far. - Radix sort's key-range reduction and the predicates' reduction path pass their backend and algorithm through the new keywords. - The domain and capability checks shared by all families move next to the algorithm types.
Reductions get a contract of their own, modelled on CUB's rather than on
Base's, so that a front-end such as GPUArrays can build Base's API on it
and decide Base's rules (empty results, result types) itself. `op` must be
associative and commutative, as for every GPU reduction.
`mapreducedim!(f, op, R, A; init, neutral, overwrite, acctype, alg)`
reduces into an existing array: the dimensions where `R` has size 1 are
reduced, and each output whose slice is not empty becomes
`op(init, partial)` with `init` (applied once), `partial` with
`overwrite=true`, and `op(R[i], partial)` otherwise, folding as
`Base.mapreducedim!` does. An output whose slice is empty is set to `init`,
or not written. The allocating reductions along `dims` are built on it.
`init` is optional, and `init = nothing` is an initial value; the
keyword's default is an internal sentinel. Without `init`, an empty
reduction is an `ArgumentError`, of a whole array and along `dims` of any
output whose slice is empty. `sum` and `prod` give zero and one of the
accumulator type instead (not as an `init`, so that a sum of `-0.0`s stays
`-0.0`), and `count` keeps its `init=0`.
Partial results have one accumulator type and are converted to the
destination's element type only when stored: `acctype` when it is given,
else the type the fold of `op` settles on from `init`'s type (or
`eltype(R)`) and the mapped elements. Every element is also a one-element
partial result (`Base.reduce_first(op, x)`), so that type is joined in too,
and `sum(Int8[...])` accumulates in `Int`. Results have that type, with
`op(init, partial)` converted to it: scalar results, arrays along `dims`
also with an `init` (not `typeof(init)`), and a single element
(`AK.reduce((a, b) -> a + b, [true]) === 1`). An `acctype` is rejected only
when a partial result cannot convert to it at all, or `op` cannot combine
it (whatever the input); otherwise an `op` that inference shows always
throws is an error only where something is combined.
`neutral` only seeds partial results. When it is not given, AK uses
GPUArraysCore's `neutral_element(op, T)` where one is defined; for any other
operator, each partial result starts from its first element
(`Base.mapreduce_first`), so anonymous operators and tuple accumulators need
no neutral element. The same kernels serve both cases: without a neutral
element, they reduce `(value, valid)` lanes through wrappers of `f` and `op`
that the kernels build on the device.
A new documentation page lists where these rules differ from Base's.
- `sum` and `prod` use Base's `add_sum` and `mul_prod`, so small integers
are summed as `Int`; `count` requires `f` to return a `Bool`.
- `temp` no longer doubles as the destination of reductions along `dims`;
it is scratch for whole-array reductions only, with the accumulator
element type.
- `BlockReduce` requires a bits-type accumulator and says so, instead of
failing in the compiler.
- Every kernel shape is tested with a reshaped view and a
`Union{Missing, Bool}` source, which failed to compile on CUDA before the
fixes in KernelAbstractions 0.9.43 and CUDA.jl (#134).
- `neutral_element` comes from GPUArraysCore 0.2.1; AK's copy is gone.
`accumulate!`, `accumulate`, `cumsum` and `cumprod` take `backend` and `alg` keywords, like sorting and reductions. The whole-array GPU scans `ScanPrefixes` and `DecoupledLookback` carry `block_size` and `items_per_thread`; the per-slice kernels used along `dims` become an algorithm of their own, `SliceScan(block_size)`; the threaded host scan is `CPUThreads.Partitioned`. `Auto` picks `Partitioned` on the host backend, `SliceScan` along `dims`, and `ScanPrefixes` for whole arrays, or `DecoupledLookback` where the backend supports it and the tuning prefers it (no default tuning does, so every call launches the same kernels as before). `scan_tuning(backend, T)` returns a `ScanTuning`. Its `items_per_thread` default is derived, as before, from the effective block size and a local-memory budget, so an explicit `block_size=1024` gets a matching default. oneAPI gets the default tuning too: oneAPI.jl passes 64-thread blocks on every call, to avoid wrong results with larger blocks, which no longer reproduce (thousands of random scans with 128, 256 and 512 threads are correct on an Iris Xe). `DecoupledLookback` needs a device-scope fence, atomics and forward progress between blocks. That is a capability, `_supports_lookback`, which the CUDA and AMDGPU extensions declare; elsewhere an explicit `DecoupledLookback` is an `ArgumentError`, as is `ScanPrefixes` or `DecoupledLookback` along `dims` (they were silently replaced by the per-slice kernels) and `SliceScan` without `dims`. - The positional `backend` and the loose `block_size`, `items_per_thread`, `max_tasks`, `min_elems` and `prefer_threads` keywords are gone; the `AccumulateAlgorithm` supertype is now `ScanAlgorithm`. - `temp` and `temp_flags` with `dims` are an `ArgumentError`; they were ignored. - `accumulate!(op, dst, src)` with `dst === src` no longer copies. - The scan contract is unchanged here: `init` is still required. The next commit changes that.
Scans get a contract of their own, like the reductions of commit 4: `op` must be associative (not commutative), elements keep their order, and a front-end decides Base's rules itself. The running value has one type and is converted to the destination's element type only when stored: `acctype` when it is given, else the accumulator type of commit 4's reductions, starting from the destination's element type joined with `init`'s type, as `mapreducedim!` starts from `eltype(R)`. So with the usual operators and no `acctype`, a scan does not run in a type narrower than its destination: `accumulate!(*, zeros(Int, 4), fill(0x10, 4))` gives `[16, 256, 4096, 65536]`, and `Float32`s scanned into a `Float64` array are summed in `Float64`. Elements enter as one-element reductions (`Base.reduce_first`): `AK.accumulate(*, ['a', 'b'])` gives strings. Only where the running type differs from the destination's element type does the scan run in a scratch array of that type. `accumulate` allocates the fold type from `init`'s type and the elements (`accumulate(+, Int8[1, 2]; init=0)` is an `Int` vector), or `acctype`. `init` is optional. Without it, an inclusive scan starts from the first element; with it, `init` is applied as `op(init, x)` to the first element of each slice (or seeds the kernels directly when it has the running type), so `init = nothing` is an initial value. A `dims` beyond the array's dimensions makes every slice one element long: inclusive scans apply `init` to every element, and exclusive ones fill every element with their seed. A scan that never calls `op` only copies. An exclusive scan starts from `init`, or from the neutral element of `op` without it; when none is known that is an `ArgumentError`. `neutral` defaults to GPUArraysCore's `neutral_element` where one is defined (`-0.0` for floating-point sums, which unlike `0.0` is an exact identity). For other operators every scan kernel, and the threaded host scan, keeps its partial results as `(value, valid)` lanes, like the reductions of commit 4, so inclusive scans need no neutral element: anonymous operators, and non-commutative ones such as matrix products, just work. Elements are lifted into lanes as they are loaded and lowered as they are stored; an empty lane is never stored. `cumsum` and `cumprod` use Base's `add_sum` and `mul_prod`, so `AK.cumsum(Int8[...])` is an `Int` vector. The documentation's differences from Base gain the scans' rows. - `accumulate!(op, dst, src)` requires `dst` to be `src` or not to overlap it, and to have its axes. - `temp` requires a known neutral element, like the reductions' `temp`, and has the running type. - Scans along `dims` step through the array with the strides of its linear indices rather than `strides(v)`, which differ for wrappers such as a `PermutedDimsArray` or a strided view: those were scanned wrongly before. - `accumulate`, `cumsum` and `cumprod` of an input without a backend (a range) allocate their result on the `backend` given.
`findall`, `any` and `all` take `backend` and `alg` keywords, like the other
families, which completes the calling convention: no operation takes a
positional backend or `prefer_threads` any more.
`findall` has `ScanScatter(block_size, items_per_thread)` on GPUs and
`CPUThreads.Partitioned` on the host, filled from `findall_tuning`. `any`
and `all` have `ConcurrentWrite(block_size)`, which stores one flag from
many threads, and `ViaReduce(reduce)`, a reduction with `|` or `&` through a
`ReduceAlgorithm` (it was `MapReduce(temp, switch_below)`, whose fields were
the reduction's), filled from `predicate_tuning`. Every predicate algorithm
requires the predicate to return a `Bool`; the reduction used to accept
integers. Where inference shows the predicate returns something else, the
call throws an `ArgumentError` before launching (except for an empty array,
where the predicate is never called), rather than failing in the kernel,
which GPU backends report without Base's `TypeError`.
`findall` checks its predicate, or its non-`Bool`
values, the same way. `Auto` picks `ConcurrentWrite`, or `ViaReduce()`
where the tuning prefers it.
`findall` selects items rather than computing indices by Base's rules: the
new `items` keyword (default `keys(A)`) is any array of `A`'s length, and
the `k`-th element of `A` selects its `k`-th element. `keys(A)` gives
`Int`s for vectors and `CartesianIndex`es otherwise, now also for the
predicate form of a 0-dimensional array; `LinearIndices(A)` gives linear
indices; an array of values selects them, so `findall(mask; items=A)` is
`A[mask]` in one pass. Positions are ordinal, so offset axes work.
oneAPI's extension redefined `AK.any` and `AK.all` to default to the
reduction, because some Intel GPUs were reported to hang on concurrent
writes to one location. It now sets that default through the tuning
(`prefer_concurrent_write=false`), so an explicit `ConcurrentWrite` runs
there as on every GPU; 900 such calls over up to 10^8 elements, all of
which match, ran correctly on an Iris Xe.
- The positional `backend` and the loose `max_tasks`, `min_elems`,
`block_size` and `prefer_threads` keywords are gone, and so are
`use_gpu_algorithm` and the test harness's `prefer_threads` global.
- `FindallAlgorithm` and `PredicateAlgorithm` (was `PredicatesAlgorithm`) are
`Algorithm`s, documented with the other families.
- `findall`'s mask lives on the resolved backend, so a range with an
explicit GPU backend works.
- `any` of a `Union{Missing, Bool}` array, and `any` or `findall` of a
reshaped view, are tested; they failed to compile on CUDA before the
fixes in KernelAbstractions 0.9.43 and CUDA.jl.
Every operation that needs scratch memory takes a `workspace` keyword
instead of its `temp`, `temp_flags`, `temp_ix`, `temp_keys`, `temp_values`
and `temp_bools` keywords, which are gone. `AK.workspace(op, args...;
kwargs...)` allocates the scratch of the call `op(args...; kwargs...)`, and
`AK.workspace_size` returns it as a `NamedTuple` of `(eltype, dims)` pairs
without allocating:
ws = AK.workspace(AK.sort!, v)
AK.sort!(v; workspace=ws) # allocates no scratch memory
Each operation now plans its scratch in one function, `_plan`, which
resolves the algorithm and lists the buffers the implementation uses,
including those of the operations it calls (radix sort's key-range
reduction and histogram scan, findall's scan of its counts, the predicates'
reduction). The implementation takes its buffers from the plan, allocated
or from the workspace, and allocates no scratch of its own, so the query
and the call cannot disagree. Merge and radix sorts, sample sort, whole
and dimensional reductions, scans (including the scratch array of a scan
whose destination has another element type), findall and the predicates
have plans; `workspace` works for `sort`, `sortperm`, `accumulate`,
`cumsum`, `sum` and the other allocating and convenience forms too. A plan
takes the operation's keywords, including those that change the buffers:
`acctype` sets the type of a reduction's partial results and of a scan's
scratch array, and `findall`'s `items` is checked against the input. The
empty `sum` and `prod`, which do not reduce, check their workspace too.
A `Workspace` records the backend, the device, the resolved algorithm and
those of the operations it calls, and the buffer sizes. A call checks all of
them, and that no buffer aliases its arrays (including the arrays of a fused
or `Broadcasted` source, and the input of an allocating `sort`), and throws
an `ArgumentError` otherwise: `Auto()` may resolve differently for another
length, so a workspace made for one call is only accepted by calls that plan
the same scratch.
The workspace covers device memory, with two exceptions. On the host, the
threaded algorithms keep their small per-task bookkeeping, and
`CPUThreads.SampleSort` leaves `Base.sort!` to allocate the scratch of its
per-task sorts and of the slices of a sort along `dims` (its serial path
takes the workspace's buffer). Before Julia 1.12, a reduction of several
arrays or of a `Broadcasted` object still materializes it; the plan and the
workspace checks see the original arrays, and materializing happens only
when the reduction runs.
- The reduction's launch-shape decision along `dims` moves into a function
that the plan and the kernel dispatch share.
- `DecoupledLookback`'s flags and every other buffer come from the plan;
none needs initialising between calls.
- Buffer requirements carry their element type in the type domain, and the
setup of reductions and scans carries types as `Val`s, so results still
infer.
- A `Broadcasted` reduction source below `switch_below` is evaluated on the
host from host copies of its arrays, instead of being materialized on the
device.
- Docs: a "Scratch memory" page, and the sorting and performance pages use
workspaces.
Member
|
I very much like this direction - however it seems that it's not a full generalisation of the algorithms and their specific requirements (temporaries, block / thread sizes, etc.); in particular, why hardcode the stability of sorting into the |
Member
Author
|
Full generalization wasn't really the goal, just a step towards a more principled API instead of loose keywords and ad-hoc arguments. It's still 0.5, but good enough to start integrating into GPUArrays. Some of what you mention is already there: there are families ( |
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
2 participants
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.
This PR reworks how AcceleratedKernels' operations are called, how they choose an algorithm, and what they promise about their results. It breaks almost every call site, so it is meant for AK 0.5.
Fixes #134.
Why
AK's operations grew one at a time, and it shows in their interfaces. Some algorithms carry their settings, while others take loose keywords on the operation (
block_size,items_per_thread,switch_below,max_tasks, ...), which some algorithms silently ignore.tempis scratch memory in one function and the destination in another. The backend is an optional trailing positional argument, which makesmapreduce(f, op, A, B, ...)ambiguous. Reductions requireinit, need a neutral element for every operator, and return AK'sneutralfor empty inputs.That makes AK hard to build on. GPUArrays wants to implement
sort!,mapreduce,accumulateandfindallfor every GPU array on top of AK (JuliaGPU/GPUArrays.jl#790), so that CUDA.jl, AMDGPU.jl, Metal.jl, oneAPI.jl and OpenCL.jl can drop their own copies of these kernels. For that, AK needs one calling convention and primitives whose behaviour is written down.So this PR gives every algorithmic operation (sorting, reductions, scans,
findall,any/all) the same shape:The launch wrappers (
foreachindex,foraxes,map!,reverse!and the batched searches) take the samebackendkeyword, plus their launch settings (block_size,max_tasks,min_elems) as keywords.It also draws a line between AK and GPUArrays. AK provides primitives with contracts of their own, modelled on CUB's rather than on Base's. GPUArrays builds Base's API on top of them, and decides everything that exists only because Base says so: result types, what an empty reduction returns, index types.
What it looks like
Most calls just drop the positional backend:
To control the algorithm, pass one. Algorithms are values that carry their settings, and fields left out are filled in for the device:
An explicit algorithm is used as given or rejected. It is never silently replaced, and bad settings fail on the host before any data is touched:
Reductions no longer need
initor a neutral element, and closures and tuple accumulators work:A new
mapreducedim!reduces into an existing array, which is what GPUArrays needs forsum!,maximum!and friends:By default, scans run in the destination's type (or wider), not in the source's:
findallselects items: by default the keys, but an array of values works too, soA[mask]takes a single pass:The backend comes from the arrays. Ranges and other lazy collections have no backend, so they run on the host unless you say otherwise:
Scratch memory can be allocated once and reused. This replaces every
temp*keyword:Design
Each family has a few algorithm structs:
sort!,sortperm!,sort_by_key!MergeSort,RadixSort,BitonicSortCPUThreads.SampleSortreduce,mapreduce,mapreducedim!BlockReduceCPUThreads.Partitionedaccumulate!ScanPrefixes,DecoupledLookback,SliceScanCPUThreads.PartitionedfindallScanScatterCPUThreads.Partitionedany,allConcurrentWrite,ViaReduceCPUThreads.PartitionedAuto(), the default, picks one for the device, andAutonever reads array contents. The values it fills in come from one internal tuning per family (sort_tuning(backend, T)and friends), which a backend extension can specialise per device. The default tunings reproduce the settings AK uses today, so calls without explicit settings launch the same kernels as before.Tunings only affect speed. Correctness facts about a backend are separate capabilities: whether it runs AK's kernels at all, or has the forward progress that
DecoupledLookbackneeds. These are declared once per backend and checked forAutoand explicit algorithms alike, so no tuning can enable something unsafe.The host is an algorithm rather than a flag. On host arrays,
Autopicks theCPUThreadsalgorithm (Julia threads), as before.prefer_threadsis gone, and so are themax_tasks/min_elemskeywords on the algorithmic operations; they are fields of theCPUThreadsalgorithms now. With KernelAbstractions 0.10, whose host backend runs kernels through PoCL, an explicit GPU algorithm also runs AK's kernels on host arrays (AK.sort!(v; alg=AK.MergeSort())). A new CI job tests this.A workspace records the backend, device, resolved algorithm and buffer sizes. A call that would need anything else throws an
ArgumentErrorrather than using the wrong buffers. That includesAutopicking another algorithm for another length, and a workspace whose buffers alias the call's arrays. On the host, the threaded algorithms still allocate small per-task bookkeeping.The contracts
AK keeps Base's function names, since the functions take Base's arguments, but it does not promise Base-identical results. The docstrings have the details, and a new "Differences from Base" documentation page lists the main semantic differences. In short:
op; scans need an associative one, and keep element order, so matrix products work. A reduction's combination order depends on the algorithm and its settings, but is fixed for a given setting, shape and device.initis optional. When given, it is applied exactly once, asop(init, partial), andinit = nothingis an initial value like any other.initis anArgumentError, also alongdimsfor every output whose slice is empty.sumandprodgive zero and one, andcounthasinit=0. An empty scan writes nothing.GPUArraysCore.neutral_elementwhere one is known. Otherwise each partial result starts from its first element, so closures and tuple accumulators such asfindmin's need no neutral element.acctypewhen given; otherwise it is the type the fold ofopsettles on, starting frominit's type (whole arrays),eltype(R)(mapreducedim!) oreltype(dst)joined withinit's type (scans). Anacctypenarrower than that is the caller's choice.findallselectsitems. Thek-th element ofAselects thek-th element ofitems, which defaults tokeys(A). Positions are ordinal, so offset axes work.Bool. Where inference shows one cannot,any,allandfindallfail on the host with anArgumentErrorrather than in a kernel.Auto(stable=true)only picks an unstable algorithm where equal elements are indistinguishable (integers,BoolandCharunder the default ordering, but not floats, whose NaNs can differ).BlockReduce(items_per_thread=...)alongdims, orScanPrefixesalongdims.Back-end fixes this relies on
Testing this on every back-end turned up bugs outside AK. They are fixed upstream rather than worked around here:
@Constloads of bits-union element types such asUnion{Missing, Bool}(Constloads (@Constin KernelAbstractions) fail to compile for bits-union element types CUDA.jl#3290, fixed by Support bits-union element types inConstloads CUDA.jl#3297 in CUDACore 6.4.1, which AK now requires). AK's reduction, predicate,findalland merge sort kernels mark their inputs@Const, soAK.any(ismissing, CuArray(Union{Missing, Bool}[true, missing]))failed onmain. AK now requires that release and tests every affected kernel.@Constcould not rebuild a reshaped view on the device (@Conston a reshaped view fails to compile on CUDA KernelAbstractions.jl#792, fixed in KernelAbstractions 0.9.43), so e.g.AK.sum(vec(view(A, 1:40, 1:30)))failed on CUDA. AK now requires KernelAbstractions 0.9.43.KernelAbstractions.zerosof a bits-union element type throws KernelAbstractions.jl#791, fixed onmain, which the KA 0.10 CI job uses).@localmem Union{Missing, Int32} (n,)did not compile on that backend (Local memory of a bits-union element type does not compile OpenCL.jl#515, fixed in SPIRVIntrinsics 1.1.4).sort_by_key!of bits unions hit a bug in LLVM's SPIR-V back-end, which emitted invalid SPIR-V (OpCompositeExtractwith a mismatched pointer type) for a pointer extracted from an aggregate (LLVM's SPIR-V back-end emits invalid SPIR-V for pointers extracted from aggregates (bits-union values) OpenCL.jl#517). The fix, [SPIR-V] Fix pointee types of pointers extracted from or inserted into aggregates llvm/llvm-project#227599, is carried in SPIRV_LLVM_Backend_jll 23.1.1+2.(flag, (a, b))struct (Intel's OpenCL driver (NEO) stores a wrongBoolfield in a struct followed by a tuple OpenCL.jl#502). AK's reduction lanes put the value first, which avoids it; GPUCompiler 2.9 works around the bug for every field order, which OpenCL.jl 0.10.12 enables.oneAPI needed two decisions. oneAPI.jl calls AK's scans with 64-thread blocks, because larger blocks once gave wrong results. That no longer reproduces (2,400 random scans with 64 to 512 threads on an Iris Xe were all correct), so AK's oneAPI extension does not carry it over and uses the default scan tuning. On
main, the extension also makesany/alldefault to the reduction, because some Intel GPUs were reported to hang when many threads write one location. An Iris Xe runsConcurrentWritefine (900 calls over up to 10^8 elements), but we don't know which devices were affected, soAutokeeps pickingViaReduceon oneAPI, through its predicate tuning. An explicitConcurrentWritestill runs there, as onmain.Testing
Every commit passes the test suite on its own, with the released back-ends it requires: CUDACore 6.4.1, GPUArrays 11.5.15, KernelAbstractions 0.9.43, OpenCL.jl 0.10.12 with
pocl_jll7.2.0+3 (the local machine has an RTX 5080). The host-backend columns use KernelAbstractionsmainwith SPIRVIntrinsics 1.1.4 and SPIRV_LLVM_Backend_jll 23.1.1+2:310f1055ffbdb5e6de24db8082d5489bbfc473d42c135188c0bdbd3fAt the head, the suite also passes with the released back-ends on Metal (Apple M1, Metal.jl 1.11.1; Julia 1.13: 51,548, 1.10: 51,546), on an Intel Iris Xe with oneAPI.jl 2.9.2 (51,543 and 51,541) and with OpenCL on Intel's driver (51,541 and 51,539). The CUDA CI job, run locally with
Pkg.teston Julia 1.13.1, passes (56,932). AMDGPU is left to CI.The bits-union tests run where the arrays can hold bits unions: CUDA, AMDGPU, host arrays and KernelAbstractions 0.10's host backend. Metal.jl's, oneAPI.jl's and OpenCL.jl's arrays cannot hold them yet.
Before merging
Follow-ups, all of which fit behind this API: measured sorting thresholds and per-device tunings; faster small and whole-array reductions (#135); an order-preserving reduction and a sequential scan for non-commutative or non-associative operators; device-side work such as sub-group collectives, onesweep and segmented radix sort.
Reviewing
Each of the eight commits passes the tests on its own. Commit 1 introduces algorithm selection and moves sorting onto it; it carries most of the design. Commits 4 and 6 define the reduction and scan contracts, and are easiest to read through their docstrings (
mapreducedim!,accumulate!) anddocs/src/api/differences.md. Commits 2, 3, 5 and 7 convert the remaining operations (7 also givesfindallitsitems), and 8 replaces the scratch keywords with workspaces.Migration guide: every changed or removed entry point
backend(every operation)backend=keyword; usually omit itprefer_threads(every operation)alg=AK.CPUThreads.…(or a kernel algorithm on KA 0.10); the launch wrappers always use Julia threads on the hostmax_tasks,min_elemson algorithmic operationsalg=AK.CPUThreads.SampleSort(; max_tasks, min_elems)/CPUThreads.Partitioned(...)block_size,items_per_thread,switch_belowon algorithmic operationsalg=AK.BlockReduce(block_size=512)temp,temp_flags,temp_ix,temp_v,temp_keys,temp_values,temp_bools,MapReduce(temp=…)workspace=AK.workspace(op, args...; kwargs...)foreachindex(f, itr, backend; …),foraxes(f, itr, dims, backend; …)foreachindex(f, itr; backend, block_size, max_tasks, min_elems),foraxes(f, itr, dims; …)map!(f, dst, src, backend; …),map(f, src, backend; …)map!(f, dst, src; backend, …)(returnsdst),map(f, src; …)reverse!(v, backend; …),reverse(v, backend; …)reverse!(v; backend, dims, …),reverse!(dst, src; …),reverse(v; …)searchsortedfirst!(ix, v, x, backend; rev::Bool)(returnednothing)searchsortedfirst!(ix, v, xs; backend, lt, by, rev, order)(returnsix); likewisesearchsortedlast!searchsortedfirst(v, x),searchsortedlast(v, x)(batched, allocating)searchsortedfirst!(similar(xs, Int), v, xs)sort!(v, backend; alg=nothing, block_size, …),sort(v, …)sort!(v; backend, alg=Auto(), dims, lt, by, rev, order, workspace),sort(v; …)sortperm!(ix, v, backend; …),sortperm(v, …)sortperm!(ix, v; …)(always overwritesix),sortperm(v; …)merge_sort!(v; block_size),merge_sort(v; …)sort!(v; alg=AK.MergeSort(block_size=…)),sort(v; alg=…)merge_sortperm!(ix, v; …),merge_sortperm(v; …)sortperm!(ix, v; alg=AK.MergeSort()),sortperm(v; alg=…)merge_sortperm_lowmem!(ix, v; …),merge_sortperm_lowmem(v; …)sortperm!(ix, v; alg=AK.MergeSort(lowmem=true)),sortperm(v; alg=…)merge_sortperm!(ix, v; inplace=true)(sortedvas well)ix .= 1:length(v); sort_by_key!(v, ix)merge_sort_by_key!(k, v; …),merge_sort_by_key(k, v; …)sort_by_key!(k, v; …),sort_by_key!(copy(k), copy(v); …)bitonic_sort!(v; …)sort!(v; alg=AK.BitonicSort())sample_sort!(v; …),sample_sortperm!(ix, v; …)sort!(v; alg=AK.CPUThreads.SampleSort()),sortperm!(ix, v; alg=…)alg=AK.SampleSort()alg=AK.CPUThreads.SampleSort()AK.bitonic_defaults(backend)AK.sort_tuning(backend, T)(internal, like every tuning)reduce(op, v, backend; init, neutral=…, block_size, items_per_thread, switch_below, temp)reduce(op, v; init=<optional>, neutral=nothing, acctype=nothing, dims, alg, workspace); an empty input needsinitmapreduce(f, op, v, backend; …),mapreduce(f, op, a, b, …, backend; …)mapreduce(f, op, v, vs...; backend, …)reduce(op, v; dims, temp=dst)(intodst)mapreducedim!(identity, op, dst, v; overwrite=true)sum(v; init=zero(T)),prod,maximum,minimum,countadd_sum/mul_prod;initoptional (count:init=0); results have the accumulator type;sum/prodof empty inputs give zero/oneaccumulate!(op, v, backend; init, block_size, items_per_thread, alg, temp, temp_flags)accumulate!(op, v; backend, init=<optional>, dims, inclusive, acctype, alg, workspace);accumulate!(op, dst, src; …)accumulate(op, v, backend; …),cumsum(v, backend; …),cumprodaccumulate(op, v; …),cumsum(v; …),cumprod(v; …)AccumulateAlgorithm, field-lessScanPrefixes()/DecoupledLookback()ScanAlgorithm; the structs haveblock_size/items_per_threadfields;SliceScanfordimsfindall(v, backend; alg=ScanScatter(block_size=256, items_per_thread=16), max_tasks, …)findall(v; items=keys(v), backend, alg=Auto(), workspace);ScanScatterfields default from the tuningany(pred, v, backend; alg=ConcurrentWrite(), block_size, …),allany(pred, v; backend, alg=Auto(), workspace);ConcurrentWrite(block_size)MapReduce(temp, switch_below),PredicatesAlgorithmViaReduce(reduce::ReduceAlgorithm),PredicateAlgorithmAK.neutral_elementGPUArraysCore.neutral_element(AK's name refers to the same function, but is not public)default_items_per_thread,default_scan_items_per_thread,_radix_defaultsreduce_tuning,scan_tuning,sort_tuning)