From 310f105713abafd6a75e7e6c272e37b916dcdee7 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 23 Sep 2026 17:12:53 +0200 Subject: [PATCH 1/8] Add algorithm selection and move sorting onto it 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). --- .github/workflows/CI-CPU.yml | 51 +-- Project.toml | 4 +- README.md | 12 +- docs/make.jl | 2 + docs/src/api/algorithms.md | 51 +++ docs/src/api/sort.md | 59 ++-- docs/src/tuning.md | 51 +++ src/AcceleratedKernels.jl | 1 + src/algorithms.jl | 155 ++++++++++ src/sort/bitonic_sort.jl | 11 +- src/sort/cpu_sample_sort.jl | 81 +++-- src/sort/merge_sort.jl | 53 +--- src/sort/merge_sort_by_key.jl | 55 +--- src/sort/merge_sortperm.jl | 110 +------ src/sort/radix_sort.jl | 6 +- src/sort/sort.jl | 567 +++++++++++++--------------------- src/sort/tuning.jl | 221 +++++++++++++ test/generic/sort.jl | 565 +++++++++++++++++++-------------- test/generic/sort_resolve.jl | 313 +++++++++++++++++++ test/runtests.jl | 23 +- 20 files changed, 1506 insertions(+), 885 deletions(-) create mode 100644 docs/src/api/algorithms.md create mode 100644 docs/src/tuning.md create mode 100644 src/algorithms.jl create mode 100644 src/sort/tuning.jl create mode 100644 test/generic/sort_resolve.jl diff --git a/.github/workflows/CI-CPU.yml b/.github/workflows/CI-CPU.yml index d64d79f5..a7e99782 100644 --- a/.github/workflows/CI-CPU.yml +++ b/.github/workflows/CI-CPU.yml @@ -102,26 +102,37 @@ jobs: - name: Promote OpenCL to test dependency run: julia test/promote.jl opencl - uses: julia-actions/julia-runtest@v1 - # cpuKA: - # name: KA CPU Backend - # runs-on: ubuntu-latest - # timeout-minutes: 60 - # permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created - # actions: write - # contents: read - # strategy: - # fail-fast: true - # steps: - # - uses: actions/checkout@v7 - # - uses: julia-actions/setup-julia@v3 - # with: - # version: 1 - # arch: x64 - # - uses: julia-actions/cache@v3 - # - uses: julia-actions/julia-buildpkg@v1 - # - uses: julia-actions/julia-runtest@v1 - # with: - # test_args: '--cpu-ka' + cpuKA: + # KernelAbstractions 0.10 (unreleased, `main`) runs AcceleratedKernels' kernels on the host + # backend, through PoCL; `--cpu-ka` tests them there. + name: KA 0.10 host kernels - Julia ${{ matrix.version }} + runs-on: ubuntu-latest + timeout-minutes: 60 + permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created + actions: write + contents: read + strategy: + fail-fast: false + matrix: + version: + - '1.10' + - '1.13' + steps: + - uses: actions/checkout@v7 + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + arch: x64 + - uses: julia-actions/cache@v3 + - name: Develop AcceleratedKernels and KernelAbstractions main into the test environment + run: | + julia --project=test -e ' + using Pkg + Pkg.develop([PackageSpec(path="."), + PackageSpec(url="https://github.com/JuliaGPU/KernelAbstractions.jl")])' + - name: Run tests + # only the generic tests: the others, Aqua included, do not depend on the backend + run: julia --project=test test/runtests.jl --cpu-ka cpu-ka/ docs: name: Documentation runs-on: ubuntu-latest diff --git a/Project.toml b/Project.toml index 3e296b4e..25db4364 100644 --- a/Project.toml +++ b/Project.toml @@ -32,9 +32,9 @@ AMDGPU = "1.3.4, 2" ArgCheck = "2" Atomix = "0.1, 1" CUDA = "6" -CUDACore = "6" +CUDACore = "6.4.1" GPUArraysCore = "0.2.0" -KernelAbstractions = "0.9.34, 0.10" +KernelAbstractions = "0.9.43, 0.10" Markdown = "1" Metal = "1.10.1" OpenCL = "0.10" diff --git a/README.md b/README.md index 0e52c494..261f97f4 100644 --- a/README.md +++ b/README.md @@ -270,12 +270,8 @@ If you need other algorithms in your work that may be of general use, please ope | [General Looping](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/foreachindex/) | `foreachindex`, `foraxes` | `Kokkos::parallel_for` `RAJA::forall` `thrust::transform` | | [Mapping](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/map/) | `map` `map!` | `thrust::transform` | | [Sorting](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/sort/) | `sort` `sort!` | `sort` `sort_team` `stable_sort` | -| | `sample_sort!` `sample_sortperm!` | | -| | `merge_sort` `merge_sort!` | | -| | `merge_sort_by_key` `merge_sort_by_key!` | `sort_team_by_key` | | | `sortperm` `sortperm!` | `sort_permutation` `index_permutation` | -| | `merge_sortperm` `merge_sortperm!` | | -| | `merge_sortperm_lowmem` `merge_sortperm_lowmem!` | | +| | `sort_by_key!` | `sort_by_key` `sort_team_by_key` `cub::DeviceRadixSort::SortPairs` | | [Reduction](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/reduce/) | `reduce` | `Kokkos:parallel_reduce` `fold` `aggregate` | | [MapReduce](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/mapreduce/) | `mapreduce` | `transform_reduce` `fold` | | [Accumulation](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/accumulate/) | `accumulate` `accumulate!` | `prefix_sum` `thrust::scan` `cumsum` | @@ -345,6 +341,12 @@ Start Julia with multiple threads to run the tests on a multithreaded CPU backen $> julia --threads=4 -e 'import Pkg; Pkg.test("AcceleratedKernels.jl")' ``` +With KernelAbstractions 0.10, whose host backend runs kernels on PoCL, `--cpu-ka` tests +AcceleratedKernels' GPU kernels on host arrays instead of the threaded CPU algorithms: +```bash +$> julia -e 'import Pkg; Pkg.test("AcceleratedKernels"; test_args=["--cpu-ka"])' +``` + ## 8. Issues and Debugging As the compilation pipeline of GPU kernels is different to that of base Julia, error messages also look different - for example, where Julia would insert an exception when a variable name was not defined (e.g. we had a typo), a GPU kernel throwing exceptions cannot be compiled and instead you'll see some cascading errors like `"[...] compiling [...] resulted in invalid LLVM IR"` caused by `"Reason: unsupported use of an undefined name"` resulting in `"Reason: unsupported dynamic function invocation"`, etc. diff --git a/docs/make.jl b/docs/make.jl index b399c79b..9c340a93 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -21,6 +21,7 @@ makedocs(; "Performance Tips" => "performance.md", "Manual" =>[ "Using Different Backends" => "api/using_backends.md", + "Algorithms and Backends" => "api/algorithms.md", "General Loops" => "api/foreachindex.md", "Map" => "api/map.md", "Sorting" => "api/sort.md", @@ -37,6 +38,7 @@ makedocs(; "Utilities" => "api/utilities.md", ], "Testing" => "testing.md", + "Tuning and Capabilities" => "tuning.md", "Debugging Kernels" => "debugging.md", "Roadmap" => "roadmap.md", "References" => "references.md", diff --git a/docs/src/api/algorithms.md b/docs/src/api/algorithms.md new file mode 100644 index 00000000..16973847 --- /dev/null +++ b/docs/src/api/algorithms.md @@ -0,0 +1,51 @@ +### Algorithms and backends + +Every algorithmic operation takes an `alg` keyword. Its default, +[`Auto()`](@ref AcceleratedKernels.Auto), lets AcceleratedKernels choose the algorithm and its +settings; any other value is an algorithm you choose, with the settings you give it. + +```@docs +AcceleratedKernels.Algorithm +AcceleratedKernels.Auto +AcceleratedKernels.SortAlgorithm +AcceleratedKernels.CPUThreads +``` + +#### How `Auto` chooses + +`Auto` looks at the backend and its current device, the element type, the layout (a whole array or +slices along `dims`, and their length) and the ordering, never at the array's contents. On the +host backend it chooses the family's `CPUThreads` algorithm, which runs on Julia threads. On a GPU +it follows per-device tuning values: for sorting, `BitonicSort` for short arrays and slices, +`RadixSort` for long whole arrays of the element types it supports, and `MergeSort` otherwise, +subject to `Auto`'s requirements (`stable=true` by default). + +Algorithms carry their settings as fields, and a field left at `nothing` takes the device's +tuned value: +```julia +AK.sort!(v; alg=AK.RadixSort()) # the device's block size and items per thread +AK.sort!(v; alg=AK.RadixSort(block_size=512)) # 512 threads, the device's items per thread +``` + +Explicit algorithms are checked before any data is touched: an algorithm the operation, the +element type, the ordering or the backend cannot run, or an invalid setting, is an +`ArgumentError`. Each new combination of settings compiles new kernels. + +#### The `backend` keyword + +Operations run on the backend of their arrays: `backend` is derived from every array argument, +the destination first, and all of them must agree. Ranges, `CartesianIndices`, `LinearIndices`, +numbers and other non-array arguments do not count. If no argument determines the backend (e.g. +a loop over a range), the operation runs on the host. + +Pass `backend` explicitly for arrays that cannot tell (a range to be processed on a GPU, for +instance), or for memory that several backends can access. AcceleratedKernels does not check that +the backend can reach the arrays. The device and stream are those of the calling task, as set by +the backend package (e.g. `CUDA.device!`); make sure the arrays live on that device. + +#### The host backend + +Arrays in host memory (`Array` and views of it) are on the host backend. `Auto` processes them on +Julia threads with the `CPUThreads` algorithms. On KernelAbstractions 0.10 the host backend also +runs AcceleratedKernels' GPU kernels (on PoCL), so an explicit kernel algorithm such as +`MergeSort()` works on host arrays there; on KernelAbstractions 0.9 it is an `ArgumentError`. diff --git a/docs/src/api/sort.md b/docs/src/api/sort.md index e01d9c6d..fb4fe391 100644 --- a/docs/src/api/sort.md +++ b/docs/src/api/sort.md @@ -1,9 +1,10 @@ ### `sort` and friends -Sorting algorithms with similar interface and default settings as the Julia Base ones, on GPUs: +Sorting with the interface and results of the Julia Base functions, on GPUs and on multithreaded CPUs: - `sort!` (in-place), `sort` (out-of-place) - `sortperm!`, `sortperm` -- **Other names**: `sort`, `sort_team`, `sort_team_by_key`, `stable_sort` or variations in Kokkos, RAJA, Thrust that I know of. +- `sort_by_key!`, which sorts keys and applies the same permutation to values +- **Other names**: `sort`, `sort_team`, `sort_team_by_key`, `stable_sort`, `sort_by_key` or variations in Kokkos, RAJA, Thrust, oneDPL and CUB. Function signatures: ```@docs @@ -11,27 +12,45 @@ AcceleratedKernels.sort! AcceleratedKernels.sort AcceleratedKernels.sortperm! AcceleratedKernels.sortperm +AcceleratedKernels.sort_by_key! ``` -Algorithm choice is available on `sort!` / `sort` / `sortperm!` / `sortperm` with `alg=AK.MergeSort()`, -`alg=AK.MergeSort(lowmem=true)`, `alg=AK.RadixSort()`, `alg=AK.BitonicSort()`, or -`alg=AK.SampleSort()`, depending on the backend and operation. +Example: +```julia +import AcceleratedKernels as AK +using AMDGPU + +v = ROCArray(rand(Int32, 100_000)) +AK.sort!(v) +``` + +Like Base, the sorts are stable by default: elements that compare equal keep their order. +[`Auto()`](@ref AcceleratedKernels.Auto), the default `alg`, picks an algorithm for the backend, +the device, the element type, the ordering and the length. Pass an algorithm to choose it and its +settings yourself; an algorithm that cannot run the call is an `ArgumentError`, never replaced: + +| Algorithm | Stable | `dims` | `sortperm!` | `sort_by_key!` | Element types, orderings | Backends | +|---|---|---|---|---|---|---| +| [`MergeSort`](@ref AcceleratedKernels.MergeSort) | yes | yes | yes | yes | all | kernels | +| [`RadixSort`](@ref AcceleratedKernels.RadixSort) | yes | no | no | no | 32/64-bit integers and floats, default ordering or its reverse | kernels | +| [`BitonicSort`](@ref AcceleratedKernels.BitonicSort) | no | yes | no | no | all | kernels | +| [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort) | yes | yes | yes | yes | all | host | + +"Kernels" means GPU backends, and the host backend on KernelAbstractions 0.10, which runs +AcceleratedKernels' kernels on PoCL. See [Algorithms and backends](@ref) for how `Auto` chooses. -Function signatures: ```@docs AcceleratedKernels.MergeSort AcceleratedKernels.RadixSort AcceleratedKernels.BitonicSort -AcceleratedKernels.SampleSort +AcceleratedKernels.CPUThreads.SampleSort ``` -Example: ```julia -import AcceleratedKernels as AK -using AMDGPU - -v = ROCArray(rand(Int32, 100_000)) -AK.sort!(v) +v = ROCArray(rand(Float32, 1_000_000)) +AK.sort!(v; alg=AK.RadixSort()) # this algorithm, with the device's settings +AK.sort!(v; alg=AK.RadixSort(items_per_thread=4)) # and this setting +AK.sort!(v; alg=AK.Auto(stable=false)) # allow unstable algorithms (Base's QuickSort) ``` Multidimensional arrays are sorted as one flat vector by default; pass `dims` to sort each 1D slice @@ -43,18 +62,10 @@ AK.sort!(A; dims=1) # each column sorted ix = AK.sortperm(A; dims=2) # A[ix] has each row sorted ``` -On GPU backends `dims` uses merge sort by default (`RadixSort()` does not support it); on CPU -backends each slice is sorted with `Base.sort!`. `BitonicSort()` is unstable and supports `sort!` -and `sort`, including `dims`, but not `sortperm!` or `sortperm`. It is fastest for small arrays -and short slices, slower than `MergeSort`/`RadixSort` for large whole-array sorts: -```julia -A = ROCArray(rand(Float32, 64, 100_000)) -AK.sort!(A; dims=1, alg=AK.BitonicSort()) -``` - -As GPU memory is more expensive, all functions in AcceleratedKernels.jl expose any temporary arrays they will use (the `temp` argument); you can supply your own buffers to make the algorithms not allocate additional GPU storage, e.g.: +As GPU memory is more expensive, the sorts accept their temporary arrays (`temp`, and `temp_keys` +and `temp_values` for `sort_by_key!`), so that you can supply your own buffers, e.g.: ```julia v = ROCArray(rand(Float32, 100_000)) temp = similar(v) -AK.sort!(v, temp=temp) +AK.sort!(v; temp) ``` diff --git a/docs/src/tuning.md b/docs/src/tuning.md new file mode 100644 index 00000000..3cfcd14c --- /dev/null +++ b/docs/src/tuning.md @@ -0,0 +1,51 @@ +# Tuning and capabilities (developer notes) + +This page is for contributors and for AcceleratedKernels' own package extensions. Nothing here is +public API: the hooks and tuning structs can change in any release, including patch releases. + +## Resolution + +Every operation resolves its `alg` keyword once, on the host, before touching any data. For +sorting, `_resolve_sort` runs four steps: + +1. `_checkdomain` checks the fields the caller set (e.g. a `block_size` that is not a power of + two), before any arithmetic uses them; +2. for `Auto`, `_select_sort` picks an algorithm from the device's tuning, the backend's + capabilities and the call's static facts, returning it with unset fields; +3. `_fill` fills the unset fields from the tuning; +4. `_check` checks the complete algorithm against the capabilities, the operation and the + arguments, and throws an `ArgumentError` for anything it cannot run. + +```@docs +AcceleratedKernels._resolve_sort +``` + +`Auto` and explicit algorithms share steps 3 and 4, so a tuning cannot make an invalid algorithm +run, and an explicit setting always wins over the tuning. + +## Tunings + +One plain struct per operation family holds the values that drive selection and fill unset +fields, and one hook returns it for a backend and element type: + +```@docs +AcceleratedKernels.SortTuning +AcceleratedKernels.sort_tuning +``` + +AcceleratedKernels defines the hook's generic method, whose values reproduce the library's +historical defaults. A package extension adds one method for its backend type and may choose +values per device, e.g. from the compute capability of the current CUDA device. Tuning queries run +on the calling task, under its current device, never at load or precompilation time. Record the +measurement behind each value (see `benchmark/tune_sort.jl`) next to it. + +## Capabilities + +Capabilities are correctness facts about a backend, checked for `Auto` and explicit algorithms +alike; tunings cannot change them. + +```@docs +AcceleratedKernels._runs_threads +AcceleratedKernels._runs_kernels +AcceleratedKernels._resolve_backend +``` diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index c0205561..fc091af8 100644 --- a/src/AcceleratedKernels.jl +++ b/src/AcceleratedKernels.jl @@ -26,6 +26,7 @@ const get_backend = KernelAbstractions.get_backend # Include code from other files include("utils.jl") +include("algorithms.jl") include("task_partitioner.jl") include("foreachindex.jl") include("map.jl") diff --git a/src/algorithms.jl b/src/algorithms.jl new file mode 100644 index 00000000..8fdc07c3 --- /dev/null +++ b/src/algorithms.jl @@ -0,0 +1,155 @@ +# Algorithms, capabilities and backend resolution: the layer every operation's host API shares. + + +""" + Algorithm + +Supertype of the algorithms accepted by AcceleratedKernels' operations through their `alg` +keyword. [`Auto`](@ref) lets AK choose; a concrete algorithm (e.g. [`MergeSort`](@ref)) is +honoured or rejected with an `ArgumentError`, never silently replaced. + +Algorithms carry their tunable settings as fields; a field left at `nothing` is filled in from +AK's defaults for the device the operation runs on. +""" +abstract type Algorithm end + + +""" + Auto(; stable=true) + +Let AcceleratedKernels choose the algorithm and its settings for the current device, subject to +the requirements given as fields. This is the default `alg` of every operation. + +`stable` applies to sorting only (other operations ignore it): with `stable=true` the result is +the one a stable sort would produce. An unstable algorithm is chosen only where no one can tell +the difference, i.e. where elements that compare equal are bitwise identical (integers, `Bool` and +`Char` under the default ordering). `stable=false` also allows unstable algorithms for other +element types; this is what Base's `alg=QuickSort` means. + +Selection never reads the array's contents. +""" +Base.@kwdef struct Auto <: Algorithm + stable::Bool = true +end + + +# Algorithm families, for documentation and dispatch. An operation's `alg` keyword accepts any +# `Algorithm`, and rejects those that do not implement it. + +""" + SortAlgorithm <: Algorithm + +Supertype of the sorting algorithms: [`MergeSort`](@ref), [`RadixSort`](@ref), +[`BitonicSort`](@ref) and [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort). +""" +abstract type SortAlgorithm <: Algorithm end + + +""" + AcceleratedKernels.CPUThreads + +Algorithms that run on Julia threads, for arrays on the host backend. +[`Auto`](@ref AcceleratedKernels.Auto) chooses them for host arrays; they are rejected for any +other backend. +""" +module CPUThreads + +import ..AcceleratedKernels: SortAlgorithm + +""" + CPUThreads.SampleSort(; max_tasks=nothing, min_elems=nothing) + +Parallel sample sort on Julia threads, deferring to `Base.sort!` for the local sorts; it is +stable, and also provides `sortperm!` and [`sort_by_key!`](@ref AcceleratedKernels.sort_by_key!). +Uses at most `max_tasks` tasks (default `Threads.nthreads()`), each with at least `min_elems` +elements (default 1). Only runs on the host backend. +""" +Base.@kwdef struct SampleSort <: SortAlgorithm + max_tasks::Union{Nothing, Int} = nothing + min_elems::Union{Nothing, Int} = nothing +end + +end # module CPUThreads + + +# Capabilities: correctness facts about a backend, checked for `Auto` and explicit algorithms +# alike. Unlike tunings, they cannot be changed per device. + +# The host backend: `KernelAbstractions.CPU` on KernelAbstractions 0.9, its PoCL backend on 0.10. +const HOST_BACKEND = get_backend(Int[]) +const HostBackend = typeof(HOST_BACKEND) + +""" + _runs_threads(backend) + +Whether `backend` is the host backend, whose arrays the `CPUThreads` algorithms can process on +Julia threads. +""" +_runs_threads(::Backend) = false +_runs_threads(::HostBackend) = true + +""" + _runs_kernels(backend) + +Whether AK's kernels (`@kernel cpu=false`) run on `backend`: every GPU backend, and the host +backend of KernelAbstractions 0.10, which compiles kernels for PoCL. KernelAbstractions 0.9's +`CPU` backend cannot run them. +""" +_runs_kernels(::Backend) = true +_runs_kernels(::HostBackend) = nameof(HostBackend) !== :CPU + + +# Backend resolution + +_backend_name(b::Backend) = nameof(typeof(b)) + +# Values that never determine the backend: lazy index collections, scalars, and any other value +# that is not an array. Every other array votes with `get_backend`, which throws for array types +# that do not implement it. +_backend_vote(_) = nothing +_backend_vote(x::AbstractArray) = get_backend(x) +_backend_vote(::Union{AbstractRange, CartesianIndices, LinearIndices}) = nothing +_backend_vote(x::Tuple) = _backend_votes(x...) +_backend_vote(bc::Base.Broadcast.Broadcasted) = _backend_votes(bc.args...) +_backend_vote(x::Base.Broadcast.Extruded) = _backend_vote(x.x) + +_backend_votes() = nothing +_backend_votes(x, xs...) = _backend_merge(_backend_vote(x), _backend_votes(xs...)) + +_backend_merge(::Nothing, ::Nothing) = nothing +_backend_merge(a, ::Nothing) = a +_backend_merge(::Nothing, b) = b +function _backend_merge(a, b) + a == b || throw(ArgumentError( + "the arguments live on different backends ($(_backend_name(a)) and " * + "$(_backend_name(b))); pass `backend` explicitly if they are accessible from one of them")) + a +end + +# Lazy index collections, and Base's views, reshapes and permutations of them, have no backend +# (unlike `_backend_vote`, this never asks an array for its backend, which an explicit `backend` +# makes unnecessary) +_backend_free(_) = false +_backend_free(::Union{AbstractRange, CartesianIndices, LinearIndices}) = true +_backend_free(x::Union{SubArray, Base.ReshapedArray, PermutedDimsArray}) = _backend_free(parent(x)) + +# An array for an operation's result: `similar(v, ...)`, unless `v` has no backend (a range, for +# instance), which gives a new array on `backend` +_similar(backend, v, ::Type{T}=eltype(v), dims=size(v)) where {T} = + _backend_free(v) ? KernelAbstractions.allocate(backend, T, dims) : similar(v, T, dims) +# A copy of `v`, likewise +_copy(backend, v) = _backend_free(v) ? copyto!(_similar(backend, v), collect(v)) : copy(v) + +""" + _resolve_backend(backend, args...) + +The backend an operation runs on: `backend` if given, else the one every array in `args` +(destination first) agrees on, recursing into `Broadcasted` trees; ranges, `CartesianIndices`, +`LinearIndices` and non-array values do not count. Arguments on different backends are an +`ArgumentError`. If no argument determines it, the host backend is used. +""" +_resolve_backend(backend::Backend, args...) = backend +function _resolve_backend(::Nothing, args...) + b = _backend_votes(args...) + b === nothing ? HOST_BACKEND : b +end diff --git a/src/sort/bitonic_sort.jl b/src/sort/bitonic_sort.jl index 60678bbd..c60f56a7 100644 --- a/src/sort/bitonic_sort.jl +++ b/src/sort/bitonic_sort.jl @@ -98,15 +98,8 @@ end end -""" - bitonic_defaults(backend::Backend) - -Default `block_size` and `items_per_thread` for [`BitonicSort`](@ref) on `backend`. -""" -bitonic_defaults(::Backend) = (block_size=256, items_per_thread=8) - - -function bitonic_sort!( +# GPU bitonic sort of `v` (or of each slice along `dims`), in place. +function _bitonic_sort!( v::AbstractArray, backend::Backend=get_backend(v); lt=isless, diff --git a/src/sort/cpu_sample_sort.jl b/src/sort/cpu_sample_sort.jl index 84f03042..f5c39952 100644 --- a/src/sort/cpu_sample_sort.jl +++ b/src/sort/cpu_sample_sort.jl @@ -171,21 +171,8 @@ end -""" - sample_sort!( - v::AbstractArray; - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - max_tasks=Threads.nthreads(), - min_elems=1, - temp::Union{Nothing, AbstractArray}=nothing, - ) -""" -function sample_sort!( +# Parallel sample sort on Julia threads, deferring to `Base.sort!` for the local sorts. +function _sample_sort!( v::AbstractArray; lt=isless, @@ -257,21 +244,8 @@ end -""" - sample_sortperm!( - ix::AbstractArray, v::AbstractArray; - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - max_tasks=Threads.nthreads(), - min_elems=1, - temp::Union{Nothing, AbstractArray}=nothing, - ) -""" -function sample_sortperm!( +# Sort permutation on Julia threads: sample sort of the indices, comparing the values. +function _sample_sortperm!( ix::AbstractArray, v::AbstractArray; lt=isless, @@ -288,7 +262,7 @@ function sample_sortperm!( @argcheck length(ix) == length(v) # Initialise indices that will be sorted by the keys in v - foreachindex(ix; max_tasks, min_elems) do i + foreachindex(ix, HOST_BACKEND; max_tasks, min_elems) do i @inbounds ix[i] = i end @@ -305,7 +279,7 @@ end function _sample_sort_barrier!(ix, v, ord; max_tasks, min_elems, temp) # Construct custom comparator indexing into global array v for every index comparison comp = (ix, iy) -> Base.Order.lt(ord, v[ix], v[iy]) - sample_sort!( + _sample_sort!( ix; lt=comp, @@ -315,3 +289,46 @@ function _sample_sort_barrier!(ix, v, ord; max_tasks, min_elems, temp) max_tasks, min_elems, temp, ) end + + +# Sort permutation of each slice along `dims`, like Base: sort the linear indices of each slice by +# the values they point to. +function _sample_sortperm_dims!(ix, v, ord, dims; max_tasks, min_elems) + perm_ord = Base.Order.Perm(ord, vec(v)) + copyto!(ix, LinearIndices(v)) + foreach_slice(ix, dims; max_tasks, min_elems) do slice + Base.sort!(slice; order=perm_ord) + end + ix +end + + +# Key/value sort on Julia threads: find the stable sorting permutation of the keys (of each slice +# along `dims`), then gather keys and values through it. +function _sample_sort_by_key!(keys, values, ord, dims; max_tasks, min_elems, temp_keys, temp_values) + for (v, temp) in ((keys, temp_keys), (values, temp_values)) + isnothing(temp) && continue + @argcheck length(temp) == length(v) + @argcheck eltype(temp) === eltype(v) + end + ix = similar(keys, Int) + if dims isa Colon + _sample_sortperm!(vec(ix), vec(keys); order=ord, max_tasks, min_elems) + else + _sample_sortperm_dims!(ix, keys, ord, dims; max_tasks, min_elems) + end + _gather_through!(keys, ix, temp_keys; max_tasks, min_elems) + _gather_through!(values, ix, temp_values; max_tasks, min_elems) + keys, values +end + +# v[i] = v_old[ix[i]] for every linear index i +function _gather_through!(v, ix, temp; max_tasks, min_elems) + old = isnothing(temp) ? copy(v) : copyto!(temp, v) + task_partition(length(v), max_tasks, min_elems) do irange + @inbounds for i in irange + v[i] = old[ix[i]] + end + end + v +end diff --git a/src/sort/merge_sort.jl b/src/sort/merge_sort.jl index 72a04dc3..732f5d7c 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -128,23 +128,8 @@ end end -""" - merge_sort!( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - - # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sort!( +# GPU merge sort of `v` (or of each slice along `dims`), in place. +function _merge_sort!( v::AbstractArray, backend::Backend=get_backend(v); lt=isless, @@ -166,10 +151,12 @@ function merge_sort!( end (isempty(v) || layout.len <= 1) && return v - # Compute keys once instead of evaluating `by` in every comparison. + # Compute keys once instead of evaluating `by` in every comparison, into an array of the + # keys' type (broadcasting would make a `BitArray` of `Bool` keys on the host) if by !== identity - keys = by.(v) - merge_sort_by_key!( + keys = similar(v, Base.promote_op(by, eltype(v))) + map!(by, keys, v, backend; block_size) + _merge_sort_by_key!( keys, v, backend; lt, rev, order, block_size, dims, temp_values=temp, # temp was for v swap buffer; maps to temp_values here @@ -219,29 +206,3 @@ function merge_sort!( v end - - -""" - merge_sort( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sort( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - v_copy = copy(v) - merge_sort!( - v_copy, backend; - kwargs... - ) -end diff --git a/src/sort/merge_sort_by_key.jl b/src/sort/merge_sort_by_key.jl index 6d10b47b..844c1ad1 100644 --- a/src/sort/merge_sort_by_key.jl +++ b/src/sort/merge_sort_by_key.jl @@ -161,26 +161,8 @@ end end -""" - merge_sort_by_key!( - keys::AbstractArray, - values::AbstractArray, - backend::Backend=get_backend(keys); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp_keys::Union{Nothing, AbstractArray}=nothing, - temp_values::Union{Nothing, AbstractArray}=nothing, - - # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sort_by_key!( +# GPU merge sort of `keys` (or of each slice along `dims`), permuting `values` alike. +function _merge_sort_by_key!( keys::AbstractArray, values::AbstractArray, backend::Backend=get_backend(keys); @@ -260,36 +242,3 @@ function merge_sort_by_key!( keys, values end - - -""" - merge_sort_by_key( - keys::AbstractArray, - values::AbstractArray, - backend::Backend=get_backend(keys); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp_keys::Union{Nothing, AbstractArray}=nothing, - temp_values::Union{Nothing, AbstractArray}=nothing, - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sort_by_key( - keys::AbstractArray, - values::AbstractArray, - backend::Backend=get_backend(keys); - kwargs... -) - keys_copy = copy(keys) - values_copy = copy(values) - - merge_sort_by_key!( - keys_copy, values_copy, backend; - kwargs... - ) -end diff --git a/src/sort/merge_sortperm.jl b/src/sort/merge_sortperm.jl index 44b50f46..1500047a 100644 --- a/src/sort/merge_sortperm.jl +++ b/src/sort/merge_sortperm.jl @@ -1,24 +1,5 @@ -""" - merge_sortperm!( - ix::AbstractArray, - v::AbstractArray, - backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - inplace::Bool=false, - block_size::Int=256, - temp_ix::Union{Nothing, AbstractArray}=nothing, - temp_v::Union{Nothing, AbstractArray}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sortperm!( +# GPU merge sort permutation: sorts a copy of the keys, carrying the indices along. +function _merge_sortperm!( ix::AbstractArray, v::AbstractArray, backend::Backend=get_backend(v); @@ -51,13 +32,13 @@ function merge_sortperm!( end # Initialise the linear indices that will be sorted by the keys in v - foreachindex(ix, block_size=block_size) do i + foreachindex(ix, backend; block_size) do i @inbounds ix[i] = i end (isempty(v) || layout.len <= 1) && return ix - keys = inplace ? v : copy(v) + keys = inplace ? v : _copy(backend, v) - merge_sort_by_key!( + _merge_sort_by_key!( keys, ix, backend; lt, by, rev, order, block_size, dims, temp_keys=temp_v, temp_values=temp_ix, @@ -67,55 +48,8 @@ function merge_sortperm!( end -""" - merge_sortperm( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - inplace::Bool=false, - block_size::Int=256, - temp_ix::Union{Nothing, AbstractArray}=nothing, - temp_v::Union{Nothing, AbstractArray}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sortperm( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - ix = similar(v, Int) - merge_sortperm!( - ix, v, backend; - kwargs... - ) -end - - -""" - merge_sortperm_lowmem!( - ix::AbstractArray, - v::AbstractArray, - backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sortperm_lowmem!( +# GPU merge sort permutation comparing the keys in global memory, without copying them. +function _merge_sortperm_lowmem!( ix::AbstractArray, v::AbstractArray, backend::Backend=get_backend(v); @@ -141,7 +75,7 @@ function merge_sortperm_lowmem!( end # Initialise the linear indices that will be sorted by the keys in v - foreachindex(ix, block_size=block_size) do i + foreachindex(ix, backend; block_size) do i @inbounds ix[i] = i end (isempty(ix) || layout.len <= 1) && return ix @@ -188,31 +122,3 @@ function merge_sortperm_lowmem!( ix end - - -""" - merge_sortperm_lowmem( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - ) -""" -function merge_sortperm_lowmem( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - ix = similar(v, Int) - merge_sortperm_lowmem!( - ix, v, backend; - kwargs... - ) -end diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index c538c607..db8c9eed 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -367,13 +367,13 @@ end # Return the extrema of the transformed sort keys. -function _rs_key_range(v::AbstractArray{T}, descending::Bool) where T +function _rs_key_range(v::AbstractArray{T}, backend::Backend, descending::Bool) where T K = typeof(_to_sort_key(zero(T))) ident = (typemax(K), typemin(K)) min_k, max_k = mapreduce( x -> (k = _to_sort_key(x); (k, k)), (a, b) -> (min(a[1], b[1]), max(a[2], b[2])), - v; + v, backend; init=ident, neutral=ident, ) @@ -435,7 +435,7 @@ function _radix_sort!( ndrange = (block_size * num_blocks,) - min_key, max_key = _rs_key_range(p1, descending) + min_key, max_key = _rs_key_range(p1, backend, descending) vitems = Val(items) hist_kern! = has_atomics ? diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 7876b8d9..9d347f7e 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -8,456 +8,311 @@ include("radix_sort.jl") include("bitonic_sort.jl") -# Available sorting algorithms -abstract type SortAlgorithm end +# Sorting algorithms (`CPUThreads.SampleSort` is defined with the other CPUThreads algorithms) """ - MergeSort(; lowmem=false) + MergeSort(; block_size=nothing, lowmem=false) -Use GPU merge sort for `sort!` and `sort`. For `sortperm!`, `lowmem=true` selects the -lower-memory permutation path. +GPU merge sort: stable, for every element type and ordering, whole arrays and `dims`, and the only +kernel algorithm for `sortperm!` and [`sort_by_key!`](@ref). Each block of `block_size` threads +(any positive number) sorts a tile of `2 * block_size` elements in local memory, then global +passes merge the tiles. `lowmem=true` selects a `sortperm!` path that does not copy the keys, at +the cost of reading them from global memory in every comparison; other operations reject it. """ Base.@kwdef struct MergeSort <: SortAlgorithm + block_size::Union{Nothing, Int} = nothing lowmem::Bool = false end """ RadixSort(; block_size=nothing, items_per_thread=nothing) -Use GPU radix sort for `sort!` and `sort`. Supports `UInt32`, `Int32`, `Float32`, `UInt64`, -`Int64`, and `Float64` with forward or reverse ordering. This algorithm does not support -`sortperm!`. +GPU LSD radix sort for whole arrays of 32- and 64-bit integers and floats (`UInt32`, `Int32`, +`Float32`, `UInt64`, `Int64`, `Float64`) under the default ordering or its reverse; it is stable +and orders floats like `isless`. It does not support `dims`, custom `lt`/`by`, `sortperm!` or +`sort_by_key!`. `block_size` must be a power of two up to 1024, `items_per_thread` between 1 and +64. `items_per_thread` applies where radix sort's chunked kernels run, which needs atomics and a +`block_size` that is a multiple of 32 whose tiles fit local memory; elsewhere its portable kernels +sort one item per thread. """ Base.@kwdef struct RadixSort <: SortAlgorithm block_size::Union{Nothing, Int} = nothing items_per_thread::Union{Nothing, Int} = nothing end -_radix_defaults(::Backend) = (block_size=256, items_per_thread=2) - """ BitonicSort(; block_size=nothing, items_per_thread=nothing) -Use GPU bitonic sort for `sort!` and `sort`, whole-array or along `dims`. Supports GPU-compatible -element types and `lt`/`by`/`rev`/`order`; `by` is evaluated at each comparison. The sort is -unstable and does not support `sortperm!` or `sortperm`. - -Fastest for small arrays and short slices, slower than `MergeSort`/`RadixSort` for large -whole-array sorts. Tiles of up to `block_size * items_per_thread` elements sort in local memory; -larger inputs need global passes. - -Both settings must be positive powers of two and default to `bitonic_defaults(backend)`. -For `block_size`, the `sort!` keyword takes precedence over the backend default. +GPU bitonic sorting network, for whole arrays and `dims`, every element type and ordering. It is +unstable and does not support `sortperm!` or `sort_by_key!`. Fastest for small arrays and short +slices: tiles of `block_size * items_per_thread` elements sort in local memory, larger inputs need +global passes. Both settings must be positive powers of two. """ Base.@kwdef struct BitonicSort <: SortAlgorithm block_size::Union{Nothing, Int} = nothing items_per_thread::Union{Nothing, Int} = nothing end -""" - SampleSort() - -Use CPU sample sort for `sort!`, `sort`, `sortperm!`, and `sortperm`. -""" -struct SampleSort <: SortAlgorithm end - -# All other algorithms have the same naming convention as Julia Base ones; provide similar -# interface here too. +include("tuning.jl") """ sort!( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, + v::AbstractArray; + backend=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=:, + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, + temp::Union{Nothing, AbstractArray}=nothing, + ) -> v - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # Algorithm choice - alg::Union{Nothing, SortAlgorithm}=nothing, +Sort `v` in place. `lt`, `by`, `rev` and `order` are those of `Base.sort!`, and so is the result: +by default the sort is stable. `dims=:` sorts the whole array as one vector (`Base.sort!` requires +`dims` for arrays of more than one dimension); an integer `dims` sorts each slice along that +dimension independently. - # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector - dims::Union{Colon, Integer}=Colon(), +`alg` is [`Auto()`](@ref Auto) by default, which chooses an algorithm for the backend and device, +the element type, the ordering, and the length of the array or slices: on the host, +[`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort); on GPUs, +[`MergeSort`](@ref), or [`BitonicSort`](@ref) for short inputs or [`RadixSort`](@ref) for long +ones where the device's tuning enables them and they give the same result. `Auto(stable=false)` +also allows the unstable `BitonicSort` for elements that compare equal without being identical +(e.g. floats). An explicit algorithm, with any of its settings, is used as given or rejected with +an `ArgumentError`. - # GPU settings - block_size::Union{Nothing, Int}=nothing, +`backend` is derived from `v`; pass it only for arrays that do not determine their backend. - # Temporary buffer, same size as `v` - temp::Union{Nothing, AbstractArray}=nothing, - ) - -Sorts the array `v` in-place using the specified backend. The `lt`, `by`, `rev`, and `order` -arguments are the same as for `Base.sort`. - -By default (`dims=:`) the whole array is sorted as one vector, whatever its shape. Pass an integer -`dims` to sort each 1D slice along that dimension independently, like `Base.sort!(A; dims)`. On GPU -backends this uses merge sort by default, or [`BitonicSort`](@ref) when requested; `RadixSort` -does not support `dims`. On CPU backends each slice is sorted with -`Base.sort!` (also with `alg=SampleSort()`, and without using `temp`), slices spread over the tasks. - -## CPU -CPU settings: use at most `max_tasks` threads to sort the array such that at least `min_elems` -elements are sorted by each thread. A parallel sample sort is used, processing -independent slices of the array and deferring to `Base.sort!` for the final local sorts. - -Note that the Base Julia `sort!` is mainly memory-bound, so multithreaded sorting only becomes -faster if it is a more compute-heavy operation to hide memory latency - that includes: -- Sorting more complex types, e.g. lexicographic sorting of tuples / structs / strings. -- More complex comparators, e.g. `by=custom_complex_function` or `lt=custom_lt_function`. -- Less cache-predictable data movement, e.g. `sortperm`. - -## GPU -GPU settings: `block_size` sets the number of threads per block. For `RadixSort` and `BitonicSort`, -fields on the algorithm take precedence over this keyword, then backend defaults. -`items_per_thread` is set on the algorithm and defaults to 2 for `RadixSort`, 8 for `BitonicSort`. - -## Algorithm choice -By default, `sort!` uses sample sort on CPU backends and merge sort on GPU -backends. Pass `alg=SampleSort()` for the CPU path, `alg=MergeSort()` for the GPU merge-sort path, -`alg=RadixSort()` to opt into GPU radix sorting, or `alg=BitonicSort()` for the GPU sorting -network. `RadixSort()` supports 32-bit and 64-bit integers and floats with default `lt`/`by`. -`BitonicSort()` is unstable: fastest for small arrays and short slices, slower than -`MergeSort`/`RadixSort` for large whole-array sorts. - -For both CPU and GPU backends, the `temp` argument can be used to reuse a temporary buffer of the -same size as `v` to store the sorted output. +`temp` is an optional scratch array with the same length and element type as `v`, used by +`MergeSort`, `RadixSort` and, for `dims=:`, `CPUThreads.SampleSort`. # Examples -Simple parallel CPU sort using all available threads (as given by `julia --threads N`): ```julia import AcceleratedKernels as AK -v = rand(1000) -AK.sort!(v) -``` +using CUDA -Parallel GPU sorting, passing a temporary buffer to avoid allocating a new one: -```julia -using oneAPI -import AcceleratedKernels as AK -v = oneArray(rand(1000)) -temp = similar(v) -AK.sort!(v, temp=temp) +v = CuArray(rand(Float32, 100_000)) +AK.sort!(v) # Auto: stable, chosen for the device +AK.sort!(v; rev=true, alg=AK.Auto(stable=false)) # may use an unstable algorithm +AK.sort!(v; alg=AK.RadixSort(block_size=512)) # this algorithm, with this setting + +A = CuArray(rand(Int32, 64, 10_000)) +AK.sort!(A; dims=1) # each column + +AK.sort!(rand(1000)) # host array: Julia threads ``` """ function sort!( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - _sort_impl!( - v, backend; - kwargs... - ) -end - - -function _sort_impl!( - v::AbstractArray, backend::Backend; - + v::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=Colon(), lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Forward, - - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, + order::Base.Order.Ordering=Base.Order.Forward, + temp::Union{Nothing, AbstractArray}=nothing, +) + backend = _resolve_backend(backend, v, temp) + ord = Base.Order.ord(lt, by, rev, order) + a = _resolve_sort(alg, backend, v, dims, ord) + _sort_impl!(a, v, backend, dims, ord; lt, by, rev, order, temp) + return v +end - alg::Union{Nothing, SortAlgorithm}=nothing, +function _sort_impl!(a::MergeSort, v, backend, dims, ord; lt, by, rev, order, temp) + _merge_sort!(v, backend; lt, by, rev, order, block_size=a.block_size, temp, dims) +end - # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector - dims::Union{Colon, Integer}=Colon(), +function _sort_impl!(a::RadixSort, v, backend, dims, ord; lt, by, rev, order, temp) + _radix_sort!(v, backend; descending=ord === Base.Order.Reverse, + block_size=a.block_size, items_per_thread=a.items_per_thread, temp) +end - # GPU settings; nothing => each GPU algorithm picks its own tuned default - block_size::Union{Nothing, Int}=nothing, +function _sort_impl!(a::BitonicSort, v, backend, dims, ord; lt, by, rev, order, temp) + _bitonic_sort!(v, backend; lt, by, rev, order, dims, + block_size=a.block_size, items_per_thread=a.items_per_thread) +end - # Temporary buffer, same size as `v` - temp::Union{Nothing, AbstractArray}=nothing, -) - if use_gpu_algorithm(backend, prefer_threads) - alg = isnothing(alg) ? MergeSort() : alg - if alg isa MergeSort - merge_sort!( - v, backend; - lt, by, rev, order, - block_size=isnothing(block_size) ? 256 : block_size, - temp, dims, - ) - elseif alg isa RadixSort - dims isa Colon || throw(ArgumentError("RadixSort does not support sorting along `dims`")) - _rs_supported(eltype(v)) || throw(ArgumentError("RadixSort is not supported for eltype \"$(eltype(v))\"")) - ordering = Base.Order.ord(lt, by, rev, order) - ordering === Base.Order.Forward || ordering === Base.Order.Reverse || - throw(ArgumentError("RadixSort only supports forward or reverse ordering")) - defaults = _radix_defaults(backend) - radix_block_size = isnothing(alg.block_size) ? - (isnothing(block_size) ? defaults.block_size : block_size) : alg.block_size - radix_items = isnothing(alg.items_per_thread) ? - defaults.items_per_thread : alg.items_per_thread - _radix_sort!( - v, backend; - descending=ordering === Base.Order.Reverse, - block_size=radix_block_size, - items_per_thread=radix_items, - temp, - ) - elseif alg isa BitonicSort - defaults = bitonic_defaults(backend) - bitonic_sort!( - v, backend; - lt, by, rev, order, dims, - block_size=something(alg.block_size, block_size, defaults.block_size), - items_per_thread=something(alg.items_per_thread, defaults.items_per_thread), - ) - else - throw(ArgumentError("$(typeof(alg)) is not supported by sort! on GPU backends")) - end +function _sort_impl!(a::CPUThreads.SampleSort, v, backend, dims, ord; lt, by, rev, order, temp) + if dims isa Colon + # `vec`: the local sorts use `Base.sort!`, which needs `dims` for other arrays + _sample_sort!(vec(v); lt, by, rev, order, max_tasks=a.max_tasks, min_elems=a.min_elems, temp) else - alg = isnothing(alg) ? SampleSort() : alg - if !(alg isa SampleSort) - throw(ArgumentError("$(typeof(alg)) is not supported by sort! on CPU backends")) - elseif dims isa Colon - sample_sort!( - v; - lt, by, rev, order, - max_tasks, min_elems, - temp, - ) - else - ord = Base.Order.ord(lt, by, rev, order) - foreach_slice(v, dims; max_tasks, min_elems) do slice - Base.sort!(slice; order=ord) - end + foreach_slice(v, dims; max_tasks=a.max_tasks, min_elems=a.min_elems) do slice + Base.sort!(slice; order=ord) end end - v end """ - sort( - v::AbstractArray, backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, + sort(v::AbstractArray; kwargs...) - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # Algorithm choice - alg::Union{Nothing, SortAlgorithm}=nothing, - - # Sort each 1D slice along this dimension; `:` sorts the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - - # GPU settings - block_size::Union{Nothing, Int}=nothing, - - # Temporary buffer, same size as `v` - temp::Union{Nothing, AbstractArray}=nothing, - ) - -Out-of-place sort, same settings as [`sort!`](@ref). +Out-of-place [`sort!`](@ref): sort a copy of `v`, with the same keywords. """ -function sort( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - vcopy = copy(v) - sort!( - vcopy, backend; - kwargs... - ) +function sort(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) + backend = _resolve_backend(backend, v) + return sort!(_copy(backend, v); backend, kwargs...) end """ sortperm!( ix::AbstractArray, - v::AbstractArray, - backend::Backend=get_backend(v); - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, + v::AbstractArray; + backend=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=:, + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # Algorithm choice - alg::Union{Nothing, SortAlgorithm}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - - # GPU settings - block_size::Union{Nothing, Int}=nothing, - - # Temporary buffer, same size as `v` temp::Union{Nothing, AbstractArray}=nothing, - ) + ) -> ix -Save into `ix` the index permutation of `v` such that `v[ix]` is sorted. The `lt`, `by`, `rev`, and -`order` arguments are the same as for `Base.sortperm`. The same algorithms are used as for -[`sort!`](@ref) with custom by-index comparators. +Write into `ix` the stable permutation that sorts `v`, so that `v[ix]` is sorted; `ix` is always +overwritten. The keywords are those of [`sort!`](@ref). -By default (`dims=:`) the whole array is permuted as one vector. Pass an integer `dims` to permute -each 1D slice along that dimension independently, like `Base.sortperm!(ix, A; dims)`: `ix` must -then have the same axes as `v` and receives linear indices into `v`, so that `v[ix]` is sorted -along `dims`. The permutation is stable in both cases. +With `dims=:`, `ix` needs as many elements as `v` and receives indices `1:length(v)`. With an +integer `dims`, `ix` must have the same axes as `v` and receives linear indices into `v`, so that +`v[ix]` is sorted along `dims`, like `Base.sortperm!(ix, A; dims)`. -## Algorithm choice -By default, `sortperm!` uses sample sort on CPU backends and merge sort on GPU -backends. Pass `alg=MergeSort(lowmem=true)` to use the lower-memory GPU permutation path. -`RadixSort()` and `BitonicSort()` do not provide a permutation path. +`Auto()` chooses [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort) on the +host and [`MergeSort`](@ref) on GPUs; `MergeSort(lowmem=true)` avoids copying the keys. +`RadixSort` and `BitonicSort` have no permutation path. `backend` is derived from `ix` and `v`. +`temp` is an optional scratch array like `ix`. """ function sortperm!( ix::AbstractArray, - v::AbstractArray, - backend::Backend=get_backend(v); - kwargs... -) - _sortperm_impl!( - ix, v, backend; - kwargs... - ) -end - - -function _sortperm_impl!( - ix::AbstractArray, - v::AbstractArray, - backend::Backend; - + v::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=Colon(), lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Forward, - - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, - - alg::Union{Nothing, SortAlgorithm}=nothing, - - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), - - # GPU settings; nothing => merge sort's tuned default (sortperm is merge-only) - block_size::Union{Nothing, Int}=nothing, - - # Temporary buffer, same size as `v` + order::Base.Order.Ordering=Base.Order.Forward, temp::Union{Nothing, AbstractArray}=nothing, ) - if !(dims isa Colon) - 1 <= dims <= ndims(v) || throw(ArgumentError("dimension out of range")) - axes(ix) == axes(v) || - throw(ArgumentError("index array must have the same axes as the input, $(axes(ix)) != $(axes(v))")) + backend = _resolve_backend(backend, ix, v, temp) + ord = Base.Order.ord(lt, by, rev, order) + if dims isa Colon + length(ix) == length(v) || throw(ArgumentError( + "index array must have as many elements as the input, $(length(ix)) != $(length(v))")) + else + axes(ix) == axes(v) || throw(ArgumentError( + "index array must have the same axes as the input, $(axes(ix)) != $(axes(v))")) end + a = _resolve_sort(alg, backend, v, dims, ord; perm=true) + _sortperm_impl!(a, ix, v, backend, dims, ord; lt, by, rev, order, temp) + return ix +end - if use_gpu_algorithm(backend, prefer_threads) - alg = isnothing(alg) ? MergeSort() : alg - bs = isnothing(block_size) ? 256 : block_size - if alg isa MergeSort - if alg.lowmem - merge_sortperm_lowmem!( - ix, v, backend; - lt, by, rev, order, - block_size=bs, - temp, dims, - ) - else - # merge_sortperm! copies keys alongside indices in shared memory so comparisons - # never touch global memory during the binary-search step. - # merge_sortperm_lowmem! avoids the key copy but its comparator does two global - # loads per comparison, making it O(n log²n) in global traffic at large n. - merge_sortperm!( - ix, v, backend; - lt, by, rev, order, - block_size=bs, - temp_ix=temp, # old `temp` was the index buffer; maps directly to temp_ix - dims, - ) - end - elseif alg isa RadixSort - throw(ArgumentError("RadixSort does not support sortperm")) - elseif alg isa BitonicSort - throw(ArgumentError("BitonicSort does not support sortperm")) - else - throw(ArgumentError("$(typeof(alg)) is not supported by sortperm! on GPU backends")) - end +function _sortperm_impl!(a::MergeSort, ix, v, backend, dims, ord; lt, by, rev, order, temp) + if a.lowmem + _merge_sortperm_lowmem!(ix, v, backend; lt, by, rev, order, + block_size=a.block_size, temp, dims) else - alg = isnothing(alg) ? SampleSort() : alg - if !(alg isa SampleSort) - throw(ArgumentError("$(typeof(alg)) is not supported by sortperm! on CPU backends")) - elseif dims isa Colon - sample_sortperm!( - ix, v; - lt, by, rev, order, - max_tasks, - min_elems, - temp, - ) - else - # Like Base: sort the linear indices of each slice by the values they point to - ord = Base.Order.Perm(Base.Order.ord(lt, by, rev, order), vec(v)) - copyto!(ix, LinearIndices(v)) - foreach_slice(ix, dims; max_tasks, min_elems) do slice - Base.sort!(slice; order=ord) - end - end + # Copies keys alongside indices, so comparisons never read global memory; the low-memory + # path does two global loads per comparison, O(n log²n) global traffic at large n. + _merge_sortperm!(ix, v, backend; lt, by, rev, order, + block_size=a.block_size, temp_ix=temp, dims) + end +end + +function _sortperm_impl!(a::CPUThreads.SampleSort, ix, v, backend, dims, ord; + lt, by, rev, order, temp) + if dims isa Colon + _sample_sortperm!(vec(ix), vec(v); lt, by, rev, order, + max_tasks=a.max_tasks, min_elems=a.min_elems, temp) + else + _sample_sortperm_dims!(ix, v, ord, dims; max_tasks=a.max_tasks, min_elems=a.min_elems) end - ix end """ - sortperm( - v::AbstractArray, - backend::Backend=get_backend(v); + sortperm(v::AbstractArray; kwargs...) + +Out-of-place [`sortperm!`](@ref): return an `Int` array shaped like `v` holding the permutation, +with the same keywords. +""" +function sortperm(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) + backend = _resolve_backend(backend, v) + return sortperm!(_similar(backend, v, Int), v; backend, kwargs...) +end - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, +""" + sort_by_key!( + keys::AbstractArray, + values::AbstractArray; + backend=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=:, + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + temp_keys::Union{Nothing, AbstractArray}=nothing, + temp_values::Union{Nothing, AbstractArray}=nothing, + ) -> (keys, values) - # Algorithm choice - alg::Union{Nothing, SortAlgorithm}=nothing, +Sort `keys` in place and apply the same permutation to `values`, stably: values with equal keys +keep their relative order. The ordering keywords and `dims` are those of [`sort!`](@ref); `values` +must have as many elements as `keys` (the same axes with an integer `dims`). - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), +`Auto()` chooses [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort) on the +host and [`MergeSort`](@ref) on GPUs, the algorithms that support key/value sorting. `backend` is +derived from `keys` and `values`. `temp_keys` and `temp_values` are optional scratch arrays like +`keys` and `values`. - # GPU settings - block_size::Union{Nothing, Int}=nothing, +Thrust, oneDPL and Kokkos call this operation `sort_by_key`, CUB `SortPairs`. - # Temporary buffer, same size as `v` - temp::Union{Nothing, AbstractArray}=nothing, - ) +# Examples +```julia +import AcceleratedKernels as AK +using Metal -Out-of-place sortperm, same settings as [`sortperm!`](@ref). +keys = MtlArray(rand(Int32(1):Int32(10), 1000)) +values = MtlArray(Int32.(1:1000)) +AK.sort_by_key!(keys, values) # values of equal keys stay in ascending order +``` """ -function sortperm( - v::AbstractArray, - backend::Backend=get_backend(v); - kwargs... +function sort_by_key!( + keys::AbstractArray, + values::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), + dims::Union{Colon, Integer}=Colon(), + lt=isless, + by=identity, + rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + temp_keys::Union{Nothing, AbstractArray}=nothing, + temp_values::Union{Nothing, AbstractArray}=nothing, ) - ix = similar(v, Int) - sortperm!( - ix, v, backend; - kwargs... - ) + backend = _resolve_backend(backend, keys, values, temp_keys, temp_values) + ord = Base.Order.ord(lt, by, rev, order) + if dims isa Colon + length(keys) == length(values) || throw(ArgumentError( + "keys and values must have the same length, $(length(keys)) != $(length(values))")) + else + axes(keys) == axes(values) || throw(ArgumentError( + "keys and values must have the same axes, $(axes(keys)) != $(axes(values))")) + end + a = _resolve_sort(alg, backend, keys, dims, ord; pairs=true) + _sort_by_key_impl!(a, keys, values, backend, dims, ord; + lt, by, rev, order, temp_keys, temp_values) + return keys, values +end + +function _sort_by_key_impl!(a::MergeSort, keys, values, backend, dims, ord; + lt, by, rev, order, temp_keys, temp_values) + _merge_sort_by_key!(keys, values, backend; lt, by, rev, order, dims, + block_size=a.block_size, temp_keys, temp_values) +end + +function _sort_by_key_impl!(a::CPUThreads.SampleSort, keys, values, backend, dims, ord; + lt, by, rev, order, temp_keys, temp_values) + _sample_sort_by_key!(keys, values, ord, dims; + max_tasks=a.max_tasks, min_elems=a.min_elems, temp_keys, temp_values) end diff --git a/src/sort/tuning.jl b/src/sort/tuning.jl new file mode 100644 index 00000000..aa5f58dd --- /dev/null +++ b/src/sort/tuning.jl @@ -0,0 +1,221 @@ +# Selection and tuning of sorting algorithms. +# +# Every sorting entry point resolves its `alg` once, before touching any data: +# 1. `_checkdomain` checks the fields the caller set explicitly; +# 2. `_select_sort` picks an algorithm for `Auto`, from the device's `SortTuning`, the +# capabilities of the backend and the call's static facts (element type, layout, ordering); +# 3. `_fill` fills the fields left at `nothing` from the tuning; +# 4. `_check` checks the complete algorithm against the backend's capabilities, the operation +# and the arguments, throwing an `ArgumentError` for anything it cannot run. +# Selection returns algorithms with unset fields, so `Auto` and explicit algorithms are filled and +# checked by the same code. + + +""" + SortTuning(; kwargs...) + +Values that drive [`Auto`](@ref) selection and fill unset algorithm fields for sorting on one +device, as returned by [`sort_tuning`](@ref). + +- `bitonic_max_len`: `Auto` picks `BitonicSort` for slices and arrays up to this length, where + its instability is allowed (see [`Auto`](@ref)). +- `radix_min_len`: `Auto` picks `RadixSort` for whole-array sorts from this length, for the + element types and orderings it supports. +- `merge_block_size`, `radix_block_size`, `radix_items_per_thread`, `bitonic_block_size`, + `bitonic_items_per_thread`: settings for the algorithms' unset fields. +- `threads_min_elems`: the default `min_elems` of `CPUThreads.SampleSort`; `max_tasks` defaults + to `Threads.nthreads()`. + +The defaults never pick `BitonicSort` or `RadixSort`, and reproduce AK's historical settings. +Internal: the fields may change in any release. +""" +Base.@kwdef struct SortTuning + bitonic_max_len::Int = 0 + radix_min_len::Int = typemax(Int) + merge_block_size::Int = 256 + radix_block_size::Int = 256 + radix_items_per_thread::Int = 2 + bitonic_block_size::Int = 256 + bitonic_items_per_thread::Int = 8 + threads_min_elems::Int = 1 +end + +""" + sort_tuning(backend, T) -> SortTuning + +The sorting tuning for element type `T` on `backend`'s current device (the device the calling +task would launch on). AK defines this generic method; AK's package extensions add one method +per backend type, which may choose different values per device. +""" +sort_tuning(::Backend, ::Type) = SortTuning() + + +# Algorithm names in error messages +_algname(a) = nameof(typeof(a)) +_algname(::CPUThreads.SampleSort) = "CPUThreads.SampleSort" + + +# The operation an algorithm is resolved for, for error messages +_sort_opname(perm, pairs) = pairs ? "sort_by_key!" : perm ? "sortperm!" : "sort!" + + +""" + _resolve_sort(alg, backend, v, dims, ord; perm=false, pairs=false) -> SortAlgorithm + +Resolve `alg` for sorting `v` (the keys) along `dims` under ordering `ord` on `backend`, for +`sort!` (`perm == pairs == false`), `sortperm!` (`perm`) or `sort_by_key!` (`pairs`). Returns a +concrete algorithm with every field set, or throws an `ArgumentError`. +""" +function _resolve_sort(alg::Algorithm, backend::Backend, v::AbstractArray, dims, + ord::Base.Order.Ordering; perm::Bool=false, pairs::Bool=false) + T = eltype(v) + layout = slice_layout(v, dims) + _checkdomain(alg) + t = sort_tuning(backend, T) + a = alg isa Auto ? _select_sort(alg, backend, t, T, layout, ord; perm, pairs) : alg + a = _fill(a, t, T) + _check(a, backend, T, layout, ord; perm, pairs) + return a +end + +function _select_sort(a::Auto, backend, t::SortTuning, ::Type{T}, layout, ord; + perm, pairs) where {T} + _runs_threads(backend) && return CPUThreads.SampleSort() + (perm || pairs) && return MergeSort() # the only stable, key/value-capable kernel algorithm + if (!a.stable || _ties_invisible(T, ord)) && layout.len <= t.bitonic_max_len + return BitonicSort() + elseif layout isa FlatLayout && _rs_supported(T) && _rs_ordering(ord) && + layout.len >= t.radix_min_len + return RadixSort() # whole arrays only; a single slice is not flat + end + return MergeSort() +end + +# Equal elements under `ord` are bitwise identical, so an unstable sort's result equals a +# stable sort's. Floats are excluded: `isless` treats NaNs with different bit patterns as equal. +# `Base.Order.ord(isless, identity, rev, Forward)` is one of the two orderings accepted here; +# custom `lt`/`by` give other ordering types. `isconcretetype` excludes union element types, +# whose equal values can differ in type. +_ties_invisible(::Type{T}, ord) where {T} = + isconcretetype(T) && T <: Union{Base.BitInteger, Bool, Char} && + ord isa Union{Base.Order.ForwardOrdering, Base.Order.ReverseOrdering{Base.Order.ForwardOrdering}} + +# The orderings radix sort supports: the default one and its reverse. +_rs_ordering(ord) = ord === Base.Order.Forward || ord === Base.Order.Reverse + + +# Domain checks of explicitly set fields, before any arithmetic uses them + +_checkdomain(::Algorithm) = nothing + +function _check_positive(a, field) + x = getfield(a, field) + x === nothing || x >= 1 || + throw(ArgumentError("$(_algname(a)): `$field` must be positive, got $x")) + nothing +end + +function _check_pow2(a, field) + x = getfield(a, field) + x === nothing || (x >= 1 && ispow2(x)) || + throw(ArgumentError("$(_algname(a)): `$field` must be a positive power of two, got $x")) + nothing +end + +_checkdomain(a::MergeSort) = _check_positive(a, :block_size) + +function _checkdomain(a::RadixSort) + # Bounded, so that the local-memory footprints computed from them cannot overflow + _check_pow2(a, :block_size) + _check_positive(a, :items_per_thread) + a.block_size === nothing || a.block_size <= 1024 || throw(ArgumentError( + "RadixSort: `block_size` must be at most 1024, got $(a.block_size)")) + a.items_per_thread === nothing || a.items_per_thread <= 64 || throw(ArgumentError( + "RadixSort: `items_per_thread` must be at most 64, got $(a.items_per_thread)")) + nothing +end + +function _checkdomain(a::BitonicSort) + _check_pow2(a, :block_size) + _check_pow2(a, :items_per_thread) + if a.block_size !== nothing && a.items_per_thread !== nothing + a.block_size <= typemax(Int) ÷ a.items_per_thread || throw(ArgumentError( + "BitonicSort: `block_size * items_per_thread` overflows " * + "($(a.block_size) * $(a.items_per_thread))")) + end + nothing +end + +function _checkdomain(a::CPUThreads.SampleSort) + _check_positive(a, :max_tasks) + _check_positive(a, :min_elems) +end + + +# Fill unset fields from the tuning + +_fill(a::MergeSort, t::SortTuning, T) = + MergeSort(something(a.block_size, t.merge_block_size), a.lowmem) +_fill(a::RadixSort, t::SortTuning, T) = + RadixSort(something(a.block_size, t.radix_block_size), + something(a.items_per_thread, t.radix_items_per_thread)) +_fill(a::BitonicSort, t::SortTuning, T) = + BitonicSort(something(a.block_size, t.bitonic_block_size), + something(a.items_per_thread, t.bitonic_items_per_thread)) +_fill(a::CPUThreads.SampleSort, t::SortTuning, T) = + CPUThreads.SampleSort(something(a.max_tasks, Threads.nthreads()), + something(a.min_elems, t.threads_min_elems)) +_fill(a::Algorithm, t::SortTuning, T) = + throw(ArgumentError("$(_algname(a)) is not a sorting algorithm")) + + +# Check a complete algorithm against the backend, the operation and the arguments + +function _require_kernels(a, backend) + _runs_kernels(backend) || throw(ArgumentError( + "$(_algname(a)) runs AcceleratedKernels' GPU kernels, which " * + "$(_backend_name(backend)) cannot run (on the host, this needs KernelAbstractions 0.10); " * + "use `Auto()` or a `CPUThreads` algorithm")) + nothing +end + +function _check(a::MergeSort, backend, T, layout, ord; perm, pairs) + _checkdomain(a) + _require_kernels(a, backend) + !a.lowmem || (perm && !pairs) || throw(ArgumentError( + "MergeSort(lowmem=true) is only supported by sortperm!, not by $(_sort_opname(perm, pairs))")) + nothing +end + +function _check(a::RadixSort, backend, ::Type{T}, layout, ord; perm, pairs) where {T} + _checkdomain(a) + _require_kernels(a, backend) + (perm || pairs) && throw(ArgumentError( + "RadixSort does not support $(_sort_opname(perm, pairs))")) + layout isa FlatLayout || throw(ArgumentError( + "RadixSort does not support sorting along `dims`")) + _rs_supported(T) || throw(ArgumentError( + "RadixSort does not support element type $T; it supports 32- and 64-bit integers " * + "and floats")) + _rs_ordering(ord) || throw(ArgumentError( + "RadixSort only supports the default ordering and its reverse (no custom `lt` or `by`)")) + _rs_portable_local_memory(T, a.block_size) <= LOCAL_MEMORY_BUDGET || throw(ArgumentError( + "RadixSort: block_size=$(a.block_size) needs more than $(LOCAL_MEMORY_BUDGET) bytes " * + "of local memory for element type $T")) + nothing +end + +function _check(a::BitonicSort, backend, T, layout, ord; perm, pairs) + _checkdomain(a) + _require_kernels(a, backend) + (perm || pairs) && throw(ArgumentError( + "BitonicSort is unstable and does not support $(_sort_opname(perm, pairs))")) + nothing +end + +function _check(a::CPUThreads.SampleSort, backend, T, layout, ord; perm, pairs) + _checkdomain(a) + _runs_threads(backend) || throw(ArgumentError( + "CPUThreads.SampleSort only runs on the host backend, not on $(_backend_name(backend))")) + nothing +end diff --git a/test/generic/sort.jl b/test/generic/sort.jl index de7bd002..28593fc0 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -1,4 +1,11 @@ -if !prefer_threads +# Tests that do not choose an algorithm use `SORT_ALG`: `Auto()`, except in the `--cpu-ka` +# configuration, whose point is to run AK's kernels on the host backend. `SETTINGS_ALG` is an +# explicitly tuned algorithm for the configuration. +SORT_ALG = HOST_KERNELS ? AK.MergeSort() : AK.Auto() +SETTINGS_ALG = TEST_KERNELS ? AK.MergeSort(block_size=64) : + AK.CPUThreads.SampleSort(max_tasks=64, min_elems=8) + +if TEST_KERNELS @testset "merge_sort" begin Random.seed!(0) @@ -6,7 +13,7 @@ if !prefer_threads for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(Int32, num_elems)) - AK.merge_sort!(v) + AK.sort!(v; alg=AK.MergeSort()) vh = Array(v) @test issorted(vh) end @@ -14,7 +21,7 @@ if !prefer_threads for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(UInt32, num_elems)) - AK.merge_sort!(v) + AK.sort!(v; alg=AK.MergeSort()) vh = Array(v) @test issorted(vh) end @@ -22,35 +29,35 @@ if !prefer_threads for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - AK.merge_sort!(v) + AK.sort!(v; alg=AK.MergeSort()) vh = Array(v) @test issorted(vh) end # Testing different settings v = array_from_host(1:10_000, Float32) - AK.merge_sort!(v, lt=(>), by=abs, rev=true, - block_size=64, temp=array_from_host(1:10_000, Float32)) + AK.sort!(v, lt=(>), by=abs, rev=true, + alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Float32)) @test issorted(Array(v)) v = array_from_host(1:10_000, Int32) - AK.merge_sort!(v, lt=(>), rev=true, - block_size=64, temp=array_from_host(1:10_000, Int32)) + AK.sort!(v, lt=(>), rev=true, + alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Int32)) @test issorted(Array(v)) v = array_from_host(1:10_000, Float32) - v = AK.merge_sort(v, lt=(>), by=abs, rev=true, - block_size=64, temp=array_from_host(1:10_000, Float32)) + v = AK.sort(v, lt=(>), by=abs, rev=true, + alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Float32)) @test issorted(Array(v)) v = array_from_host(1:10_000, Int32) - v = AK.merge_sort(v, lt=(>), by=abs, rev=true, - block_size=64, temp=array_from_host(1:10_000, Int32)) + v = AK.sort(v, lt=(>), by=abs, rev=true, + alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Int32)) @test issorted(Array(v)) end @testset "sort_by_transform" begin - # Tests for the by= hoisting optimisation: by(elem) is broadcast once before + # Tests for the by= hoisting optimisation: by(elem) is mapped once before # sorting rather than being called inside every merge comparison. # Checks exact output match against Base.sort so we catch ordering regressions. Random.seed!(42) @@ -66,7 +73,7 @@ end ) v = array_from_host(v_h) tmp = copy(v) - AK.merge_sort!(tmp; kw...) + AK.sort!(tmp; alg=AK.MergeSort(), kw...) @test Array(tmp) == sort(v_h; base_kw...) end end @@ -75,7 +82,7 @@ end n = 10_000 v_h = randn(Float32, n) v = array_from_host(v_h); tmp = copy(v) - AK.merge_sort!(tmp; rev=true) + AK.sort!(tmp; alg=AK.MergeSort(), rev=true) @test Array(tmp) == sort(v_h; rev=true) # Edge sizes under by= hoisting @@ -83,7 +90,7 @@ end v_h = randn(Float32, n) v = array_from_host(v_h) tmp = copy(v) - AK.merge_sort!(tmp; by=abs) + AK.sort!(tmp; alg=AK.MergeSort(), by=abs) @test Array(tmp) == sort(v_h; by=abs) end @@ -93,7 +100,7 @@ end v = array_from_host(v_h) tmp = copy(v) temp = array_from_host(zeros(Float32, n)) - AK.merge_sort!(tmp; by=abs, temp) + AK.sort!(tmp; alg=AK.MergeSort(), by=abs, temp) @test Array(tmp) == sort(v_h; by=abs) # sort! (public API) routes through the same hoisting path @@ -109,27 +116,54 @@ end v_h = randn(Float32, n) v = array_from_host(v_h) tmp = copy(v) - AK.merge_sort!(tmp; by=x->x>0) + AK.sort!(tmp; alg=AK.MergeSort(), by=x->x>0) @test Array(tmp) == sort(v_h; by=x->x>0) + # Keys keep their own type: transforms to wider floats, fractions and tuples must not be + # converted back to the element type + if KernelAbstractions.supports_float64(BACKEND) # the keys are Float64 + v_h = Float32[2, 1, 3] + by_wide = x -> 1.0 + Float64(x) * eps(Float64) + @test Array(AK.sort!(array_from_host(v_h); alg=AK.MergeSort(), by=by_wide)) == sort(v_h; by=by_wide) + v_h = Int32[3, 1, 2, 4] + @test Array(AK.sort!(array_from_host(v_h); alg=AK.MergeSort(), by=x -> x / 2)) == sort(v_h; by=x -> x / 2) + end + v_h = Int32[3, 1, 2, 4] + @test Array(AK.sort!(array_from_host(v_h); alg=AK.MergeSort(), by=x -> (x % 2, x))) == + sort(v_h; by=x -> (x % 2, x)) + # identity path unchanged: verify no regression from the early-return guard n = 10_000 v_h = rand(Float32, n) v = array_from_host(v_h) tmp = copy(v) - AK.merge_sort!(tmp) + AK.sort!(tmp; alg=AK.MergeSort()) @test Array(tmp) == sort(v_h) end -else # CPU backend +end + +if AK._runs_threads(BACKEND) @testset "sample_sort" begin Random.seed!(0) + # Stable, also when several tasks sort buckets in parallel (at least 16 elements per task): + # elements tagged with their position, sorted by a key with many ties, keep their order + th = [(rand(Int32(1):Int32(50)), Int32(i)) for i in 1:100_000] + for max_tasks in (1, 4, 16) + t = copy(th) + AK.sort!(t; by=first, alg=AK.CPUThreads.SampleSort(; max_tasks)) + @test t == sort(th; by=first) + ix = zeros(Int, length(th)) + AK.sortperm!(ix, th; by=first, alg=AK.CPUThreads.SampleSort(; max_tasks)) + @test ix == sortperm(th; by=first) + end + # Fuzzy correctness testing for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(Int32, num_elems)) - AK.sample_sort!(v) + AK.sort!(v; alg=AK.CPUThreads.SampleSort()) vh = Array(v) @test issorted(vh) end @@ -137,7 +171,7 @@ else # CPU backend for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(UInt32, num_elems)) - AK.sample_sort!(v) + AK.sort!(v; alg=AK.CPUThreads.SampleSort()) vh = Array(v) @test issorted(vh) end @@ -145,20 +179,20 @@ else # CPU backend for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - AK.sample_sort!(v) + AK.sort!(v; alg=AK.CPUThreads.SampleSort()) vh = Array(v) @test issorted(vh) end # Testing different settings v = array_from_host(rand(1:100_000, 10_000), Float32) - AK.sample_sort!(v, lt=(>), by=abs, rev=true, - max_tasks=64, temp=array_from_host(1:10_000, Float32)) + AK.sort!(v, lt=(>), by=abs, rev=true, + alg=AK.CPUThreads.SampleSort(max_tasks=64), temp=array_from_host(1:10_000, Float32)) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - AK.sample_sort!(v, lt=(>), rev=true, - max_tasks=64, temp=array_from_host(1:10_000, Int32)) + AK.sort!(v, lt=(>), rev=true, + alg=AK.CPUThreads.SampleSort(max_tasks=64), temp=array_from_host(1:10_000, Int32)) @test issorted(Array(v)) end end @@ -171,7 +205,7 @@ end for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Int32, num_elems)) - AK.sort!(v; prefer_threads) + AK.sort!(v; alg=SORT_ALG) vh = Array(v) @test issorted(vh) end @@ -179,7 +213,7 @@ end for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(UInt32, num_elems)) - AK.sort!(v; prefer_threads) + AK.sort!(v; alg=SORT_ALG) vh = Array(v) @test issorted(vh) end @@ -187,33 +221,29 @@ end for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - AK.sort!(v; prefer_threads) + AK.sort!(v; alg=SORT_ALG) vh = Array(v) @test issorted(vh) end # Testing different settings v = array_from_host(rand(1:100_000, 10_000), Float32) - AK.sort!(v; prefer_threads, lt=(>), by=abs, rev=true, - max_tasks=64, min_elems=8, block_size=64, + AK.sort!(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, temp=array_from_host(1:10_000, Float32)) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - AK.sort!(v; prefer_threads, lt=(>), rev=true, - max_tasks=64, min_elems=8, block_size=64, + AK.sort!(v; alg=SETTINGS_ALG, lt=(>), rev=true, temp=array_from_host(1:10_000, Int32)) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Float32) - v = AK.sort(v; prefer_threads, lt=(>), by=abs, rev=true, - max_tasks=64, min_elems=8, block_size=64, + v = AK.sort(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, temp=array_from_host(1:10_000, Float32)) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - v = AK.sort(v; prefer_threads, lt=(>), by=abs, rev=true, - max_tasks=64, min_elems=8, block_size=64, + v = AK.sort(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, temp=array_from_host(1:10_000, Int32)) @test issorted(Array(v)) end @@ -229,20 +259,20 @@ end issorted(vh[ixh]; kwargs...) end - if !prefer_threads + if TEST_KERNELS for T in valid_backend_eltypes(BACKEND, (UInt32, Int32, Float32, UInt64, Int64, Float64)) v_h = rand(T, 10_000) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test Array(v) == sort(v_h) end v_h = rand(Int32, 10_000) v_default = array_from_host(v_h) v_merge = array_from_host(v_h) - AK.sort!(v_default; prefer_threads) - AK.sort!(v_merge; prefer_threads, alg=AK.MergeSort()) + AK.sort!(v_default; alg=SORT_ALG) + AK.sort!(v_merge; alg=AK.MergeSort()) @test Array(v_merge) == Array(v_default) perm_h = rand(Float32, 4096) @@ -250,35 +280,45 @@ end v = array_from_host(perm_h) ix = array_from_host(zeros(Int, length(perm_h))) temp = array_from_host(zeros(Int, length(perm_h))) - AK.sortperm!(ix, v; prefer_threads, alg, temp) + AK.sortperm!(ix, v; alg, temp) @test is_valid_perm(perm_h, Int.(Array(ix))) end v = array_from_host(rand(Float32, 128)) ix = array_from_host(zeros(Int, length(v))) - @test_throws ArgumentError AK.sort!(copy(v); prefer_threads, alg=AK.SampleSort()) - @test_throws ArgumentError AK.sortperm!(ix, v; prefer_threads, alg=AK.RadixSort()) - else + @test_throws ArgumentError AK.sortperm!(ix, v; alg=AK.RadixSort()) + @test_throws ArgumentError AK.sortperm!(ix, v; alg=AK.BitonicSort()) + @test_throws ArgumentError AK.sort!(copy(v); alg=AK.MergeSort(lowmem=true)) + end + + if AK._runs_threads(BACKEND) v_h = rand(Int32, 10_000) v_default = array_from_host(v_h) v_sample = array_from_host(v_h) - AK.sort!(v_default; prefer_threads) - AK.sort!(v_sample; prefer_threads, alg=AK.SampleSort()) + AK.sort!(v_default) + AK.sort!(v_sample; alg=AK.CPUThreads.SampleSort()) @test Array(v_sample) == Array(v_default) ix = array_from_host(zeros(Int, length(v_h))) - AK.sortperm!(ix, array_from_host(v_h); prefer_threads, alg=AK.SampleSort()) + AK.sortperm!(ix, array_from_host(v_h); alg=AK.CPUThreads.SampleSort()) @test is_valid_perm(v_h, Int.(Array(ix))) - @test_throws ArgumentError AK.sort!(array_from_host(v_h); prefer_threads, alg=AK.MergeSort()) - @test_throws ArgumentError AK.sort!(array_from_host(v_h); prefer_threads, alg=AK.RadixSort()) - @test_throws ArgumentError AK.sortperm!(ix, array_from_host(v_h); prefer_threads, alg=AK.RadixSort()) + @test_throws ArgumentError AK.sort!(array_from_host(v_h); alg=AK.CPUThreads.SampleSort(max_tasks=0)) + @test_throws ArgumentError AK.sortperm!(ix, array_from_host(v_h); alg=AK.RadixSort()) + else + @test_throws ArgumentError AK.sort!(array_from_host(rand(Int32, 16)); alg=AK.CPUThreads.SampleSort()) + end + + # AK's kernels need a backend that runs them (on the host, KernelAbstractions 0.10) + if !AK._runs_kernels(BACKEND) + for alg in (AK.MergeSort(), AK.RadixSort(), AK.BitonicSort()) + @test_throws ArgumentError AK.sort!(array_from_host(rand(Int32, 16)); alg) + end end end -if !prefer_threads -@testset "merge_sort_by_key" begin +@testset "sort_by_key" begin Random.seed!(0) # Fuzzy correctness testing @@ -286,7 +326,7 @@ if !prefer_threads num_elems = rand(1:100_000) k = array_from_host(rand(Int32, num_elems)) v = copy(k) .- 1 - AK.merge_sort_by_key!(k, v) + AK.sort_by_key!(k, v; alg=SORT_ALG) kh = Array(k) vh = Array(v) @test issorted(kh) @@ -297,7 +337,7 @@ if !prefer_threads num_elems = rand(1:100_000) k = array_from_host(rand(UInt32, num_elems)) v = copy(k) .- 1 - AK.merge_sort_by_key!(k, v) + AK.sort_by_key!(k, v; alg=SORT_ALG) kh = Array(k) vh = Array(v) @test issorted(kh) @@ -308,7 +348,7 @@ if !prefer_threads num_elems = rand(1:100_000) k = array_from_host(rand(Float32, num_elems)) v = copy(k) .- 1 - AK.merge_sort_by_key!(k, v) + AK.sort_by_key!(k, v; alg=SORT_ALG) kh = Array(k) vh = Array(v) @test issorted(kh) @@ -318,9 +358,9 @@ if !prefer_threads # Testing different settings k = array_from_host(1:10_000, Float32) v = array_from_host(1:10_000, Int32) - AK.merge_sort_by_key!(k, v, + AK.sort_by_key!(k, v, lt=(>), by=abs, rev=true, - block_size=64, + alg=SETTINGS_ALG, temp_keys=array_from_host(1:10_000, Float32), temp_values=array_from_host(1:10_000, Int32)) @test issorted(Array(k)) @@ -328,38 +368,65 @@ if !prefer_threads k = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.merge_sort_by_key!(k, v, + AK.sort_by_key!(k, v, lt=(>), by=abs, rev=true, - block_size=64, + alg=SETTINGS_ALG, temp_keys=array_from_host(1:10_000, Int32), temp_values=array_from_host(1:10_000, Float32)) @test issorted(Array(k)) @test issorted(Array(v)) - k = array_from_host(1:10_000, Float32) - v = array_from_host(1:10_000, Int32) - AK.merge_sort_by_key(k, v, - lt=(>), by=abs, rev=true, - block_size=64, - temp_keys=array_from_host(1:10_000, Float32), - temp_values=array_from_host(1:10_000, Int32)) - @test issorted(Array(k)) - @test issorted(Array(v)) + # Stable: the values of equal keys keep their order, whole arrays and along `dims`, including + # a sample sort that runs several tasks + kh = rand(Int32(1):Int32(20), 50_000) + k = array_from_host(kh) + v = array_from_host(Int32.(1:50_000)) + AK.sort_by_key!(k, v; alg=SORT_ALG) + @test Array(k) == sort(kh) + @test Array(v) == sortperm(kh) + if AK._runs_threads(BACKEND) + k = array_from_host(kh) + v = array_from_host(Int32.(1:50_000)) + AK.sort_by_key!(k, v; alg=AK.CPUThreads.SampleSort(max_tasks=4)) + @test Array(v) == sortperm(kh) + end + Kh = rand(Int32(1):Int32(5), 300, 7) + for dims in (1, 2) + k = array_from_host(Kh) + v = array_from_host(reshape(Int32.(1:length(Kh)), size(Kh))) + AK.sort_by_key!(k, v; alg=SORT_ALG, dims) + @test Array(k) == sort(Kh; dims) + @test Array(v) == sortperm(Kh; dims) + end - k = array_from_host(1:10_000, Int32) - v = array_from_host(1:10_000, Float32) - AK.merge_sort_by_key(k, v, - lt=(>), by=abs, rev=true, - block_size=64, - temp_keys=array_from_host(1:10_000, Int32), - temp_values=array_from_host(1:10_000, Float32)) - @test issorted(Array(k)) - @test issorted(Array(v)) -end + # Arrays of any dimension, sorted as one flat vector, small and large + for sz in ((2, 2), (3, 4, 5), (300, 70)) + Kh = rand(Int32(1):Int32(5), sz) + k = array_from_host(Kh) + v = array_from_host(reshape(Int32.(1:length(Kh)), sz)) + AK.sort_by_key!(k, v; alg=SORT_ALG) + @test vec(Array(k)) == sort(vec(Kh)) + @test vec(Array(v)) == sortperm(vec(Kh)) + A = array_from_host(Kh) + @test vec(Array(AK.sort!(A; alg=SORT_ALG))) == sort(vec(Kh)) + ix = array_from_host(zeros(Int, sz)) + @test vec(Array(AK.sortperm!(ix, array_from_host(Kh); alg=SORT_ALG))) == sortperm(vec(Kh)) + end + + # Invalid scratch buffers are rejected before anything is modified + kh = Int32[2, 1] + k = array_from_host(kh) + v = array_from_host(Int32[20, 10]) + @test_throws ArgumentError AK.sort_by_key!(k, v; alg=SORT_ALG, temp_values=array_from_host(zeros(Int32, 1))) + @test Array(k) == kh + + # Mismatched sizes + @test_throws ArgumentError AK.sort_by_key!(array_from_host(rand(Int32, 10)), + array_from_host(rand(Int32, 9)); alg=SORT_ALG) end -if !prefer_threads +if TEST_KERNELS @testset "merge_sortperm" begin Random.seed!(0) @@ -368,7 +435,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Int32, num_elems)) - AK.merge_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -378,7 +445,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(UInt32, num_elems)) - AK.merge_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -388,7 +455,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Float32, num_elems)) - AK.merge_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -397,29 +464,30 @@ if !prefer_threads # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.merge_sortperm!(ix, - v, - lt=(>), by=abs, rev=true, - inplace=true, block_size=64, - temp_ix=array_from_host(1:10_000, Int32), - temp_v=array_from_host(1:10_000, Float32)) + AK.sortperm!(ix, + v, + lt=(>), by=abs, rev=true, + alg=AK.MergeSort(block_size=64), + temp=array_from_host(1:10_000, Int32)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) v = array_from_host(1:10_000, Float32) - ix = AK.merge_sortperm(v, - lt=(>), by=abs, rev=true, - inplace=true, block_size=64, - temp_ix=array_from_host(1:10_000, Int), - temp_v=array_from_host(1:10_000, Float32)) + ix = AK.sortperm(v, + lt=(>), by=abs, rev=true, + alg=AK.MergeSort(block_size=64), + temp=array_from_host(1:10_000, Int)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) end -else # CPU backend - @testset "sample_sortperm" begin +end + + +if AK._runs_threads(BACKEND) +@testset "sample_sortperm" begin Random.seed!(0) # Fuzzy correctness testing @@ -427,7 +495,7 @@ else # CPU backend num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Int32, num_elems)) - AK.sample_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.CPUThreads.SampleSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -437,7 +505,7 @@ else # CPU backend num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(UInt32, num_elems)) - AK.sample_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.CPUThreads.SampleSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -447,7 +515,7 @@ else # CPU backend num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Float32, num_elems)) - AK.sample_sortperm!(ix, v) + AK.sortperm!(ix, v; alg=AK.CPUThreads.SampleSort()) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -456,10 +524,10 @@ else # CPU backend # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sample_sortperm!(ix, - v, - lt=(>), by=abs, rev=true, - max_tasks=64, + AK.sortperm!(ix, + v, + lt=(>), by=abs, rev=true, + alg=AK.CPUThreads.SampleSort(max_tasks=64), temp=array_from_host(1:10_000, Int32)) ixh = Array(ix) vh = Array(v) @@ -468,7 +536,7 @@ end end -if !prefer_threads +if TEST_KERNELS @testset "merge_sortperm_lowmem" begin Random.seed!(0) @@ -477,7 +545,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Int32, num_elems)) - AK.merge_sortperm_lowmem!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort(lowmem=true)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -487,7 +555,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(UInt32, num_elems)) - AK.merge_sortperm_lowmem!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort(lowmem=true)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -497,7 +565,7 @@ if !prefer_threads num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Float32, num_elems)) - AK.merge_sortperm_lowmem!(ix, v) + AK.sortperm!(ix, v; alg=AK.MergeSort(lowmem=true)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -506,19 +574,19 @@ if !prefer_threads # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.merge_sortperm_lowmem!(ix, - v, - lt=(>), by=abs, rev=true, - block_size=64, + AK.sortperm!(ix, + v, + lt=(>), by=abs, rev=true, + alg=AK.MergeSort(lowmem=true, block_size=64), temp=array_from_host(1:10_000, Int32)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) v = array_from_host(1:10_000, Float32) - ix = AK.merge_sortperm_lowmem(v, - lt=(>), by=abs, rev=true, - block_size=64, + ix = AK.sortperm(v, + lt=(>), by=abs, rev=true, + alg=AK.MergeSort(lowmem=true, block_size=64), temp=array_from_host(1:10_000, Int)) ixh = Array(ix) vh = Array(v) @@ -535,7 +603,7 @@ end num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Int32, num_elems)) - AK.sortperm!(ix, v; prefer_threads) + AK.sortperm!(ix, v; alg=SORT_ALG) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -545,7 +613,7 @@ end num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(UInt32, num_elems)) - AK.sortperm!(ix, v; prefer_threads) + AK.sortperm!(ix, v; alg=SORT_ALG) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -555,7 +623,7 @@ end num_elems = rand(1:100_000) ix = array_from_host(zeros(Int32, num_elems)) v = array_from_host(rand(Float32, num_elems)) - AK.sortperm!(ix, v; prefer_threads) + AK.sortperm!(ix, v; alg=SORT_ALG) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -566,9 +634,8 @@ end v = array_from_host(1:10_000, Float32) AK.sortperm!(ix, v; - prefer_threads, + alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, - block_size=64, temp=array_from_host(1:10_000, Int32)) ixh = Array(ix) vh = Array(v) @@ -576,9 +643,8 @@ end v = array_from_host(1:10_000, Float32) ix = AK.sortperm(v; - prefer_threads, + alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, - block_size=64, temp=array_from_host(1:10_000, Int)) ixh = Array(ix) vh = Array(v) @@ -586,7 +652,7 @@ end end -if !prefer_threads +if TEST_KERNELS @testset "sortperm_extended" begin # Helper: ix is a valid permutation of 1:n that produces a sorted order function is_valid_perm(vh, ixh; kwargs...) @@ -604,7 +670,7 @@ if !prefer_threads n = rand(1:50_000) v = array_from_host(rand(T, n)) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v) + AK.sortperm!(ix, v; alg=SORT_ALG) vh, ixh = Array(v), Array(ix) @test is_valid_perm(vh, ixh) end @@ -614,7 +680,7 @@ if !prefer_threads for n in (1, 2, 3, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049) v = array_from_host(rand(Float32, n)) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v) + AK.sortperm!(ix, v; alg=SORT_ALG) vh, ixh = Array(v), Array(ix) @test is_valid_perm(vh, ixh) end @@ -632,7 +698,7 @@ if !prefer_threads ) v = array_from_host(arr) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v) + AK.sortperm!(ix, v; alg=SORT_ALG) vh, ixh = Array(v), Array(ix) @test is_valid_perm(vh, ixh) end @@ -649,7 +715,7 @@ if !prefer_threads ) v = array_from_host(randn(Float32, n)) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v; kw...) + AK.sortperm!(ix, v; alg=SORT_ALG, kw...) vh, ixh = Array(v), Array(ix) res = is_valid_perm(vh, ixh; kw...) @test res @@ -663,8 +729,8 @@ if !prefer_threads ix1 = array_from_host(zeros(Int, n)) ix2 = array_from_host(zeros(Int, n)) temp = array_from_host(zeros(Int, n)) - AK.sortperm!(ix1, v1; temp) - AK.sortperm!(ix2, v2; temp) + AK.sortperm!(ix1, v1; alg=SORT_ALG, temp) + AK.sortperm!(ix2, v2; alg=SORT_ALG, temp) @test Array(ix1) == Array(ix2) # ── Exact match against Base.sortperm ──────────────────────────────────── @@ -674,7 +740,7 @@ if !prefer_threads ref = sortperm(v_h) v = array_from_host(v_h) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v) + AK.sortperm!(ix, v; alg=SORT_ALG) ixh = Int.(Array(ix)) @test v_h[ixh] == v_h[ref] end @@ -684,7 +750,7 @@ if !prefer_threads v_h = Int32.(mod.(1:n, 10)) # values 0..9 cycling, 1000 of each v = array_from_host(v_h) ix = array_from_host(zeros(Int, n)) - AK.sortperm!(ix, v) + AK.sortperm!(ix, v; alg=SORT_ALG) ixh = Array(ix) for k in 0:9 group = ixh[v_h[ixh] .== k] @@ -694,12 +760,12 @@ if !prefer_threads # ── sortperm does not mutate the input ─────────────────────────────────── v = array_from_host(rand(Float32, 5_000)) vbak = copy(v) - AK.sortperm(v) + AK.sortperm(v; alg=SORT_ALG) @test Array(v) == Array(vbak) end @testset "radix_sort_alg" begin - if !prefer_threads + if TEST_KERNELS Random.seed!(0) # ── Correctness: fuzzy testing across supported types ───────────────── @@ -707,7 +773,7 @@ end for _ in 1:200 n = rand(1:100_000) v = array_from_host(rand(T, n)) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test issorted(Array(v)) end end @@ -717,7 +783,7 @@ end for _ in 1:200 n = rand(1:100_000) v = array_from_host(rand(T, n)) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test issorted(Array(v)) end end @@ -728,7 +794,7 @@ end n = 10_000 v_h = rand(T, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test Array(v) == sort(v_h) end @@ -738,7 +804,7 @@ end n = 10_000 v_h = rand(T, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort(), rev=true) + AK.sort!(v; alg=AK.RadixSort(), rev=true) @test Array(v) == sort(v_h; rev=true) end @@ -750,11 +816,11 @@ end v_h = shuffle!(vcat(specials, randn(T, 10_000))) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test isequal(Array(v), sort(v_h)) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort(), rev=true) + AK.sort!(v; alg=AK.RadixSort(), rev=true) @test isequal(Array(v), sort(v_h; rev=true)) end @@ -763,7 +829,7 @@ end for (rev, order) in ((nothing, Base.Order.Reverse), (true, Base.Order.Forward), (true, Base.Order.Reverse)) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort(), rev, order) + AK.sort!(v; alg=AK.RadixSort(), rev, order) @test Array(v) == sort(v_h; rev, order) end @@ -771,27 +837,27 @@ end n = 10_000 v_h = Int32.(mod.(1:n, 100)) # 100 distinct values, 100 copies each v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + AK.sort!(v; alg=AK.RadixSort()) @test Array(v) == sort(v_h) # ── Edge cases ──────────────────────────────────────────────────────── - @test length(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg=AK.RadixSort()))) == 0 - @test Array(AK.sort!(array_from_host(Int32[42]); prefer_threads, alg=AK.RadixSort())) == Int32[42] - @test Array(AK.sort!(array_from_host(Int32[2, 1]); prefer_threads, alg=AK.RadixSort())) == Int32[1, 2] + @test length(Array(AK.sort!(array_from_host(Int32[]); alg=AK.RadixSort()))) == 0 + @test Array(AK.sort!(array_from_host(Int32[42]); alg=AK.RadixSort())) == Int32[42] + @test Array(AK.sort!(array_from_host(Int32[2, 1]); alg=AK.RadixSort())) == Int32[1, 2] # ── temp kwarg: preallocated buffer ─────────────────────────────────── n = 50_000 v_h = rand(Float32, n) v = array_from_host(v_h) temp = similar(v) - AK.sort!(v; prefer_threads, alg=AK.RadixSort(), temp) + AK.sort!(v; alg=AK.RadixSort(), temp) @test Array(v) == sort(v_h) # ── Out-of-place ────────────────────────────────────────────────────── n = 10_000 v_h = rand(Float32, n) v = array_from_host(v_h) - w = AK.sort(v; prefer_threads, alg=AK.RadixSort()) + w = AK.sort(v; alg=AK.RadixSort()) @test Array(w) == sort(v_h) @test Array(v) == v_h # input unchanged @@ -799,15 +865,14 @@ end v_h = rand(UInt32, 20_000) for block_size in (128, 256, 512), items_per_thread in (1, 2, 4) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, - alg=AK.RadixSort(; block_size, items_per_thread)) + AK.sort!(v; alg=AK.RadixSort(; block_size, items_per_thread)) @test Array(v) == sort(v_h) end for n in (255, 256, 257) v_h = rand(UInt32, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.RadixSort(block_size=128)) + AK.sort!(v; alg=AK.RadixSort(block_size=128)) @test Array(v) == sort(v_h) end @@ -815,20 +880,20 @@ end n = 10 v_h = rand(Int32, n) v = array_from_host(v_h) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.RadixSort(), by=abs) + @test_throws ArgumentError AK.sort!(v; alg=AK.RadixSort(), by=abs) v_h = rand(Int32, n) v = array_from_host(v_h) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.RadixSort(), lt=(>)) + @test_throws ArgumentError AK.sort!(v; alg=AK.RadixSort(), lt=(>)) v_h = rand(Int32, n) v = array_from_host(v_h) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.RadixSort(), + @test_throws ArgumentError AK.sort!(v; alg=AK.RadixSort(), order=Base.Order.By(abs, Base.Order.Forward)) v_h = rand(Int16, n) v = array_from_host(v_h) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.RadixSort()) + @test_throws ArgumentError AK.sort!(v; alg=AK.RadixSort()) end end @@ -846,7 +911,7 @@ end A_h = rand(T, sz...) A = array_from_host(A_h) for dim in 1:nd, rev in (false, true) - @test Array(AK.sort(A; prefer_threads, dims=dim, rev)) == sort(A_h; dims=dim, rev) + @test Array(AK.sort(A; alg=SORT_ALG, dims=dim, rev)) == sort(A_h; dims=dim, rev) end end end @@ -855,58 +920,58 @@ end # block-level sort and the global merge passes are exercised, with duplicate-heavy data for len in (511, 512, 513, 1023, 1024, 1025, 2049, 10_000), nslices in (1, 3) A_h = rand(Int32(0):Int32(7), len, nslices) - @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=1)) == sort(A_h; dims=1) - @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=1, rev=true)) == sort(A_h; dims=1, rev=true) + @test Array(AK.sort(array_from_host(A_h); alg=SORT_ALG, dims=1)) == sort(A_h; dims=1) + @test Array(AK.sort(array_from_host(A_h); alg=SORT_ALG, dims=1, rev=true)) == sort(A_h; dims=1, rev=true) A_h = rand(Float32, nslices, len) - @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=2)) == sort(A_h; dims=2) + @test Array(AK.sort(array_from_host(A_h); alg=SORT_ALG, dims=2)) == sort(A_h; dims=2) end # by, lt, order and temp act on the values within each slice A_h = rand(Float32, 300, 700) A = array_from_host(A_h) - @test Array(AK.sort(A; prefer_threads, dims=1, by=x->-x)) == sort(A_h; dims=1, by=x->-x) - @test Array(AK.sort(A; prefer_threads, dims=2, lt=(>))) == sort(A_h; dims=2, lt=(>)) - @test Array(AK.sort(A; prefer_threads, dims=2, order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) - @test Array(AK.sort(A; prefer_threads, dims=2, temp=similar(A))) == sort(A_h; dims=2) - if !prefer_threads - @test Array(AK.sort(A; prefer_threads, dims=1, block_size=64)) == sort(A_h; dims=1) - @test_throws ArgumentError AK.sort(A; prefer_threads, dims=1, alg=AK.RadixSort()) + @test Array(AK.sort(A; alg=SORT_ALG, dims=1, by=x->-x)) == sort(A_h; dims=1, by=x->-x) + @test Array(AK.sort(A; alg=SORT_ALG, dims=2, lt=(>))) == sort(A_h; dims=2, lt=(>)) + @test Array(AK.sort(A; alg=SORT_ALG, dims=2, order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sort(A; alg=SORT_ALG, dims=2, temp=similar(A))) == sort(A_h; dims=2) + if TEST_KERNELS + @test Array(AK.sort(A; alg=AK.MergeSort(block_size=64), dims=1)) == sort(A_h; dims=1) + @test_throws ArgumentError AK.sort(A; dims=1, alg=AK.RadixSort()) end # NaNs, infinities and signed zeros order like Base A_h = Float32[NaN 1 -0.0; 0.0 -Inf NaN; 2 NaN Inf] A = array_from_host(A_h) - @test isequal(Array(AK.sort(A; prefer_threads, dims=1)), sort(A_h; dims=1)) - @test isequal(Array(AK.sort(A; prefer_threads, dims=2, rev=true)), sort(A_h; dims=2, rev=true)) + @test isequal(Array(AK.sort(A; alg=SORT_ALG, dims=1)), sort(A_h; dims=1)) + @test isequal(Array(AK.sort(A; alg=SORT_ALG, dims=2, rev=true)), sort(A_h; dims=2, rev=true)) # In-place sorts each slice, leaves the array otherwise intact A_h = rand(Int32, 40, 31) A = array_from_host(A_h) - AK.sort!(A; prefer_threads, dims=2) + AK.sort!(A; alg=SORT_ALG, dims=2) @test Array(A) == sort(A_h; dims=2) # dims=1 on a vector is a full sort v_h = rand(Int32, 5000) v = array_from_host(v_h) - @test Array(AK.sort(v; prefer_threads, dims=1)) == sort(v_h) + @test Array(AK.sort(v; alg=SORT_ALG, dims=1)) == sort(v_h) # 4D arrays A_h = rand(Int32(0):Int32(3), 3, 4, 5, 6) A = array_from_host(A_h) for dim in 1:4 - @test Array(AK.sort(A; prefer_threads, dims=dim)) == sort(A_h; dims=dim) + @test Array(AK.sort(A; alg=SORT_ALG, dims=dim)) == sort(A_h; dims=dim) end # Empty and singleton slices for sz in ((0, 5), (5, 0), (1, 64), (64, 1)), dim in 1:2 A_h = rand(Float32, sz...) - @test Array(AK.sort(array_from_host(A_h); prefer_threads, dims=dim)) == sort!(copy(A_h); dims=dim) + @test Array(AK.sort(array_from_host(A_h); alg=SORT_ALG, dims=dim)) == sort!(copy(A_h); dims=dim) end # Out-of-range dimension errors A = array_from_host(rand(Float32, 8, 8)) - @test_throws ArgumentError AK.sort(A; prefer_threads, dims=3) - @test_throws ArgumentError AK.sort(A; prefer_threads, dims=0) + @test_throws ArgumentError AK.sort(A; alg=SORT_ALG, dims=3) + @test_throws ArgumentError AK.sort(A; alg=SORT_ALG, dims=0) end @@ -922,7 +987,7 @@ end A_h = T <: Integer ? rand(T(0):T(4), sz...) : rand(T, sz...) A = array_from_host(A_h) for dim in 1:nd, rev in (false, true) - ix = Array(AK.sortperm(A; prefer_threads, dims=dim, rev)) + ix = Array(AK.sortperm(A; alg=SORT_ALG, dims=dim, rev)) @test ix == sortperm(A_h; dims=dim, rev) @test A_h[ix] == sort(A_h; dims=dim, rev) end @@ -932,39 +997,39 @@ end # Ties across block tiles and global merge passes must stay stable for len in (511, 512, 513, 1025, 2049, 10_000), nslices in (1, 3) A_h = rand(Int32(0):Int32(3), len, nslices) - @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=1)) == sortperm(A_h; dims=1) - @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=1, rev=true)) == sortperm(A_h; dims=1, rev=true) + @test Array(AK.sortperm(array_from_host(A_h); alg=SORT_ALG, dims=1)) == sortperm(A_h; dims=1) + @test Array(AK.sortperm(array_from_host(A_h); alg=SORT_ALG, dims=1, rev=true)) == sortperm(A_h; dims=1, rev=true) A_h = rand(Int32(0):Int32(3), nslices, len) - @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=2)) == sortperm(A_h; dims=2) + @test Array(AK.sortperm(array_from_host(A_h); alg=SORT_ALG, dims=2)) == sortperm(A_h; dims=2) end # by, order, temp and the low-memory GPU path A_h = rand(Float32, 300, 700) A = array_from_host(A_h) - @test Array(AK.sortperm(A; prefer_threads, dims=1, by=x->-x)) == sortperm(A_h; dims=1, by=x->-x) - @test Array(AK.sortperm(A; prefer_threads, dims=2, order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) - @test Array(AK.sortperm(A; prefer_threads, dims=2, temp=similar(A, Int))) == sortperm(A_h; dims=2) - if !prefer_threads - @test Array(AK.sortperm(A; prefer_threads, dims=2, alg=AK.MergeSort(lowmem=true))) == sortperm(A_h; dims=2) - @test Array(AK.sortperm(A; prefer_threads, dims=1, alg=AK.MergeSort(lowmem=true), block_size=64)) == sortperm(A_h; dims=1) + @test Array(AK.sortperm(A; alg=SORT_ALG, dims=1, by=x->-x)) == sortperm(A_h; dims=1, by=x->-x) + @test Array(AK.sortperm(A; alg=SORT_ALG, dims=2, order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) + @test Array(AK.sortperm(A; alg=SORT_ALG, dims=2, temp=similar(A, Int))) == sortperm(A_h; dims=2) + if TEST_KERNELS + @test Array(AK.sortperm(A; dims=2, alg=AK.MergeSort(lowmem=true))) == sortperm(A_h; dims=2) + @test Array(AK.sortperm(A; dims=1, alg=AK.MergeSort(lowmem=true, block_size=64))) == sortperm(A_h; dims=1) end # In-place fills ix with the same global linear indices as Base A_h = rand(Int32(0):Int32(5), 40, 31) A = array_from_host(A_h) ix = array_from_host(zeros(Int, 40, 31)) - AK.sortperm!(ix, A; prefer_threads, dims=2) + AK.sortperm!(ix, A; alg=SORT_ALG, dims=2) @test Array(ix) == sortperm(A_h; dims=2) # dims=1 on a vector is a full sortperm v_h = rand(Int32(0):Int32(9), 5000) v = array_from_host(v_h) - @test Array(AK.sortperm(v; prefer_threads, dims=1)) == sortperm(v_h) + @test Array(AK.sortperm(v; alg=SORT_ALG, dims=1)) == sortperm(v_h) # Empty and singleton slices for sz in ((0, 5), (5, 0), (1, 64), (64, 1)), dim in 1:2 A_h = rand(Float32, sz...) - @test Array(AK.sortperm(array_from_host(A_h); prefer_threads, dims=dim)) == sortperm(A_h; dims=dim) + @test Array(AK.sortperm(array_from_host(A_h); alg=SORT_ALG, dims=dim)) == sortperm(A_h; dims=dim) end # Slice offsets refer to the view's linear indexing, not its parent's strides. @@ -972,8 +1037,8 @@ end A = array_from_host(A_h) V_h = view(A_h, 1:2:6, 1:2:1030) V = view(A, 1:2:6, 1:2:1030) - @test Array(AK.sortperm(V; prefer_threads, dims=2)) == sortperm(V_h; dims=2) - AK.sort!(V; prefer_threads, dims=2) + @test Array(AK.sortperm(V; alg=SORT_ALG, dims=2)) == sortperm(V_h; dims=2) + AK.sort!(V; alg=SORT_ALG, dims=2) sort!(V_h; dims=2) @test Array(A) == A_h @@ -982,35 +1047,35 @@ end A = array_from_host(A_h) for dim in 1:2, rev in (false, true) expected = sortperm(A_h; dims=dim, rev) - @test Array(AK.sortperm(A; prefer_threads, dims=dim, rev)) == expected - if !prefer_threads - @test Array(AK.merge_sortperm(A; dims=dim, rev)) == expected - @test Array(AK.merge_sortperm_lowmem(A; dims=dim, rev)) == expected + @test Array(AK.sortperm(A; alg=SORT_ALG, dims=dim, rev)) == expected + if TEST_KERNELS + @test Array(AK.sortperm(A; alg=AK.MergeSort(), dims=dim, rev)) == expected + @test Array(AK.sortperm(A; alg=AK.MergeSort(lowmem=true), dims=dim, rev)) == expected end end # Invalid dimensions must not overwrite the output. ix = array_from_host(fill(-1, size(A_h))) - @test_throws ArgumentError AK.sortperm!(ix, A; prefer_threads, dims=3) + @test_throws ArgumentError AK.sortperm!(ix, A; alg=SORT_ALG, dims=3) @test Array(ix) == fill(-1, size(A_h)) - if !prefer_threads - @test_throws ArgumentError AK.merge_sortperm!(ix, A; dims=3) + if TEST_KERNELS + @test_throws ArgumentError AK.sortperm!(ix, A; alg=AK.MergeSort(), dims=3) @test Array(ix) == fill(-1, size(A_h)) - @test_throws ArgumentError AK.merge_sortperm_lowmem!(ix, A; dims=3) + @test_throws ArgumentError AK.sortperm!(ix, A; alg=AK.MergeSort(lowmem=true), dims=3) @test Array(ix) == fill(-1, size(A_h)) - @test_throws ArgumentError AK.merge_sort_by_key!(copy(A), similar(ix, length(ix)); dims=1) + @test_throws ArgumentError AK.sort_by_key!(copy(A), similar(ix, length(ix)); alg=AK.MergeSort(), dims=1) end # Out-of-range dimension and mismatched index array errors A = array_from_host(rand(Float32, 8, 8)) - @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=3) - @test_throws ArgumentError AK.sortperm(A; prefer_threads, dims=0) - @test_throws ArgumentError AK.sortperm!(array_from_host(zeros(Int, 64)), A; prefer_threads, dims=1) + @test_throws ArgumentError AK.sortperm(A; alg=SORT_ALG, dims=3) + @test_throws ArgumentError AK.sortperm(A; alg=SORT_ALG, dims=0) + @test_throws ArgumentError AK.sortperm!(array_from_host(zeros(Int, 64)), A; alg=SORT_ALG, dims=1) end @testset "bitonic_sort_alg" begin - if !prefer_threads + if TEST_KERNELS Random.seed!(0) alg = AK.BitonicSort() @@ -1020,7 +1085,7 @@ end n = rand(1:100_000) v_h = rand(T, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg) + AK.sort!(v; alg) @test Array(v) == sort(v_h) end end @@ -1030,7 +1095,7 @@ end 4097, 8191, 8192, 8193, 100_000, 1_000_000) v_h = rand(Float32, n) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg) + AK.sort!(v; alg) @test Array(v) == sort(v_h) end @@ -1039,7 +1104,7 @@ end for v_h in (fill(2.5f0, n), Float32.(1:n), Float32.(n:-1:1), Float32.(rand(0:1, n)), Float32.(rand(0:3, n))) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg) + AK.sort!(v; alg) @test Array(v) == sort(v_h) end end @@ -1051,7 +1116,7 @@ end v_h[rand(1:5000, 100)] .= 0.0f0 for rev in (false, true) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg, rev) + AK.sort!(v; alg, rev) @test isequal(Array(v), sort(v_h; rev)) end @@ -1060,7 +1125,7 @@ end for kw in ((rev=true,), (order=Base.Order.Reverse,), (rev=true, order=Base.Order.Reverse), (lt=(>),), (by=abs,), (by=x -> x % Int32(7), rev=true), (lt=(a, b) -> a % 5 < b % 5,)) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg, kw...) + AK.sort!(v; alg, kw...) sorted = Array(v) ord = Base.Order.ord(get(kw, :lt, isless), get(kw, :by, identity), get(kw, :rev, nothing), get(kw, :order, Base.Order.Forward)) @@ -1072,43 +1137,40 @@ end v_h = rand(UInt32, 20_000) for block_size in (32, 128, 256, 512), items_per_thread in (1, 2, 8, 16) v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg=AK.BitonicSort(; block_size, items_per_thread)) + AK.sort!(v; alg=AK.BitonicSort(; block_size, items_per_thread)) @test Array(v) == sort(v_h) end v = array_from_host(v_h) - AK.sort!(v; prefer_threads, alg, block_size=64) + AK.sort!(v; alg=AK.BitonicSort(block_size=64)) @test Array(v) == sort(v_h) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(block_size=100)) - @test_throws ArgumentError AK.sort!(v; prefer_threads, alg=AK.BitonicSort(items_per_thread=3)) - @test_throws ArgumentError AK.sort!(v; prefer_threads, - alg=AK.BitonicSort(block_size=2, items_per_thread=1 << (Sys.WORD_SIZE - 2))) + @test_throws ArgumentError AK.sort!(v; alg=AK.BitonicSort(block_size=100)) + @test_throws ArgumentError AK.sort!(v; alg=AK.BitonicSort(items_per_thread=3)) + @test_throws ArgumentError AK.sort!(v; alg=AK.BitonicSort(block_size=2, items_per_thread=1 << (Sys.WORD_SIZE - 2))) # Empty input - @test isempty(Array(AK.sort!(array_from_host(Int32[]); prefer_threads, alg))) + @test isempty(Array(AK.sort!(array_from_host(Int32[]); alg))) # Matrices are sorted as one flat vector by default m_h = rand(Float32, 100, 30) m = array_from_host(m_h) - AK.sort!(m; prefer_threads, alg) + AK.sort!(m; alg) @test vec(Array(m)) == sort(vec(m_h)) # Out-of-place: input unchanged v_h = rand(Float32, 10_000) v = array_from_host(v_h) - w = AK.sort(v; prefer_threads, alg) + w = AK.sort(v; alg) @test Array(w) == sort(v_h) @test Array(v) == v_h # No permutation path - @test_throws ArgumentError AK.sortperm(v; prefer_threads, alg) - else - @test_throws ArgumentError AK.sort!(rand(Int32, 16); prefer_threads, alg=AK.BitonicSort()) + @test_throws ArgumentError AK.sortperm(v; alg) end end @testset "bitonic_sort_dims" begin - if !prefer_threads + if TEST_KERNELS Random.seed!(0) alg = AK.BitonicSort() @@ -1119,11 +1181,11 @@ end sz = dim == 1 ? (L, ncols) : (ncols, L) h = rand(T, sz...) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg) + AK.sort!(v; dims=dim, alg) @test Array(v) == sort(h; dims=dim) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg, rev=true) + AK.sort!(v; dims=dim, alg, rev=true) @test Array(v) == sort(h; dims=dim, rev=true) end end @@ -1131,7 +1193,7 @@ end h = rand(T, 7, 40, 5) for dim in (1, 2, 3) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg) + AK.sort!(v; dims=dim, alg) @test Array(v) == sort(h; dims=dim) end end @@ -1143,36 +1205,75 @@ end h = rand(UInt32, dim == 1 ? (len, 9) : (9, len)) for kw in ((by=x -> xor(x, mask),), (lt=(a, b) -> xor(a, mask) < xor(b, mask),)) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=dim, alg=packed_alg, kw...) + AK.sort!(v; dims=dim, alg=packed_alg, kw...) @test Array(v) == sort(h; dims=dim, kw...) end end # by / lt, NaNs, vectors, empty and singleton slices h = rand(Float32, 300, 700) - @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg, by=x -> -x)) == sort(h; dims=1, by=x -> -x) - @test Array(AK.sort(array_from_host(h); prefer_threads, dims=2, alg, lt=(>))) == sort(h; dims=2, lt=(>)) + @test Array(AK.sort(array_from_host(h); dims=1, alg, by=x -> -x)) == sort(h; dims=1, by=x -> -x) + @test Array(AK.sort(array_from_host(h); dims=2, alg, lt=(>))) == sort(h; dims=2, lt=(>)) h[rand(1:length(h), 1000)] .= NaN32 - @test isequal(Array(AK.sort(array_from_host(h); prefer_threads, dims=2, alg)), sort(h; dims=2)) + @test isequal(Array(AK.sort(array_from_host(h); dims=2, alg)), sort(h; dims=2)) h = rand(Float32, 1000) - @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg)) == sort(h) - @test size(AK.sort(array_from_host(rand(Float32, 0, 5)); prefer_threads, dims=1, alg)) == (0, 5) - @test size(AK.sort(array_from_host(rand(Float32, 5, 0)); prefer_threads, dims=1, alg)) == (5, 0) + @test Array(AK.sort(array_from_host(h); dims=1, alg)) == sort(h) + @test size(AK.sort(array_from_host(rand(Float32, 0, 5)); dims=1, alg)) == (0, 5) + @test size(AK.sort(array_from_host(rand(Float32, 5, 0)); dims=1, alg)) == (5, 0) h = rand(Float32, 1, 100) - @test Array(AK.sort(array_from_host(h); prefer_threads, dims=1, alg)) == h + @test Array(AK.sort(array_from_host(h); dims=1, alg)) == h # Tuning applies per slice; out-of-place leaves the input untouched h = rand(Float32, 3000, 10) for items_per_thread in (1, 4, 16) v = array_from_host(h) - AK.sort!(v; prefer_threads, dims=1, alg=AK.BitonicSort(; block_size=128, items_per_thread)) + AK.sort!(v; dims=1, alg=AK.BitonicSort(; block_size=128, items_per_thread)) @test Array(v) == sort(h; dims=1) end v = array_from_host(h) - w = AK.sort(v; prefer_threads, dims=1, alg) + w = AK.sort(v; dims=1, alg) @test Array(w) == sort(h; dims=1) @test Array(v) == h - @test_throws ArgumentError AK.sort!(array_from_host(h); prefer_threads, dims=3, alg) + @test_throws ArgumentError AK.sort!(array_from_host(h); dims=3, alg) + end +end + + +@testset "sort: backend-free inputs" begin + # A range has no backend: the result is allocated on the one given + r = AK.sort(5:-1:1; backend=BACKEND) + @test get_backend(r) == BACKEND + @test Array(r) == 1:5 + ix = AK.sortperm(5:-1:1; backend=BACKEND) + @test get_backend(ix) == BACKEND + @test Array(ix) == 5:-1:1 +end + + +@testset "sort: reshaped views" begin + # Merge sort's kernels read such an input through `@Const`, which rebuilds the reshape on the + # device + if TEST_KERNELS + h = rand(Int32, 50, 40) + d = array_from_host(h) + v = vec(view(d, 1:40, 1:30)) + AK.sort!(v; alg=AK.MergeSort()) + @test Array(v) == sort(vec(view(h, 1:40, 1:30))) + end +end + + +@testset "sort: bits-union elements" begin + # Merge sort's kernels read such elements through `@Const` (where the backend's arrays can + # hold them) + if TEST_KERNELS && TEST_BITS_UNIONS + Random.seed!(0) + h = rand(Union{Missing, Int32}[missing, 1, 2, 3], 3000) + @test isequal(Array(AK.sort!(array_from_host(h); alg=AK.MergeSort())), sort(h)) + k = rand(Int32(1):Int32(50), 3000) + kd, vd = array_from_host(k), array_from_host(h) + AK.sort_by_key!(kd, vd; alg=AK.MergeSort()) + @test Array(kd) == sort(k) && isequal(Array(vd), h[sortperm(k)]) end end diff --git a/test/generic/sort_resolve.jl b/test/generic/sort_resolve.jl new file mode 100644 index 00000000..db6edf43 --- /dev/null +++ b/test/generic/sort_resolve.jl @@ -0,0 +1,313 @@ +# Resolution of sorting algorithms and of the backend, without launching kernels: resolution only +# looks at types, sizes and orderings, so host arrays stand in for device arrays throughout. + +# A GPU backend with a tuning the tests can change, and arrays that live on it +struct ResolveTestBackend <: KernelAbstractions.Backend end +const RESOLVE_TUNING = Ref(AK.SortTuning()) +AK.sort_tuning(::ResolveTestBackend, ::Type) = RESOLVE_TUNING[] + +struct ResolveTestArray{T, N} <: AbstractArray{T, N} + data::Array{T, N} +end +Base.size(a::ResolveTestArray) = size(a.data) +Base.getindex(a::ResolveTestArray, i::Int...) = a.data[i...] +KernelAbstractions.get_backend(::ResolveTestArray) = ResolveTestBackend() + +# A backend that cannot run AK's kernels, like KernelAbstractions 0.9's `CPU` +struct NoKernelsTestBackend <: KernelAbstractions.Backend end +AK._runs_kernels(::NoKernelsTestBackend) = false + +# An algorithm that is not a sorting algorithm +struct NotASortTestAlgorithm <: AK.Algorithm end + +# An array type that does not implement `get_backend` +struct NoBackendTestArray <: AbstractVector{Int} end +Base.size(::NoBackendTestArray) = (3,) +Base.getindex(::NoBackendTestArray, i::Int) = i + +# A host vector that must not be asked for its backend +struct NoVoteTestVector{T} <: AbstractVector{T} + data::Vector{T} +end +Base.size(a::NoVoteTestVector) = size(a.data) +Base.getindex(a::NoVoteTestVector, i::Int) = a.data[i] +Base.setindex!(a::NoVoteTestVector, x, i::Int) = (a.data[i] = x) +KernelAbstractions.get_backend(::NoVoteTestVector) = error("backend queried despite an explicit backend") + +const RB = ResolveTestBackend() +const FWD = Base.Order.Forward +const REV = Base.Order.Reverse + +resolve(alg, v; dims=:, ord=FWD, backend=RB, kw...) = AK._resolve_sort(alg, backend, v, dims, ord; kw...) + +function with_tuning(f; kwargs...) + old = RESOLVE_TUNING[] + RESOLVE_TUNING[] = AK.SortTuning(; kwargs...) + try + f() + finally + RESOLVE_TUNING[] = old + end +end + + +@testset "sort resolution: Auto" begin + v = zeros(Float32, 1000) + + # The default tuning reproduces AK's historical choices and settings + @test resolve(AK.Auto(), v) === AK.MergeSort(256, false) + @test resolve(AK.Auto(), v; perm=true) === AK.MergeSort(256, false) + @test resolve(AK.Auto(), v; pairs=true) === AK.MergeSort(256, false) + @test resolve(AK.Auto(), zeros(Int32, 10^8 ÷ 100)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 10, 10); dims=1) isa AK.MergeSort + + # On the host backend: the threaded sample sort, filled + host = AK.HOST_BACKEND + a = resolve(AK.Auto(), v; backend=host) + @test a isa AK.CPUThreads.SampleSort + @test a.max_tasks == Threads.nthreads() && a.min_elems == 1 + @test resolve(AK.Auto(), v; backend=host, perm=true) isa AK.CPUThreads.SampleSort + @test resolve(AK.Auto(), v; backend=host, pairs=true) isa AK.CPUThreads.SampleSort + + # Thresholds: bitonic sort up to and including `bitonic_max_len`, radix sort from + # `radix_min_len`, for element types and orderings where they give the same result + with_tuning(; bitonic_max_len=1024, radix_min_len=4096, radix_block_size=128, + radix_items_per_thread=4, bitonic_block_size=64, bitonic_items_per_thread=4) do + @test resolve(AK.Auto(), zeros(Int32, 1024)) === AK.BitonicSort(64, 4) + @test resolve(AK.Auto(), zeros(Int32, 1025)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 4095)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 4096)) === AK.RadixSort(128, 4) + @test resolve(AK.Auto(), zeros(Int32, 4096); ord=REV) isa AK.RadixSort + @test resolve(AK.Auto(), zeros(Int32, 64, 16); dims=1) isa AK.BitonicSort + @test resolve(AK.Auto(), zeros(Int32, 2048, 16); dims=1) isa AK.MergeSort + + # Radix sort is for whole arrays only: a vector or an N×1 matrix sorted along `dims` + # is a slice, however long + @test resolve(AK.Auto(), zeros(Int32, 10_000); dims=1) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 10_000, 1); dims=1) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 10_000, 1)) isa AK.RadixSort + + # Radix sort needs a supported element type and the default ordering or its reverse + @test resolve(AK.Auto(), zeros(Int16, 10_000)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 10_000); ord=Base.Order.ord(isless, abs, nothing)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 10_000); ord=Base.Order.ord(>, identity, nothing)) isa AK.MergeSort + + # Stable by default: bitonic sort only where equal elements are bitwise identical + for T in (Int8, UInt16, Int32, UInt64, Int128, Bool, Char) + @test resolve(AK.Auto(), Vector{T}(undef, 100)) isa AK.BitonicSort + @test resolve(AK.Auto(), Vector{T}(undef, 100); ord=REV) isa AK.BitonicSort + end + for T in (Float16, Float32, Float64, Tuple{Int32, Int32}) + @test resolve(AK.Auto(), Vector{T}(undef, 100)) isa AK.MergeSort + @test resolve(AK.Auto(stable=false), Vector{T}(undef, 100)) isa AK.BitonicSort + end + @test resolve(AK.Auto(), Vector{Union{Int32, UInt32}}(undef, 100)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 100); ord=Base.Order.ord(isless, abs, nothing)) isa AK.MergeSort + @test resolve(AK.Auto(), zeros(Int32, 100); ord=Base.Order.ord((a, b) -> a > b, identity, nothing)) isa AK.MergeSort + @test resolve(AK.Auto(stable=false), zeros(Int32, 100); ord=Base.Order.ord(isless, abs, nothing)) isa AK.BitonicSort + + # sortperm! and key/value sorting always observe stability, and need a permutation path + @test resolve(AK.Auto(), zeros(Int32, 100); perm=true) isa AK.MergeSort + @test resolve(AK.Auto(stable=false), zeros(Float32, 100); perm=true) isa AK.MergeSort + @test resolve(AK.Auto(stable=false), zeros(Int32, 10_000); pairs=true) isa AK.MergeSort + end + + # A different tuning changes the selection + with_tuning(; radix_min_len=0) do + @test resolve(AK.Auto(), zeros(Float32, 10)) isa AK.RadixSort + end + @test resolve(AK.Auto(), zeros(Float32, 10)) isa AK.MergeSort + + # Selection never depends on the data + @test resolve(AK.Auto(), fill(NaN32, 100)) === resolve(AK.Auto(), zeros(Float32, 100)) +end + + +@testset "sort resolution: explicit algorithms" begin + v = zeros(Float32, 1000) + + # Explicit fields win over the tuning, unset ones come from it + with_tuning(; merge_block_size=128, radix_block_size=512, radix_items_per_thread=8, + bitonic_block_size=32, bitonic_items_per_thread=16, threads_min_elems=7) do + @test resolve(AK.MergeSort(), v) === AK.MergeSort(128, false) + @test resolve(AK.MergeSort(block_size=64), v) === AK.MergeSort(64, false) + @test resolve(AK.MergeSort(lowmem=true), v; perm=true) === AK.MergeSort(128, true) + @test resolve(AK.RadixSort(), v) === AK.RadixSort(512, 8) + @test resolve(AK.RadixSort(items_per_thread=1), v) === AK.RadixSort(512, 1) + @test resolve(AK.BitonicSort(block_size=128), v) === AK.BitonicSort(128, 16) + end + host = AK.HOST_BACKEND + @test resolve(AK.CPUThreads.SampleSort(max_tasks=3), v; backend=host) === + AK.CPUThreads.SampleSort(3, 1) + @test resolve(AK.CPUThreads.SampleSort(min_elems=5), v; backend=host) === + AK.CPUThreads.SampleSort(Threads.nthreads(), 5) + + # Rejections, per algorithm (and none for valid combinations) + rejects(alg, v=v; kw...) = try + resolve(alg, v; kw...) + false + catch err + err isa ArgumentError || rethrow() + true + end + + # MergeSort + @test rejects(AK.MergeSort(block_size=0)) + @test !rejects(AK.MergeSort(block_size=100)) # any positive size + @test rejects(AK.MergeSort(lowmem=true)) # sort! + @test rejects(AK.MergeSort(lowmem=true); pairs=true) # sort_by_key! + @test !rejects(AK.MergeSort(lowmem=true); perm=true) + @test !rejects(AK.MergeSort(), zeros(Float32, 4, 4); dims=2, perm=true) + + # RadixSort + @test rejects(AK.RadixSort(block_size=0)) + @test rejects(AK.RadixSort(block_size=100)) + @test rejects(AK.RadixSort(items_per_thread=0)) + @test rejects(AK.RadixSort(); perm=true) + @test rejects(AK.RadixSort(); pairs=true) + @test rejects(AK.RadixSort(), zeros(Float32, 4, 4); dims=1) + @test rejects(AK.RadixSort(), zeros(Float32, 16); dims=1) + @test rejects(AK.RadixSort(), zeros(Int16, 16)) + @test rejects(AK.RadixSort(), Vector{Tuple{Int32, Int32}}(undef, 16)) + @test rejects(AK.RadixSort(); ord=Base.Order.ord(isless, abs, nothing)) + @test rejects(AK.RadixSort(); ord=Base.Order.ord(>, identity, nothing)) + @test rejects(AK.RadixSort(block_size=8192)) # local memory + @test rejects(AK.RadixSort(block_size=1 << 62)) # no overflowing footprints + @test rejects(AK.RadixSort(items_per_thread=1 << 62)) + @test rejects(AK.RadixSort(items_per_thread=65)) + for T in (UInt32, Int32, Float32, UInt64, Int64, Float64), ord in (FWD, REV) + @test !rejects(AK.RadixSort(), zeros(T, 16); ord) + end + @test !rejects(AK.RadixSort(), zeros(Float32, 4, 4)) # dims=: sorts flat + + # BitonicSort + @test rejects(AK.BitonicSort(block_size=100)) + @test rejects(AK.BitonicSort(items_per_thread=3)) + @test rejects(AK.BitonicSort(block_size=2, items_per_thread=1 << (Sys.WORD_SIZE - 2))) + @test rejects(AK.BitonicSort(); perm=true) + @test rejects(AK.BitonicSort(); pairs=true) + @test !rejects(AK.BitonicSort(), zeros(Float32, 4, 4); dims=2, ord=Base.Order.ord(isless, abs, nothing)) + + # CPUThreads.SampleSort only on the host, where every operation supports it + @test rejects(AK.CPUThreads.SampleSort()) + @test rejects(AK.CPUThreads.SampleSort(max_tasks=0); backend=host) + @test rejects(AK.CPUThreads.SampleSort(min_elems=0); backend=host) + @test !rejects(AK.CPUThreads.SampleSort(); backend=host, perm=true) + @test !rejects(AK.CPUThreads.SampleSort(); backend=host, pairs=true, dims=1) + + # Kernel algorithms need a backend that runs AK's kernels + for alg in (AK.MergeSort(), AK.RadixSort(), AK.BitonicSort()) + @test rejects(alg; backend=NoKernelsTestBackend()) + @test rejects(alg; backend=host) == !AK._runs_kernels(host) + end + + # Domain checks run before anything else, whatever the input length + @test rejects(AK.MergeSort(block_size=-1), zeros(Float32, 0)) + @test rejects(AK.RadixSort(block_size=3), zeros(Float32, 1)) + + # Invalid `dims` + @test rejects(AK.Auto(), zeros(Float32, 4, 4); dims=3) + @test rejects(AK.Auto(), zeros(Float32, 4, 4); dims=0) + + # Not a sorting algorithm + @test rejects(NotASortTestAlgorithm()) +end + + +@testset "sort resolution: inference" begin + for T in (Int32, Float32), (A, D) in ((Vector{T}, Colon), (Matrix{T}, Int)) + # Auto resolves to a small union of concrete algorithms, not `Any` + rt = only(Base.return_types(AK._resolve_sort, + (AK.Auto, ResolveTestBackend, A, D, Base.Order.ForwardOrdering))) + @test rt <: Union{AK.MergeSort, AK.RadixSort, AK.BitonicSort} + @test rt isa Union || isconcretetype(rt) + rt = only(Base.return_types(AK._resolve_sort, + (AK.Auto, typeof(AK.HOST_BACKEND), A, D, Base.Order.ForwardOrdering))) + @test rt === AK.CPUThreads.SampleSort + # An explicit algorithm resolves to itself, or always throws (no radix sort along `dims`) + rt = only(Base.return_types(AK._resolve_sort, + (AK.RadixSort, ResolveTestBackend, A, D, Base.Order.ForwardOrdering))) + @test rt === (D === Colon ? AK.RadixSort : Union{}) + end +end + + +@testset "sort resolution: tunings of $(nameof(typeof(BACKEND)))" begin + # Every tuning this backend's device returns resolves to algorithms that pass the checks + for T in valid_backend_eltypes(BACKEND, (UInt8, Int16, Int32, UInt32, Float32, Int64, Float64)) + for len in (1, 100, 10_000, 10^7), stable in (true, false) + v = Vector{T}(undef, len) # only the type and length matter + @test resolve(AK.Auto(; stable), v; backend=BACKEND) isa AK.SortAlgorithm + @test resolve(AK.Auto(; stable), v; backend=BACKEND, perm=true) isa AK.SortAlgorithm + end + if AK._runs_kernels(BACKEND) + for alg in (AK.MergeSort(), AK.BitonicSort(), AK.RadixSort()) + alg isa AK.RadixSort && !AK._rs_supported(T) && continue + @test resolve(alg, zeros(T, 16); backend=BACKEND) isa typeof(alg) + end + end + end +end + + +@testset "backend resolution" begin + d = ResolveTestArray(zeros(Float32, 4)) + h = zeros(Float32, 4) + host = AK.HOST_BACKEND + + # Every array votes, destination first; backend-free leaves do not + @test AK._resolve_backend(nothing, d) === RB + @test AK._resolve_backend(nothing, d, 1:4, CartesianIndices((2, 2)), LinearIndices(h), + 1.0, Ref(2), (3, 1:2), nothing) === RB + @test AK._resolve_backend(nothing, 1:4, d) === RB + @test AK._resolve_backend(nothing, h, 1:4) == host + + # All votes must agree, including the non-destination inputs + @test_throws ArgumentError AK._resolve_backend(nothing, d, h) + @test_throws ArgumentError AK._resolve_backend(nothing, h, d) + @test_throws ArgumentError AK._resolve_backend(nothing, h, (1, d)) + # ... unless the backend is explicit + @test AK._resolve_backend(RB, h) === RB + @test AK._resolve_backend(host, d, h) == host + + # Broadcasted trees are walked + bc = Base.Broadcast.broadcasted(+, d, Base.Broadcast.broadcasted(*, 2, 1:4)) + @test AK._resolve_backend(nothing, bc) === RB + @test_throws ArgumentError AK._resolve_backend(nothing, h, bc) + # ... including the arrays of a preprocessed one + pbc = Base.Broadcast.preprocess(nothing, + Base.Broadcast.instantiate(Base.Broadcast.broadcasted(identity, d))) + @test AK._resolve_backend(nothing, pbc) === RB + @test_throws ArgumentError AK._resolve_backend(nothing, h, pbc) + + # Nothing votes: the host backend + @test AK._resolve_backend(nothing) == host + @test AK._resolve_backend(nothing, 1:10, 2.0) == host + + # Array types without `get_backend` keep raising KernelAbstractions' error + @test_throws ArgumentError AK._resolve_backend(nothing, NoBackendTestArray()) + @test AK._resolve_backend(RB, NoBackendTestArray()) === RB + # ... and allocating operations, given the backend, do not ask for it either + @test AK.sortperm(NoBackendTestArray(); backend=host) == [1, 2, 3] + @test AK._backend_free(reshape(1:6, 2, 3)) && !AK._backend_free(NoBackendTestArray()) + + # The sorting entry points resolve from all their arrays, and accept only backends + @test_throws ArgumentError AK.sort!(d; temp=h) + @test_throws ArgumentError AK.sortperm!(zeros(Int, 4), d) + @test_throws ArgumentError AK.sort_by_key!(d, h) + @test_throws TypeError AK.sort!(h; backend=:cpu) + # ... and never take a positional backend + @test_throws MethodError AK.sort!(h, host) + + # Nested operations receive the resolved backend instead of querying the arrays again + alg = AK.CPUThreads.SampleSort(max_tasks=4) + v = NoVoteTestVector(Int32[3, 1, 2]) + @test AK.sort!(v; backend=host, alg).data == [1, 2, 3] + ix = NoVoteTestVector(zeros(Int, 100)) + vals = rand(Int32, 100) + @test AK.sortperm!(ix, vals; backend=host, alg).data == sortperm(vals) + k = NoVoteTestVector(rand(Int32(1):Int32(3), 100)) + kv = copy(k.data) + AK.sort_by_key!(k, NoVoteTestVector(collect(1:100)); backend=host, alg) + @test k.data == sort(kv) +end diff --git a/test/runtests.jl b/test/runtests.jl index e901ccc6..27a48107 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -38,6 +38,10 @@ const _testset_setup_code = quote global valid_backend_eltypes valid_backend_eltypes(backend, eltypes) = filter(T -> T !== Float64 || KernelAbstractions.supports_float64(backend), eltypes) + + # Whether the backend's arrays can hold bits unions such as `Union{Missing, Bool}` (Metal.jl's, + # oneAPI.jl's and OpenCL.jl's cannot yet) + global TEST_BITS_UNIONS = nameof(typeof(BACKEND)) ∉ (:MetalBackend, :oneAPIBackend, :OpenCLBackend) end # Build list of active backends, each with setup code @@ -53,6 +57,8 @@ if in_test_env("CUDACore") global BACKEND = CUDABackend() global MAX_BLOCK_SIZE = Int(CUDACore.attribute(CUDACore.device(), CUDACore.DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)) global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = false # ... on the host backend global TEST_DL = true $_testset_setup_code end) @@ -68,6 +74,8 @@ if in_test_env("AMDGPU") global BACKEND = ROCBackend() global MAX_BLOCK_SIZE = 1024 global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = false # ... on the host backend global TEST_DL = true $_testset_setup_code end) @@ -82,6 +90,8 @@ if in_test_env("Metal") global BACKEND = MetalBackend() global MAX_BLOCK_SIZE = Int(Metal.device().maxThreadsPerThreadgroup.width) global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = false # ... on the host backend global TEST_DL = false $_testset_setup_code end) @@ -96,6 +106,8 @@ if in_test_env("oneAPI") global BACKEND = oneAPIBackend() global MAX_BLOCK_SIZE = Int(oneAPI.oneL0.compute_properties(oneAPI.device()).maxTotalGroupSize) global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = false # ... on the host backend global TEST_DL = false $_testset_setup_code end) @@ -111,17 +123,24 @@ if in_test_env("OpenCL") global BACKEND = OpenCLBackend() global MAX_BLOCK_SIZE = Int(OpenCL.cl.device().max_work_group_size) global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = false # ... on the host backend global TEST_DL = false $_testset_setup_code end) end -# cpu-ka only when --cpu-ka flag passed +# cpu-ka only when --cpu-ka flag passed: AK's kernels on the host backend, which runs them on +# KernelAbstractions 0.10 (on PoCL) but not on 0.9 if args.custom["cpu-ka"] !== nothing + AK._runs_kernels(AK.HOST_BACKEND) || + error("--cpu-ka needs KernelAbstractions 0.10, whose host backend runs AK's kernels") push!(backends, "cpu-ka" => quote global BACKEND = get_backend([]) global MAX_BLOCK_SIZE = 1024 global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = true # AK's kernels run in this configuration + global HOST_KERNELS = true # ... on the host backend global TEST_DL = false $_testset_setup_code end) @@ -133,6 +152,8 @@ if args.custom["cpu"] !== nothing || isempty(backends) global BACKEND = get_backend([]) global MAX_BLOCK_SIZE = 1024 global prefer_threads = true # Also used to determine whether to run the CPU or GPU tests + global TEST_KERNELS = false + global HOST_KERNELS = false global TEST_DL = false $_testset_setup_code end) From 5ffbdb56b2ab2138b10caaf40fddbecbaa48180b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 23 Sep 2026 17:23:34 +0200 Subject: [PATCH 2/8] Move the launch wrappers and searches to the new calling convention `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`. --- docs/src/api/binarysearch.md | 12 +- docs/src/api/foreachindex.md | 12 +- src/accumulate/accumulate_nd.jl | 2 +- src/findall.jl | 3 +- src/foreachindex.jl | 160 +++++++++++-------------- src/map.jl | 65 ++++------ src/reduce/mapreduce_nd.jl | 2 +- src/reduce/utilities.jl | 2 +- src/reverse.jl | 98 ++++++++-------- src/searchsorted.jl | 202 ++++++++++---------------------- src/sort/cpu_sample_sort.jl | 2 +- src/sort/merge_sort.jl | 2 +- src/sort/merge_sortperm.jl | 4 +- test/generic/binarysearch.jl | 109 ++++++----------- test/generic/looping.jl | 44 +++++-- test/generic/map.jl | 26 +++- test/generic/reverse.jl | 102 +++++++++------- 17 files changed, 374 insertions(+), 473 deletions(-) diff --git a/docs/src/api/binarysearch.md b/docs/src/api/binarysearch.md index 6f5ef485..37429253 100644 --- a/docs/src/api/binarysearch.md +++ b/docs/src/api/binarysearch.md @@ -1,9 +1,9 @@ ### Binary Search -Find the indices where some elements `x` should be inserted into a sorted sequence `v` to maintain the sorted order. Effectively applying the Julia.Base functions in parallel on a GPU using `foreachindex`. -- `searchsortedfirst!` (in-place), `searchsortedfirst` (allocating): index of first element in `v` >= `x[j]`. -- `searchsortedlast!`, `searchsortedlast`: index of last element in `v` <= `x[j]`. -- **Other names**: `thrust::upper_bound`, `std::lower_bound`. +Find the indices where many elements `x[j]` would be inserted into a sorted sequence `v` to maintain the sorted order: the Julia Base functions applied to each query in parallel. There are no allocating forms, because `Base.searchsortedfirst(v, x)` treats a vector `x` as one value. +- `searchsortedfirst!`: index of the first element of `v` not ordered before `x[j]` (`>= x[j]` by default). +- `searchsortedlast!`: index of the last element of `v` not ordered after `x[j]` (`<= x[j]` by default). +- **Other names**: `thrust::lower_bound`, `thrust::upper_bound`, `std::lower_bound`. Example: @@ -13,7 +13,7 @@ using Metal # Sorted array v = MtlArray(rand(Float32, 100_000)) -AK.sort!(v; alg=AK.MergeSort()) +AK.sort!(v) # Elements `x` to place within `v` at indices `ix` x = MtlArray(rand(Float32, 10_000)) @@ -25,7 +25,5 @@ AK.searchsortedfirst!(ix, v, x) ```@docs AcceleratedKernels.searchsortedfirst! -AcceleratedKernels.searchsortedfirst AcceleratedKernels.searchsortedlast! -AcceleratedKernels.searchsortedlast ``` diff --git a/docs/src/api/foreachindex.md b/docs/src/api/foreachindex.md index 46246cef..b49bd992 100644 --- a/docs/src/api/foreachindex.md +++ b/docs/src/api/foreachindex.md @@ -56,14 +56,22 @@ v2 = oneArray(rand(Float32, 100_000)) f(v1, v2) ``` +The loop runs on the backend of the iterable. A range does not determine one, so pass `backend` when a loop over a range accesses GPU arrays: +```julia +AK.foreachindex(1:n; backend=AK.get_backend(v)) do i + v[i] = i +end +``` +Host arrays and ranges run on Julia threads. + All GPU functions allow you to specify a block size - this is often a power of two (mostly 64, 128, 256, 512); the optimum depends on the algorithm, input data and hardware - you can try the different values and `@time` or `@benchmark` them: ```julia -@time AK.foreachindex(f, itr_gpu, block_size=512) +@time AK.foreachindex(f, itr_gpu; block_size=512) ``` Similarly, for performance on the CPU the overhead of spawning threads should be masked by processing more elements per thread (but there is no reason here to launch more threads than `Threads.nthreads()`, the number of threads Julia was started with); the optimum depends on how expensive `f` is - again, benchmarking is your friend: ```julia -@time AK.foreachindex(f, itr_cpu, max_tasks=16, min_elems=1000) +@time AK.foreachindex(f, itr_cpu; max_tasks=16, min_elems=1000) ``` diff --git a/src/accumulate/accumulate_nd.jl b/src/accumulate/accumulate_nd.jl index 5a44c34f..3dea501a 100644 --- a/src/accumulate/accumulate_nd.jl +++ b/src/accumulate/accumulate_nd.jl @@ -86,7 +86,7 @@ function _accumulate_nd_cpu_sections!( # Each thread handles a section of the output array - i.e. reducing along the dims, for # multiple output strides - foreachindex(1:length_outer, CPU(), max_tasks=max_tasks, min_elems=min_elems) do idst + _foreachindex(1:length_outer, HOST_BACKEND; max_tasks, min_elems) do idst @inbounds begin # Compute the base index in v for this outer axis diff --git a/src/findall.jl b/src/findall.jl index 055a161d..486733d1 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -206,8 +206,7 @@ function findall_bools( end input_indices = eachindex(v) bool_indices = eachindex(bools) - foreachindex(Base.OneTo(length(v)), backend; - max_tasks, min_elems, prefer_threads, block_size) do position + _foreachindex(Base.OneTo(length(v)), backend; max_tasks, min_elems, block_size) do position input_index = findall_index(input_indices, position) bool_index = findall_index(bool_indices, position) @inbounds bools[bool_index] = pred(v[input_index]) ? true : false diff --git a/src/foreachindex.jl b/src/foreachindex.jl index ab46cff7..bf4bb588 100644 --- a/src/foreachindex.jl +++ b/src/foreachindex.jl @@ -41,27 +41,50 @@ function _forindices_threads(f, indices; max_tasks, min_elems) end -""" - foreachindex( - f, itr, backend::Backend=get_backend(itr); +# Launch settings of the wrappers, checked on every backend (only one of them applies to each) +function _check_launch(; block_size, max_tasks, min_elems) + block_size >= 1 || throw(ArgumentError("`block_size` must be positive, got $block_size")) + max_tasks >= 1 || throw(ArgumentError("`max_tasks` must be positive, got $max_tasks")) + min_elems >= 1 || throw(ArgumentError("`min_elems` must be positive, got $min_elems")) + nothing +end + +# Run `f` over `indices`: on Julia threads on the host backend, as a kernel elsewhere +function _foreachindex(f, indices, backend::Backend; + block_size=256, max_tasks=Threads.nthreads(), min_elems=1) + _check_launch(; block_size, max_tasks, min_elems) + if _runs_threads(backend) + _forindices_threads(f, indices; max_tasks, min_elems) + else + _forindices_gpu(f, indices, backend; block_size) + end + nothing +end - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - # GPU settings - block_size=256, +""" + foreachindex( + f, itr; + backend=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) -Parallelised `for` loop over the indices of an iterable. +Parallelised `for` loop over the indices of an iterable: call `f(i)` for every `i` in +`eachindex(itr)`. It allows you to run normal Julia code on a GPU over multiple arrays - e.g. CuArray, ROCArray, -MtlArray, oneArray - with one GPU thread per index. +MtlArray, oneArray - with one GPU thread per index, in blocks of `block_size` threads. -On CPUs at most `max_tasks` threads are launched, or fewer such that each thread processes at least -`min_elems` indices; if a single task ends up being needed, `f` is inlined and no thread is -launched. Tune it to your function - the more expensive it is, the fewer elements are needed to -amortise the cost of launching a thread (which is a few μs). +On the host backend, the loop runs on Julia threads: at most `max_tasks` tasks, or fewer such that +each task processes at least `min_elems` indices; if a single task ends up being needed, `f` is +inlined and no task is launched. Tune it to your function - the more expensive it is, the fewer +elements are needed to amortise the cost of launching a task (which is a few μs). + +`backend` is derived from `itr`. A range or other index collection does not determine a backend, +so loops over one run on the host unless you pass `backend`, e.g. the backend of the arrays that +`f` accesses. # Examples Normally you would write a for loop like this: @@ -89,6 +112,15 @@ function f() end ``` +A loop over a range needs the backend of the arrays it accesses: +```julia +function g!(y, x) + AK.foreachindex(1:length(x) ÷ 2; backend=AK.get_backend(x)) do i + @inbounds y[i] = x[2i] + end +end +``` + Note that the above code is pure arithmetic, which you can write directly (and on some platforms it may be faster) as: ```julia @@ -121,45 +153,29 @@ somecopy!(x) # This works ``` """ function foreachindex( - f, itr, backend::Backend=get_backend(itr); - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, - - # GPU settings - block_size=256, + f, itr; + backend::Union{Nothing, Backend}=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) - if use_gpu_algorithm(backend, prefer_threads) - _forindices_gpu(f, eachindex(itr), backend; block_size) - else - _forindices_threads(f, eachindex(itr); max_tasks, min_elems) - end + backend = _resolve_backend(backend, itr) + _foreachindex(f, eachindex(itr), backend; block_size, max_tasks, min_elems) end """ foraxes( - f, itr, dims::Union{Nothing, <:Integer}=nothing, backend::Backend=get_backend(itr); - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, + f, itr, dims::Union{Nothing, Integer}=nothing; + backend=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) -Parallelised `for` loop over the indices along axis `dims` of an iterable. - -It allows you to run normal Julia code on a GPU over multiple arrays - e.g. CuArray, ROCArray, -MtlArray, oneArray - with one GPU thread per index. - -On CPUs at most `max_tasks` threads are launched, or fewer such that each thread processes at least -`min_elems` indices; if a single task ends up being needed, `f` is inlined and no thread is -launched. Tune it to your function - the more expensive it is, the fewer elements are needed to -amortise the cost of launching a thread (which is a few μs). +Parallelised `for` loop over the indices along axis `dims` of an iterable: call `f(i)` for every +`i` in `axes(itr, dims)`, or in `eachindex(itr)` for `dims=nothing`. `dims` mirrors +`axes(itr, dims)`. The keywords are those of [`foreachindex`](@ref). # Examples Normally you would write a for loop like this: @@ -191,51 +207,17 @@ function f() end ``` -**Important note**: to use this function on a GPU, the objects referenced inside the loop body must -have known types - i.e. be inside a function. For example: -```julia -using oneAPI -import AcceleratedKernels as AK - -x = oneArray(reshape(1:3000, 3, 1000)) - -# CRASHES - typical error message: "Reason: unsupported dynamic function invocation" -# AK.foraxes(x) do i -# x[i] = i -# end - -function somecopy!(v) - # Because it is inside a function, the type of `v` will be known - AK.foraxes(v) do i - v[i] = i - end -end - -somecopy!(x) # This works -``` +As with [`foreachindex`](@ref), the objects referenced inside the loop body must have known types +on a GPU, i.e. be inside a function. """ function foraxes( - f, itr, dims::Union{Nothing, <:Integer}=nothing, backend::Backend=get_backend(itr); - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, - - # GPU settings - block_size=256, + f, itr, dims::Union{Nothing, Integer}=nothing; + backend::Union{Nothing, Backend}=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) - if isnothing(dims) - return foreachindex( - f, itr, backend; - max_tasks, min_elems, - prefer_threads, block_size, - ) - end - - if use_gpu_algorithm(backend, prefer_threads) - _forindices_gpu(f, axes(itr, dims), backend; block_size) - else - _forindices_threads(f, axes(itr, dims); max_tasks, min_elems) - end + backend = _resolve_backend(backend, itr) + indices = isnothing(dims) ? eachindex(itr) : axes(itr, dims) + _foreachindex(f, indices, backend; block_size, max_tasks, min_elems) end diff --git a/src/map.jl b/src/map.jl index 10d690da..eb1dd094 100644 --- a/src/map.jl +++ b/src/map.jl @@ -1,17 +1,15 @@ """ map!( - f, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); + f, dst::AbstractArray, src::AbstractArray; + backend=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, + ) -> dst - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, - ) - -Apply the function `f` to each element of `src` in parallel and store the result in `dst`. The -CPU and GPU settings are the same as for [`foreachindex`](@ref). +Apply the function `f` to each element of `src` in parallel and store the result in `dst`, which +must have as many elements. `backend` is derived from `dst` and `src`; the other keywords are +those of [`foreachindex`](@ref). On CPUs, multithreading only improves performance when complex computation hides the memory latency and the overhead of spawning tasks - that includes more complex functions and less @@ -20,7 +18,7 @@ threads. # Examples ```julia -import Metal +using Metal import AcceleratedKernels as AK x = MtlArray(rand(Float32, 100_000)) @@ -32,14 +30,16 @@ end ``` """ function map!( - f, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); - kwargs... + f, dst::AbstractArray, src::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) - @argcheck length(dst) == length(src) - foreachindex( - src, backend; - kwargs... - ) do idx + backend = _resolve_backend(backend, dst, src) + length(dst) == length(src) || throw(ArgumentError( + "destination and source must have the same length, $(length(dst)) != $(length(src))")) + _foreachindex(eachindex(src), backend; block_size, max_tasks, min_elems) do idx dst[idx] = f(src[idx]) end dst @@ -47,28 +47,13 @@ end """ - map( - f, src::AbstractArray, backend::Backend=get_backend(src); - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, - ) + map(f, src::AbstractArray; kwargs...) Apply the function `f` to each element of `src` and store the results in a copy of `src` (if `f` -changes the `eltype`, allocate `dst` separately and call [`map!`](@ref)). The CPU and GPU -settings are the same as for [`foreachindex`](@ref). +changes the `eltype`, allocate `dst` separately and call [`map!`](@ref)). The keywords are those +of [`map!`](@ref). """ -function map( - f, src::AbstractArray, backend::Backend=get_backend(src); - kwargs... -) - dst = similar(src) - map!( - f, dst, src, backend; - kwargs... - ) +function map(f, src::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) + backend = _resolve_backend(backend, src) + return map!(f, _similar(backend, src), src; backend, kwargs...) end diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 5304cce5..f2b6c696 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -487,7 +487,7 @@ function _mapreduce_nd_cpu_sections!( Rother = CartesianIndices(dst) Rreduce = _mapreduce_reduce_indices(src, dst) - foreachindex(dst, max_tasks=max_tasks, min_elems=min_elems) do idst + _foreachindex(eachindex(dst), HOST_BACKEND; max_tasks, min_elems) do idst @inbounds begin Iother = Rother[idst] res = init diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index d3000532..452c2b98 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -39,7 +39,7 @@ function _mapreduce_nd_apply_init!( min_elems=1, block_size=256, ) - foreachindex(dst, backend; max_tasks, min_elems, block_size) do i + _foreachindex(eachindex(dst), backend; max_tasks, min_elems, block_size) do i dst[i] = op(init, f(src[i])) end end diff --git a/src/reverse.jl b/src/reverse.jl index a47dc333..694247fc 100644 --- a/src/reverse.jl +++ b/src/reverse.jl @@ -30,7 +30,7 @@ function reverse_dims!( reduced_size = ntuple(d -> ifelse(d == half_dim, cld(size(v, d), 2), size(v, d)), N) nd_idx = CartesianIndices(reduced_size) - foreachindex(1:Base.prod(reduced_size), backend; kwargs...) do i + _foreachindex(1:Base.prod(reduced_size), backend; kwargs...) do i idx = Tuple(nd_idx[i]) index_in = lin_idx[idx...] idx_mirror = ifelse.(rev_dims, ref .- idx, idx) @@ -55,7 +55,7 @@ function reverse_dims!( lin_idx = LinearIndices(src) nd_idx = CartesianIndices(src) - foreachindex(src, backend; kwargs...) do i + _foreachindex(eachindex(src), backend; kwargs...) do i idx = Tuple(nd_idx[i]) idx_mirror = ifelse.(rev_dims, ref .- idx, idx) index_out = lin_idx[idx_mirror...] @@ -68,22 +68,18 @@ end """ reverse!( - v::AbstractArray, backend::Backend=get_backend(v); - + v::AbstractArray; + backend=nothing, dims=:, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, - ) + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, + ) -> v Reverse `v` in-place and return it. With `dims=:` (the default) the whole array is reversed; pass `dims=d` (an integer or an iterable of distinct integers in `1:ndims(v)`) to reverse only along -those dimensions. `dims=()` leaves `v` unchanged. The CPU and GPU settings are the same as for -[`foreachindex`](@ref). +those dimensions. `dims=()` leaves `v` unchanged. `backend` is derived from `v`; the other +keywords are those of [`foreachindex`](@ref). No temporary array is allocated. @@ -102,12 +98,19 @@ AK.reverse!(m; dims=2) # reverse the columns ``` """ function reverse!( - v::AbstractArray, backend::Backend=get_backend(v); - dims=:, kwargs... + v::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + dims=:, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) + backend = _resolve_backend(backend, v) + launch = (; block_size, max_tasks, min_elems) + _check_launch(; launch...) # before the early returns dims = check_reverse_dims(v, dims) if !(dims isa Colon) - return reverse_dims!(v, dims, backend; kwargs...) + return reverse_dims!(v, dims, backend; launch...) end len = length(v) @@ -117,7 +120,7 @@ function reverse!( hi = lastindex(v) # Swap each pair once; an odd-length array keeps its middle element. - foreachindex(1:(len ÷ 2), backend; kwargs...) do i + _foreachindex(1:(len ÷ 2), backend; launch...) do i left = lo + i - 1 right = hi - i + 1 @inbounds begin @@ -133,33 +136,37 @@ end """ reverse!( - dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); - + dst::AbstractArray, src::AbstractArray; + backend=nothing, dims=:, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, - ) + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, + ) -> dst Write the reverse of `src` into `dst` and return `dst`; `src` is left unchanged. `dst` and `src` must not alias. With `dims=:` (the default), the whole array is reversed and only the lengths must match. With an integer or iterable of distinct dimensions, the sizes must match. `dims=()` copies `src` unchanged. Elements are converted to the destination element type on assignment. -The CPU and GPU settings are the same as for [`foreachindex`](@ref). +`backend` is derived from `dst` and `src`; the other keywords are those of +[`foreachindex`](@ref). """ function reverse!( - dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); - dims=:, kwargs... + dst::AbstractArray, src::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + dims=:, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, ) + backend = _resolve_backend(backend, dst, src) + launch = (; block_size, max_tasks, min_elems) + _check_launch(; launch...) # before the early returns dims = check_reverse_dims(src, dims) if !(dims isa Colon) @argcheck size(dst) == size(src) length(src) == 0 && return dst - return reverse_dims!(dst, src, dims, backend; kwargs...) + return reverse_dims!(dst, src, dims, backend; launch...) end @argcheck length(dst) == length(src) @@ -168,7 +175,7 @@ function reverse!( hi_src = lastindex(src) lo_dst = firstindex(dst) - foreachindex(src, backend; kwargs...) do i + _foreachindex(eachindex(src), backend; launch...) do i @inbounds dst[lo_dst + (hi_src - i)] = src[i] end @@ -177,28 +184,15 @@ end """ - reverse( - v::AbstractArray, backend::Backend=get_backend(v); - - dims=:, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size=256, - ) + reverse(v::AbstractArray; kwargs...) Return a reversed copy of `v`, leaving `v` unchanged. With `dims=:` (the default) the whole array is -reversed; pass `dims=d` to reverse only along those dimensions, matching `Base.reverse`. The CPU and -GPU settings are the same as for [`foreachindex`](@ref). +reversed; pass `dims=d` to reverse only along those dimensions, matching `Base.reverse`. The +keywords are those of [`reverse!`](@ref). Prefer [`reverse!`](@ref) when you do not need to keep `v`; it avoids the allocation. """ -function reverse( - v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - reverse!(similar(v), v, backend; kwargs...) +function reverse(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) + backend = _resolve_backend(backend, v) + return reverse!(_similar(backend, v), v; backend, kwargs...) end diff --git a/src/searchsorted.jl b/src/searchsorted.jl index 9039f12a..c283b693 100644 --- a/src/searchsorted.jl +++ b/src/searchsorted.jl @@ -70,165 +70,89 @@ end """ searchsortedfirst!( - ix::AbstractVector, - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings + ix::AbstractVector, v::AbstractVector, xs::AbstractVector; + backend=nothing, + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + block_size::Int=256, max_tasks::Int=Threads.nthreads(), min_elems::Int=1000, - - # GPU settings - block_size::Int=256, - ) - -Equivalent to applying `searchsortedfirst!` element-wise to each element of `x`. The CPU and GPU -settings are the same as for [`foreachindex`](@ref). + ) -> ix + +Batched binary search: for each element `xs[i]`, write `searchsortedfirst(v, xs[i]; lt, by, rev, +order)` into `ix[i]`, the index of the first element of the sorted vector `v` not ordered before +`xs[i]`. Unlike `Base.searchsortedfirst`, which would treat `xs` as a single value, every element +of `xs` is a separate query. `ix` needs as many elements as `xs`. + +`backend` is derived from `ix`, `v` and `xs`; the launch keywords are those of +[`foreachindex`](@ref), except that on the host at least `min_elems=1000` queries go to each +task. + +# Examples +```julia +import AcceleratedKernels as AK +using CUDA + +v = CuArray(sort(rand(Float32, 1_000_000))) +xs = CuArray(rand(Float32, 10_000)) +ix = similar(xs, Int) +AK.searchsortedfirst!(ix, v, xs) +``` """ function searchsortedfirst!( - ix::AbstractVector, - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings with different default from `foreachindex` + ix::AbstractVector, v::AbstractVector, xs::AbstractVector; + backend::Union{Nothing, Backend}=nothing, + lt=isless, + by=identity, + rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), min_elems::Int=1000, - - kwargs... ) - # Simple sanity checks - @argcheck length(ix) == length(x) - - # Construct comparator - ord = Base.Order.ord(lt, by, rev) - comp = (x, y) -> Base.Order.lt(ord, x, y) - - foreachindex( - x, backend; - min_elems, kwargs... - ) do i - @inbounds ix[i] = _searchsortedfirst(v, x[i], firstindex(v), lastindex(v), comp) + backend = _resolve_backend(backend, ix, v, xs) + length(ix) == length(xs) || throw(ArgumentError( + "index array must have as many elements as the queries, $(length(ix)) != $(length(xs))")) + ord = Base.Order.ord(lt, by, rev, order) + _foreachindex(eachindex(xs), backend; block_size, max_tasks, min_elems) do i + @inbounds ix[i] = _searchsortedfirst(v, xs[i], firstindex(v), lastindex(v), ord) end -end - - -""" - searchsortedfirst( - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1000, - - # GPU settings - block_size::Int=256, - ) - -Equivalent to applying `searchsortedfirst` element-wise to each element of `x`. The CPU and GPU -settings are the same as for [`foreachindex`](@ref). -""" -function searchsortedfirst( - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - kwargs... -) - ix = similar(x, Int) - searchsortedfirst!( - ix, v, x, backend; - kwargs... - ) ix end """ searchsortedlast!( - ix::AbstractVector, - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings + ix::AbstractVector, v::AbstractVector, xs::AbstractVector; + backend=nothing, + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + block_size::Int=256, max_tasks::Int=Threads.nthreads(), min_elems::Int=1000, + ) -> ix - # GPU settings - block_size::Int=256, - ) - -Equivalent to applying `searchsortedlast!` element-wise to each element of `x`. The CPU and GPU -settings are the same as for [`foreachindex`](@ref). +Batched binary search: for each element `xs[i]`, write `searchsortedlast(v, xs[i]; lt, by, rev, +order)` into `ix[i]`, the index of the last element of the sorted vector `v` not ordered after +`xs[i]`. The keywords are those of [`searchsortedfirst!`](@ref). """ function searchsortedlast!( - ix::AbstractVector, - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings with different default from `foreachindex` + ix::AbstractVector, v::AbstractVector, xs::AbstractVector; + backend::Union{Nothing, Backend}=nothing, + lt=isless, + by=identity, + rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, + block_size::Int=256, + max_tasks::Int=Threads.nthreads(), min_elems::Int=1000, - - kwargs... ) - # Simple sanity checks - @argcheck length(ix) == length(x) - - # Construct comparator - ord = Base.Order.ord(lt, by, rev) - comp = (x, y) -> Base.Order.lt(ord, x, y) - - foreachindex( - x, backend; - min_elems, kwargs... - ) do i - @inbounds ix[i] = _searchsortedlast(v, x[i], firstindex(v), lastindex(v), comp) + backend = _resolve_backend(backend, ix, v, xs) + length(ix) == length(xs) || throw(ArgumentError( + "index array must have as many elements as the queries, $(length(ix)) != $(length(xs))")) + ord = Base.Order.ord(lt, by, rev, order) + _foreachindex(eachindex(xs), backend; block_size, max_tasks, min_elems) do i + @inbounds ix[i] = _searchsortedlast(v, xs[i], firstindex(v), lastindex(v), ord) end -end - - -""" - searchsortedlast( - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - - by=identity, lt=isless, rev::Bool=false, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1000, - - # GPU settings - block_size::Int=256, - ) - -Equivalent to applying `searchsortedlast` element-wise to each element of `x`. The CPU and GPU -settings are the same as for [`foreachindex`](@ref). -""" -function searchsortedlast( - v::AbstractVector, - x::AbstractVector, - backend::Backend=get_backend(x); - kwargs... -) - ix = similar(x, Int) - searchsortedlast!( - ix, v, x, backend; - kwargs... - ) ix end diff --git a/src/sort/cpu_sample_sort.jl b/src/sort/cpu_sample_sort.jl index f5c39952..9a145364 100644 --- a/src/sort/cpu_sample_sort.jl +++ b/src/sort/cpu_sample_sort.jl @@ -262,7 +262,7 @@ function _sample_sortperm!( @argcheck length(ix) == length(v) # Initialise indices that will be sorted by the keys in v - foreachindex(ix, HOST_BACKEND; max_tasks, min_elems) do i + _foreachindex(eachindex(ix), HOST_BACKEND; max_tasks, min_elems) do i @inbounds ix[i] = i end diff --git a/src/sort/merge_sort.jl b/src/sort/merge_sort.jl index 732f5d7c..83705328 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -155,7 +155,7 @@ function _merge_sort!( # keys' type (broadcasting would make a `BitArray` of `Bool` keys on the host) if by !== identity keys = similar(v, Base.promote_op(by, eltype(v))) - map!(by, keys, v, backend; block_size) + map!(by, keys, v; backend, block_size) _merge_sort_by_key!( keys, v, backend; lt, rev, order, block_size, dims, diff --git a/src/sort/merge_sortperm.jl b/src/sort/merge_sortperm.jl index 1500047a..88c8d570 100644 --- a/src/sort/merge_sortperm.jl +++ b/src/sort/merge_sortperm.jl @@ -32,7 +32,7 @@ function _merge_sortperm!( end # Initialise the linear indices that will be sorted by the keys in v - foreachindex(ix, backend; block_size) do i + _foreachindex(eachindex(ix), backend; block_size) do i @inbounds ix[i] = i end (isempty(v) || layout.len <= 1) && return ix @@ -75,7 +75,7 @@ function _merge_sortperm_lowmem!( end # Initialise the linear indices that will be sorted by the keys in v - foreachindex(ix, backend; block_size) do i + _foreachindex(eachindex(ix), backend; block_size) do i @inbounds ix[i] = i end (isempty(ix) || layout.len <= 1) && return ix diff --git a/test/generic/binarysearch.jl b/test/generic/binarysearch.jl index b1dbf7f9..0245b736 100644 --- a/test/generic/binarysearch.jl +++ b/test/generic/binarysearch.jl @@ -2,92 +2,51 @@ Random.seed!(0) - # Fuzzy correctness testing of searchsortedfirst - for _ in 1:100 + # Fuzzy correctness testing against Base, applied to each query + for _ in 1:100, T in (Int32, Float32) num_elems_v = rand(1:100_000) num_elems_x = rand(1:100_000) + vh = sort(rand(T, num_elems_v)) + xh = rand(T, num_elems_x) + v = array_from_host(vh) + x = array_from_host(xh) - # Ints - v = array_from_host(sort(rand(Int32, num_elems_v))) - x = array_from_host(rand(Int32, num_elems_x)) ix = similar(x, Int32) - AK.searchsortedfirst!(ix, v, x; prefer_threads) + @test AK.searchsortedfirst!(ix, v, x) === ix + @test Array(ix) == [searchsortedfirst(vh, e) for e in xh] - vh = Array(v) - xh = Array(x) - ixh = AK.searchsortedfirst(vh, xh; prefer_threads=true) - ixh_base = [searchsortedfirst(vh, e) for e in xh] - - @test all(Array(ix) .== ixh .== ixh_base) - - # Floats - v = array_from_host(sort(rand(Float32, num_elems_v))) - x = array_from_host(rand(Float32, num_elems_x)) ix = similar(x, Int32) - AK.searchsortedfirst!(ix, v, x; prefer_threads) - - vh = Array(v) - xh = Array(x) - ixh = AK.searchsortedfirst(vh, xh; prefer_threads=true) - ixh_base = [searchsortedfirst(vh, e) for e in xh] - - @test all(Array(ix) .== ixh .== ixh_base) + @test AK.searchsortedlast!(ix, v, x) === ix + @test Array(ix) == [searchsortedlast(vh, e) for e in xh] end - # Fuzzy correctness testing of searchsortedlast - for _ in 1:100 - num_elems_v = rand(1:100_000) - num_elems_x = rand(1:100_000) - - # Ints - v = array_from_host(sort(rand(Int32, num_elems_v))) - x = array_from_host(rand(Int32, num_elems_x)) - ix = similar(x, Int32) - AK.searchsortedlast!(ix, v, x; prefer_threads) - - vh = Array(v) - xh = Array(x) - ixh = AK.searchsortedlast(vh, xh; prefer_threads=true) - ixh_base = [searchsortedlast(vh, e) for e in xh] - - @test all(Array(ix) .== ixh .== ixh_base) - - # Floats - v = array_from_host(sort(rand(Float32, num_elems_v))) - x = array_from_host(rand(Float32, num_elems_x)) - ix = similar(x, Int32) - AK.searchsortedlast!(ix, v, x; prefer_threads) - - vh = Array(v) - xh = Array(x) - ixh = AK.searchsortedlast(vh, xh; prefer_threads=true) - ixh_base = [searchsortedlast(vh, e) for e in xh] - - @test all(Array(ix) .== ixh .== ixh_base) + # Orderings, as in Base + vh = rand(Int32(-1000):Int32(1000), 10_000) + xh = rand(Int32(-1000):Int32(1000), 1000) + for kw in ((rev=true,), (order=Base.Order.Reverse,), (by=abs,), (lt=(>),), + (by=abs, rev=true), (lt=(>), rev=true)) + svh = sort(vh; kw...) + v = array_from_host(svh) + x = array_from_host(xh) + ix = array_from_host(zeros(Int, length(xh))) + AK.searchsortedfirst!(ix, v, x; kw...) + @test Array(ix) == [searchsortedfirst(svh, e; kw...) for e in xh] + AK.searchsortedlast!(ix, v, x; kw...) + @test Array(ix) == [searchsortedlast(svh, e; kw...) for e in xh] end - # Testing different settings + # Launch settings v = array_from_host(sort(rand(Int32, 100_000))) x = array_from_host(rand(Int32, 10_000)) ix = similar(x, Int32) - - AK.searchsortedfirst!(ix, v, x; prefer_threads, by=abs, lt=(>), rev=true, block_size=64) - AK.searchsortedfirst(v, x; prefer_threads, by=abs, lt=(>), rev=true, block_size=64) - AK.searchsortedlast!(ix, v, x; prefer_threads, by=abs, lt=(>), rev=true, block_size=64) - AK.searchsortedlast(v, x; prefer_threads, by=abs, lt=(>), rev=true, block_size=64) - - vh = Array(v) - xh = Array(x) - ixh = similar(xh, Int32) - - AK.searchsortedfirst!(ixh, vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100) - AK.searchsortedfirst(vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100) - AK.searchsortedlast!(ixh, vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100) - AK.searchsortedlast(vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100) - - # Test that undefined kwargs are not accepted - @test_throws MethodError AK.searchsortedfirst!(ixh, vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100, bad=:kwarg) - @test_throws MethodError AK.searchsortedfirst(vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100, bad=:kwarg) - @test_throws MethodError AK.searchsortedlast!(ixh, vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100, bad=:kwarg) - @test_throws MethodError AK.searchsortedlast(vh, xh; prefer_threads=true, by=abs, lt=(>), rev=true, max_tasks=10, min_elems=100, bad=:kwarg) + AK.searchsortedfirst!(ix, v, x; block_size=64, max_tasks=10, min_elems=100) + @test Array(ix) == [searchsortedfirst(Array(v), e) for e in Array(x)] + AK.searchsortedlast!(ix, v, x; block_size=64, max_tasks=10, min_elems=100) + @test Array(ix) == [searchsortedlast(Array(v), e) for e in Array(x)] + + # Invalid arguments + @test_throws ArgumentError AK.searchsortedfirst!(similar(x, Int32, 3), v, x) + @test_throws ArgumentError AK.searchsortedfirst!(ix, v, x; block_size=0) + @test_throws MethodError AK.searchsortedfirst!(ix, v, x; bad=:kwarg) + @test_throws MethodError AK.searchsortedlast!(ix, v, x; bad=:kwarg) end diff --git a/test/generic/looping.jl b/test/generic/looping.jl index 4980bf30..9bb440cf 100644 --- a/test/generic/looping.jl +++ b/test/generic/looping.jl @@ -3,27 +3,27 @@ Random.seed!(0) # CPU - if prefer_threads + if AK._runs_threads(BACKEND) x = zeros(Int, 1000) - AK.foreachindex(x; prefer_threads) do i + AK.foreachindex(x) do i x[i] = i end @test all(x .== 1:length(x)) x = zeros(Int, 1000) - AK.foreachindex(x; prefer_threads, max_tasks=1, min_elems=1) do i + AK.foreachindex(x; max_tasks=1, min_elems=1) do i x[i] = i end @test all(x .== 1:length(x)) x = zeros(Int, 1000) - AK.foreachindex(x; prefer_threads, max_tasks=10, min_elems=1) do i + AK.foreachindex(x; max_tasks=10, min_elems=1) do i x[i] = i end @test all(x .== 1:length(x)) x = zeros(Int, 1000) - AK.foreachindex(x; prefer_threads, max_tasks=10, min_elems=10) do i + AK.foreachindex(x; max_tasks=10, min_elems=10) do i x[i] = i end @test all(x .== 1:length(x)) @@ -31,7 +31,7 @@ # GPU else x = array_from_host(zeros(Int, 10_000)) - f1(x) = AK.foreachindex(x; prefer_threads) do i # This must be inside a function to have a known type! + f1(x) = AK.foreachindex(x) do i # This must be inside a function to have a known type! x[i] = i end f1(x) @@ -39,13 +39,33 @@ @test all(xh .== 1:length(xh)) x = array_from_host(zeros(Int, 10_000)) - f2(x) = AK.foreachindex(x; prefer_threads, block_size=64) do i + f2(x) = AK.foreachindex(x; block_size=64) do i x[i] = i end f2(x) xh = Array(x) @test all(xh .== 1:length(xh)) end + + # A range does not determine the backend: pass the backend of the arrays the loop accesses + x = array_from_host(zeros(Int, 1000)) + f3(x) = AK.foreachindex(1:500; backend=AK.get_backend(x)) do i + x[i] = i + end + f3(x) + @test Array(x) == [1:500; zeros(Int, 500)] + # ... without one, it runs on the host + y = zeros(Int, 10) + AK.foreachindex(i -> (y[i] = i), 1:10) + @test y == 1:10 + + # Invalid launch settings, on every backend + for kw in ((block_size=0,), (max_tasks=0,), (min_elems=0,)) + @test_throws ArgumentError AK.foreachindex(i -> nothing, x; kw...) + end + @test_throws TypeError AK.foreachindex(i -> nothing, x; backend=:gpu) + @test_throws MethodError AK.foreachindex(i -> nothing, x; prefer_threads=true) + @test_throws MethodError AK.foreachindex(i -> nothing, x, AK.get_backend(x)) end @@ -59,12 +79,12 @@ end end x = array_from_host(zeros(Int, 10, 1000)) - f1(x; prefer_threads) + f1(x) xh = Array(x) @test all(xh .== (1:10) .+ (1:1000)') x = array_from_host(zeros(UInt32, 10, 1000)) - f1(x; prefer_threads, max_tasks=2, min_elems=100, block_size=64) + f1(x; max_tasks=2, min_elems=100, block_size=64) xh = Array(x) @test all(xh .== (1:10) .+ (1:1000)') @@ -75,12 +95,12 @@ end end x = array_from_host(zeros(Int, 10, 1000)) - f2(x; prefer_threads) + f2(x) xh = Array(x) @test all(xh .== (1:10) .+ (1:1000)') x = array_from_host(zeros(UInt32, 10, 1000)) - f2(x; prefer_threads, max_tasks=2, min_elems=100, block_size=64) + f2(x; max_tasks=2, min_elems=100, block_size=64) xh = Array(x) @test all(xh .== (1:10) .+ (1:1000)') @@ -90,7 +110,7 @@ end end x = array_from_host(zeros(Int, 10, 1000)) - f3(x; prefer_threads) + f3(x) xh = Array(x) @test all(xh[:] .== 1:length(x)) end diff --git a/test/generic/map.jl b/test/generic/map.jl index c673fe28..7f3c3ff3 100644 --- a/test/generic/map.jl +++ b/test/generic/map.jl @@ -2,33 +2,47 @@ Random.seed!(0) x = array_from_host(1:1000) - y = AK.map(x; prefer_threads) do i + y = AK.map(x) do i i^2 end @test Array(y) == map(i -> i^2, 1:1000) x = array_from_host(1:1000) y = array_from_host(zeros(Int, 1000)) - AK.map!(y, x; prefer_threads) do i + AK.map!(y, x) do i i^2 end @test Array(y) == map(i -> i^2, 1:1000) x = array_from_host(rand(Float32, 1000)) # Tests different things with GPU and CPU backends as well as irrelevant parameters being ignored - y = AK.map(x; prefer_threads, block_size=64, max_tasks=2, min_elems=100) do i + y = AK.map(x; block_size=64, max_tasks=2, min_elems=100) do i i > 0.5 ? i : 0 end @test Array(y) == map(i -> i > 0.5 ? i : 0, Array(x)) - if prefer_threads # CPU only + if AK._runs_threads(BACKEND) # host arrays x = rand(Float32, 1000) - y = AK.map(x; prefer_threads, max_tasks=4, min_elems=500) do i + y = AK.map(x; max_tasks=4, min_elems=500) do i i > 0.5 ? i : 0 end @test y == map(i -> i > 0.5 ? i : 0, x) end + # map! returns its destination, which must have as many elements as the source + x = array_from_host(1:1000) + y = array_from_host(zeros(Int, 1000)) + @test AK.map!(i -> i + 1, y, x) === y + @test_throws ArgumentError AK.map!(identity, array_from_host(zeros(Int, 3)), x) + # Test that undefined kwargs are not accepted - @test_throws MethodError AK.map(x -> x^2, x; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.map(x -> x^2, x; bad=:kwarg) + @test_throws MethodError AK.map(x -> x^2, x; prefer_threads=true) +end + +@testset "map: backend-free inputs" begin + # A range has no backend: the result is allocated on the one given + r = AK.map(x -> 2x, 1:5; backend=BACKEND) + @test get_backend(r) == BACKEND + @test Array(r) == 2:2:10 end diff --git a/test/generic/reverse.jl b/test/generic/reverse.jl index 4a23e36b..8cd39171 100644 --- a/test/generic/reverse.jl +++ b/test/generic/reverse.jl @@ -14,7 +14,7 @@ for T in test_types, n in edge_sizes h = rand(T, n) v = array_from_host(h) - AK.reverse!(v; prefer_threads) + AK.reverse!(v) @test Array(v) == reverse(h) end @@ -22,14 +22,14 @@ for _ in 1:50 h = rand(Float32, rand(1:100_000)) v = array_from_host(h) - AK.reverse!(v; prefer_threads) - AK.reverse!(v; prefer_threads) + AK.reverse!(v) + AK.reverse!(v) @test Array(v) == h end # Returns the same array it was given, not a copy v = array_from_host(rand(Float32, 1000)) - @test AK.reverse!(v; prefer_threads) === v + @test AK.reverse!(v) === v end @testset "reverse! out-of-place" begin @@ -37,23 +37,21 @@ h = rand(T, n) src = array_from_host(h) dst = array_from_host(zeros(T, n)) - AK.reverse!(dst, src; prefer_threads) + AK.reverse!(dst, src) @test Array(dst) == reverse(h) @test Array(src) == h # source left untouched end @test_throws Exception AK.reverse!( array_from_host(rand(Float32, 10)), - array_from_host(rand(Float32, 11)); - prefer_threads, - ) + array_from_host(rand(Float32, 11))) end @testset "reverse allocating" begin for T in test_types, n in edge_sizes h = rand(T, n) v = array_from_host(h) - out = AK.reverse(v; prefer_threads) + out = AK.reverse(v) @test Array(out) == reverse(h) @test Array(v) == h # source left untouched @test out !== v @@ -67,12 +65,12 @@ h = rand(Float32, n) v = array_from_host(h) - AK.reverse!(v; prefer_threads) + AK.reverse!(v) @test Array(v) == reverse(h) src = array_from_host(h) dst = array_from_host(zeros(Float32, n)) - AK.reverse!(dst, src; prefer_threads) + AK.reverse!(dst, src) @test Array(dst) == reverse(h) end end @@ -82,12 +80,12 @@ h = rand(Float32, 10_000) for block_size in (32, 64, 128, 256) v = array_from_host(h) - AK.reverse!(v; prefer_threads, block_size) + AK.reverse!(v; block_size) @test Array(v) == reverse(h) end for (max_tasks, min_elems) in ((1, 1), (2, 100), (4, 1000)) v = array_from_host(h) - AK.reverse!(v; prefer_threads, max_tasks, min_elems) + AK.reverse!(v; max_tasks, min_elems) @test Array(v) == reverse(h) end end @@ -101,16 +99,16 @@ h = rand(Float32, shape...) v = array_from_host(h) - AK.reverse!(v; dims=dim, prefer_threads) + AK.reverse!(v; dims=dim) @test Array(v) == reverse(h; dims=dim) src = array_from_host(h) - out = AK.reverse(src; dims=dim, prefer_threads) + out = AK.reverse(src; dims=dim) @test Array(out) == reverse(h; dims=dim) @test Array(src) == h # source left untouched dst = array_from_host(zeros(Float32, shape...)) - AK.reverse!(dst, src; dims=dim, prefer_threads) + AK.reverse!(dst, src; dims=dim) @test Array(dst) == reverse(h; dims=dim) end @@ -123,32 +121,32 @@ h = rand(Float32, shape...) v = array_from_host(h) - AK.reverse!(v; dims=dims, prefer_threads) + AK.reverse!(v; dims=dims) @test Array(v) == reverse(h; dims=dims) - out = AK.reverse(array_from_host(h); dims=dims, prefer_threads) + out = AK.reverse(array_from_host(h); dims=dims) @test Array(out) == reverse(h; dims=dims) src = array_from_host(h) dst = array_from_host(zeros(Float32, shape...)) - AK.reverse!(dst, src; dims=dims, prefer_threads) + AK.reverse!(dst, src; dims=dims) @test Array(dst) == reverse(h; dims=dims) end # Any iterable of integers works, e.g. a Vector (Base only accepts tuples) h = rand(Float32, 4, 5, 6) - out = AK.reverse(array_from_host(h); dims=[1, 3], prefer_threads) + out = AK.reverse(array_from_host(h); dims=[1, 3]) @test Array(out) == reverse(h; dims=(1, 3)) # Stateful iterators must survive validation, including duplicate detection. expected = reverse(h; dims=(1, 3)) v = array_from_host(h) - @test AK.reverse!(v; dims=Iterators.Stateful([1, 3]), prefer_threads) === v + @test AK.reverse!(v; dims=Iterators.Stateful([1, 3])) === v @test Array(v) == expected src = array_from_host(h) - @test Array(AK.reverse(src; dims=Iterators.Stateful([1, 3]), prefer_threads)) == expected + @test Array(AK.reverse(src; dims=Iterators.Stateful([1, 3]))) == expected dst = similar(src) - @test AK.reverse!(dst, src; dims=Iterators.Stateful([1, 3]), prefer_threads) === dst + @test AK.reverse!(dst, src; dims=Iterators.Stateful([1, 3])) === dst @test Array(dst) == expected @test Array(src) == h @@ -157,15 +155,15 @@ src_int = array_from_host(h_int) for dims in (:, (), 1, (1, 2)) dst_float = array_from_host(zeros(Float32, size(h_int))) - @test AK.reverse!(dst_float, src_int; dims, prefer_threads) === dst_float + @test AK.reverse!(dst_float, src_int; dims) === dst_float @test Array(dst_float) == reverse(h_int; dims) end @test Array(src_int) == h_int # dims=() reverses nothing v = array_from_host(h) - @test Array(AK.reverse!(v; dims=(), prefer_threads)) == h - @test Array(AK.reverse(v; dims=(), prefer_threads)) == h + @test Array(AK.reverse!(v; dims=())) == h + @test Array(AK.reverse(v; dims=())) == h # Shapes spanning many blocks, with odd extents so the in-place middle slice is # non-trivial, for integer element types too @@ -173,10 +171,10 @@ h = rand(T, shape) for dims in (1, 2, (1, 2)), block_size in (64, 256) v = array_from_host(h) - AK.reverse!(v; dims, prefer_threads, block_size) + AK.reverse!(v; dims, block_size) @test Array(v) == reverse(h; dims) - out = AK.reverse(array_from_host(h); dims, prefer_threads, block_size) + out = AK.reverse(array_from_host(h); dims, block_size) @test Array(out) == reverse(h; dims) end end @@ -185,34 +183,54 @@ h = zeros(Float32, 0, 5) for dims in (1, 2, (1, 2)) v = array_from_host(h) - @test Array(AK.reverse!(v; dims, prefer_threads)) == reverse(h; dims) + @test Array(AK.reverse!(v; dims)) == reverse(h; dims) dst = array_from_host(copy(h)) - @test Array(AK.reverse!(dst, v; dims, prefer_threads)) == reverse(h; dims) + @test Array(AK.reverse!(dst, v; dims)) == reverse(h; dims) - @test Array(AK.reverse(v; dims, prefer_threads)) == reverse(h; dims) + @test Array(AK.reverse(v; dims)) == reverse(h; dims) end end # Invalid dims arguments throw, matching Base/CUDA @testset "dims errors" begin v = array_from_host(rand(Float32, 2, 3, 4)) - @test_throws ArgumentError AK.reverse!(v; dims=0, prefer_threads) - @test_throws ArgumentError AK.reverse!(v; dims=4, prefer_threads) - @test_throws ArgumentError AK.reverse(v; dims=0, prefer_threads) - @test_throws ArgumentError AK.reverse(v; dims=4, prefer_threads) + @test_throws ArgumentError AK.reverse!(v; dims=0) + @test_throws ArgumentError AK.reverse!(v; dims=4) + @test_throws ArgumentError AK.reverse(v; dims=0) + @test_throws ArgumentError AK.reverse(v; dims=4) # Non-integer dims must throw rather than silently do nothing - @test_throws ArgumentError AK.reverse!(v; dims=1.5, prefer_threads) - @test_throws ArgumentError AK.reverse(v; dims=(1, 2.5), prefer_threads) + @test_throws ArgumentError AK.reverse!(v; dims=1.5) + @test_throws ArgumentError AK.reverse(v; dims=(1, 2.5)) dst = similar(v) - @test_throws ArgumentError AK.reverse!(dst, v; dims=Iterators.Stateful([1, 1]), prefer_threads) - @test_throws ArgumentError AK.reverse!(v; dims=nothing, prefer_threads) - @test_throws ArgumentError AK.reverse!(similar(v, 4, 3, 2), v; dims=1, prefer_threads) + @test_throws ArgumentError AK.reverse!(dst, v; dims=Iterators.Stateful([1, 1])) + @test_throws ArgumentError AK.reverse!(v; dims=nothing) + @test_throws ArgumentError AK.reverse!(similar(v, 4, 3, 2), v; dims=1) # Duplicate dims throw, as in Base - @test_throws ArgumentError AK.reverse!(v; dims=(1, 1), prefer_threads) - @test_throws ArgumentError AK.reverse(v; dims=[2, 3, 2], prefer_threads) + @test_throws ArgumentError AK.reverse!(v; dims=(1, 1)) + @test_throws ArgumentError AK.reverse(v; dims=[2, 3, 2]) end + + # Invalid launch settings throw on every backend, also where nothing needs to move + @testset "launch setting errors" begin + for x in (array_from_host(rand(Int32, 10)), array_from_host(Int32[1]), + array_from_host(Int32[])) + @test_throws ArgumentError AK.reverse!(x; block_size=0) + @test_throws ArgumentError AK.reverse!(x; max_tasks=0) + @test_throws ArgumentError AK.reverse!(similar(x), x; min_elems=0) + @test_throws ArgumentError AK.reverse(x; block_size=0) + end + m = array_from_host(rand(Int32, 2, 3)) + @test_throws ArgumentError AK.reverse!(m; dims=(), max_tasks=0) + @test_throws ArgumentError AK.reverse!(similar(m), m; dims=1, min_elems=0) + end +end + +@testset "reverse: backend-free inputs" begin + r = AK.reverse(1:5; backend=BACKEND) + @test get_backend(r) == BACKEND + @test Array(r) == 5:-1:1 end From e6de24d4c03973eb3fcd3145cd564597814b03e8 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 23 Sep 2026 18:13:58 +0200 Subject: [PATCH 3/8] Move reductions onto algorithm selection `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. --- docs/src/api/algorithms.md | 11 +- docs/src/api/reduce.md | 22 ++ docs/src/api/using_backends.md | 8 +- docs/src/tuning.md | 9 + src/algorithms.jl | 87 +++++++- src/arithmetics.jl | 260 +++--------------------- src/predicates.jl | 16 +- src/reduce/mapreduce_1d_cpu.jl | 7 - src/reduce/mapreduce_1d_gpu.jl | 6 - src/reduce/mapreduce_nd.jl | 48 ++--- src/reduce/reduce.jl | 286 ++++++-------------------- src/reduce/tuning.jl | 116 +++++++++++ src/reduce/utilities.jl | 10 +- src/sort/radix_sort.jl | 3 +- src/sort/tuning.jl | 42 +--- test/generic/reduce.jl | 359 +++++++++++++++++---------------- test/generic/reduce_resolve.jl | 158 +++++++++++++++ test/generic/sort_resolve.jl | 5 + 18 files changed, 708 insertions(+), 745 deletions(-) create mode 100644 src/reduce/tuning.jl create mode 100644 test/generic/reduce_resolve.jl diff --git a/docs/src/api/algorithms.md b/docs/src/api/algorithms.md index 16973847..09bdc288 100644 --- a/docs/src/api/algorithms.md +++ b/docs/src/api/algorithms.md @@ -11,6 +11,9 @@ AcceleratedKernels.SortAlgorithm AcceleratedKernels.CPUThreads ``` +The algorithm types of each operation are documented with it, for [sorting](sort.md) and for +[reductions](reduce.md). + #### How `Auto` chooses `Auto` looks at the backend and its current device, the element type, the layout (a whole array or @@ -18,7 +21,7 @@ slices along `dims`, and their length) and the ordering, never at the array's co host backend it chooses the family's `CPUThreads` algorithm, which runs on Julia threads. On a GPU it follows per-device tuning values: for sorting, `BitonicSort` for short arrays and slices, `RadixSort` for long whole arrays of the element types it supports, and `MergeSort` otherwise, -subject to `Auto`'s requirements (`stable=true` by default). +subject to `Auto`'s requirements (`stable=true` by default); for reductions, `BlockReduce`. Algorithms carry their settings as fields, and a field left at `nothing` takes the device's tuned value: @@ -34,9 +37,9 @@ element type, the ordering or the backend cannot run, or an invalid setting, is #### The `backend` keyword Operations run on the backend of their arrays: `backend` is derived from every array argument, -the destination first, and all of them must agree. Ranges, `CartesianIndices`, `LinearIndices`, -numbers and other non-array arguments do not count. If no argument determines the backend (e.g. -a loop over a range), the operation runs on the host. +the destination first, and all of them must agree. Ranges, `CartesianIndices`, `LinearIndices` +(also through views and reshapes), numbers and other non-array arguments do not count. If no +argument determines the backend (e.g. a loop over a range), the operation runs on the host. Pass `backend` explicitly for arrays that cannot tell (a range to be processed on a GPU, for instance), or for memory that several backends can access. AcceleratedKernels does not check that diff --git a/docs/src/api/reduce.md b/docs/src/api/reduce.md index b4d6a5ae..59304d54 100644 --- a/docs/src/api/reduce.md +++ b/docs/src/api/reduce.md @@ -8,3 +8,25 @@ Apply a custom binary operator reduction on all elements in an iterable; can be ```@docs AcceleratedKernels.reduce ``` + +The operator must be associative and commutative: like every GPU reduction, AcceleratedKernels +combines elements in an order that depends on the algorithm and its settings, not in element +order. + +[`Auto()`](@ref AcceleratedKernels.Auto), the default `alg`, uses +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host and +[`BlockReduce`](@ref AcceleratedKernels.BlockReduce) on GPUs, with the device's settings. Pass an +algorithm to choose its settings yourself: + +```@docs +AcceleratedKernels.ReduceAlgorithm +AcceleratedKernels.BlockReduce +AcceleratedKernels.CPUThreads.Partitioned +``` + +```julia +m = ROCArray(rand(Float32, 1000, 1000)) +AK.reduce(+, m; init=0f0, alg=AK.BlockReduce(block_size=512)) # 512 threads per block +AK.reduce(+, m; init=0f0, alg=AK.BlockReduce(switch_below=1024)) # finish on the host below 1024 values +AK.reduce(+, Array(m); init=0f0, alg=AK.CPUThreads.Partitioned(max_tasks=4)) +``` diff --git a/docs/src/api/using_backends.md b/docs/src/api/using_backends.md index fad7a626..936889eb 100644 --- a/docs/src/api/using_backends.md +++ b/docs/src/api/using_backends.md @@ -22,12 +22,14 @@ v = CuArray{UInt32}(0:5:100_000) # Range with explicit step size v_host = Array(v) ``` -All publicly-exposed functions have CPU implementations with unified parameter interfaces: +All publicly-exposed functions also run on host arrays, with the same interface: ```julia import AcceleratedKernels as AK v = Vector(-1000:1000) # Normal CPU array -AK.reduce(+, v, max_tasks=Threads.nthreads()) +AK.reduce(+, v; init=0) +AK.reduce(+, v; init=0, alg=AK.CPUThreads.Partitioned(max_tasks=4)) # at most 4 tasks ``` -By default all algorithms use the number of threads Julia was started with. +On the host, operations run on Julia threads and by default use as many tasks as Julia has +threads. See [Algorithms and backends](@ref) for the algorithms and their settings. diff --git a/docs/src/tuning.md b/docs/src/tuning.md index 3cfcd14c..3da43486 100644 --- a/docs/src/tuning.md +++ b/docs/src/tuning.md @@ -20,6 +20,13 @@ sorting, `_resolve_sort` runs four steps: AcceleratedKernels._resolve_sort ``` +The other families follow the same steps, with a one-line selection: reductions resolve with +`_resolve_reduce`. + +```@docs +AcceleratedKernels._resolve_reduce +``` + `Auto` and explicit algorithms share steps 3 and 4, so a tuning cannot make an invalid algorithm run, and an explicit setting always wins over the tuning. @@ -31,6 +38,8 @@ fields, and one hook returns it for a backend and element type: ```@docs AcceleratedKernels.SortTuning AcceleratedKernels.sort_tuning +AcceleratedKernels.ReduceTuning +AcceleratedKernels.reduce_tuning ``` AcceleratedKernels defines the hook's generic method, whose values reproduce the library's diff --git a/src/algorithms.jl b/src/algorithms.jl index 8fdc07c3..03235598 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -44,6 +44,14 @@ Supertype of the sorting algorithms: [`MergeSort`](@ref), [`RadixSort`](@ref), """ abstract type SortAlgorithm <: Algorithm end +""" + ReduceAlgorithm <: Algorithm + +Supertype of the reduction algorithms: [`BlockReduce`](@ref). Reductions also accept +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned). +""" +abstract type ReduceAlgorithm <: Algorithm end + """ AcceleratedKernels.CPUThreads @@ -54,7 +62,7 @@ other backend. """ module CPUThreads -import ..AcceleratedKernels: SortAlgorithm +import ..AcceleratedKernels: Algorithm, SortAlgorithm """ CPUThreads.SampleSort(; max_tasks=nothing, min_elems=nothing) @@ -69,6 +77,19 @@ Base.@kwdef struct SampleSort <: SortAlgorithm min_elems::Union{Nothing, Int} = nothing end +""" + CPUThreads.Partitioned(; max_tasks=nothing, min_elems=nothing) + +Split the input into contiguous parts, one per task, process them on Julia threads, and combine +the parts' results. Used by reductions, scans, `findall` and `any`/`all` on host arrays. Uses at +most `max_tasks` tasks (default `Threads.nthreads()`), each with at least `min_elems` elements +(default 1, or the operation's tuning). Only runs on the host backend. +""" +Base.@kwdef struct Partitioned <: Algorithm + max_tasks::Union{Nothing, Int} = nothing + min_elems::Union{Nothing, Int} = nothing +end + end # module CPUThreads @@ -99,16 +120,69 @@ _runs_kernels(::Backend) = true _runs_kernels(::HostBackend) = nameof(HostBackend) !== :CPU +# Checks shared by the algorithm families + +# Algorithm names in error messages +_algname(a) = nameof(typeof(a)) +_algname(::CPUThreads.SampleSort) = "CPUThreads.SampleSort" +_algname(::CPUThreads.Partitioned) = "CPUThreads.Partitioned" + +# Domain checks of the fields an algorithm was given explicitly, before any arithmetic uses them +_checkdomain(::Algorithm) = nothing + +function _check_positive(a, field) + x = getfield(a, field) + x === nothing || x >= 1 || + throw(ArgumentError("$(_algname(a)): `$field` must be positive, got $x")) + nothing +end + +function _check_pow2(a, field) + x = getfield(a, field) + x === nothing || (x >= 1 && ispow2(x)) || + throw(ArgumentError("$(_algname(a)): `$field` must be a positive power of two, got $x")) + nothing +end + +function _check_threads(a) + _check_positive(a, :max_tasks) + _check_positive(a, :min_elems) +end + +_checkdomain(a::CPUThreads.Partitioned) = _check_threads(a) + +# Unset fields of a threaded algorithm: all threads, and the tuning's minimum per task +_fill_threads(a::A, t) where {A} = + A(something(a.max_tasks, Threads.nthreads()), something(a.min_elems, t.threads_min_elems)) + +function _require_kernels(a, backend) + _runs_kernels(backend) || throw(ArgumentError( + "$(_algname(a)) runs AcceleratedKernels' GPU kernels, which " * + "$(_backend_name(backend)) cannot run (on the host, this needs KernelAbstractions 0.10); " * + "use `Auto()` or a `CPUThreads` algorithm")) + nothing +end + +function _require_threads(a, backend) + _runs_threads(backend) || throw(ArgumentError( + "$(_algname(a)) only runs on the host backend, not on $(_backend_name(backend))")) + nothing +end + + # Backend resolution _backend_name(b::Backend) = nameof(typeof(b)) -# Values that never determine the backend: lazy index collections, scalars, and any other value -# that is not an array. Every other array votes with `get_backend`, which throws for array types -# that do not implement it. +# Values that never determine the backend: lazy index collections (and Base's wrappers of them), +# scalars, and any other value that is not an array. Every other array votes with `get_backend`, +# which throws for array types that do not implement it. _backend_vote(_) = nothing _backend_vote(x::AbstractArray) = get_backend(x) _backend_vote(::Union{AbstractRange, CartesianIndices, LinearIndices}) = nothing +# Base's views, reshapes and permutations vote like the array they wrap, so that e.g. a reshaped +# range does not determine the backend either +_backend_vote(x::Union{SubArray, Base.ReshapedArray, PermutedDimsArray}) = _backend_vote(parent(x)) _backend_vote(x::Tuple) = _backend_votes(x...) _backend_vote(bc::Base.Broadcast.Broadcasted) = _backend_votes(bc.args...) _backend_vote(x::Base.Broadcast.Extruded) = _backend_vote(x.x) @@ -145,8 +219,9 @@ _copy(backend, v) = _backend_free(v) ? copyto!(_similar(backend, v), collect(v)) The backend an operation runs on: `backend` if given, else the one every array in `args` (destination first) agrees on, recursing into `Broadcasted` trees; ranges, `CartesianIndices`, -`LinearIndices` and non-array values do not count. Arguments on different backends are an -`ArgumentError`. If no argument determines it, the host backend is used. +`LinearIndices`, Base's views, reshapes and permutations of them, and non-array values do not +count. Arguments on different backends are an `ArgumentError`. If no argument determines it, the +host backend is used. """ _resolve_backend(backend::Backend, args...) = backend function _resolve_backend(::Nothing, args...) diff --git a/src/arithmetics.jl b/src/arithmetics.jl index 16f98761..d9211845 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -1,301 +1,101 @@ """ - sum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) + sum(src::AbstractArray; init=zero(eltype(src)), kwargs...) -Sum of elements of an array, with optional `init` and `dims`. Arguments are the same as for +Sum of the elements of an array, with optional `init` and `dims`. The keywords are those of [`reduce`](@ref). -# Examples -Simple sum of elements in a vector: ```julia import AcceleratedKernels as AK using Metal v = MtlArray(rand(Int32(1):Int32(100), 100_000)) s = AK.sum(v) -``` -Row-wise sum of a matrix: -```julia -m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) -s = AK.sum(m, dims=1) -``` - -If you know the shape of the resulting array (in case of a dimensionwise sum, i.e. `dims` is not -`nothing`), you can provide the `temp` argument to save results into and avoid allocations: -```julia m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) -temp = MtlArray(zeros(Int32, 10, 1)) -s = AK.sum(m, dims=2, temp=temp) +s = AK.sum(m; dims=1) # row-wise +s = AK.sum(m; dims=2, temp=MtlArray(zeros(Int32, 10, 1))) # into a preallocated result ``` """ -function sum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - kwargs... -) - reduce( - +, src, backend; - init, - kwargs... - ) -end +sum(src::AbstractArray; init=zero(eltype(src)), kwargs...) = reduce(+, src; init, kwargs...) """ - prod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) + prod(src::AbstractArray; init=one(eltype(src)), kwargs...) -Product of elements of an array, with optional `init` and `dims`. Arguments are the same as for +Product of the elements of an array, with optional `init` and `dims`. The keywords are those of [`reduce`](@ref). -# Examples -Simple product of elements in a vector: ```julia import AcceleratedKernels as AK using AMDGPU v = ROCArray(rand(Int32(1):Int32(100), 100_000)) p = AK.prod(v) -``` - -Row-wise product of a matrix: -```julia -m = ROCArray(rand(Int32(1):Int32(100), 10, 100_000)) -p = AK.prod(m, dims=1) -``` - -If you know the shape of the resulting array (in case of a dimensionwise product, i.e. `dims` is not -`nothing`), you can provide the `temp` argument to save results into and avoid allocations: -```julia -m = ROCArray(rand(Int32(1):Int32(100), 10, 100_000)) -temp = ROCArray(ones(Int32, 10, 1)) -p = AK.prod(m, dims=2, temp=temp) +p = AK.prod(ROCArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -function prod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - kwargs... -) - reduce( - *, src, backend; - init, - kwargs... - ) -end +prod(src::AbstractArray; init=one(eltype(src)), kwargs...) = reduce(*, src; init, kwargs...) """ - maximum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemin(eltype(src)), - dims=nothing, + maximum(src::AbstractArray; init=typemin(eltype(src)), kwargs...) - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) - -Maximum of elements of an array, with optional `init` and `dims`. Arguments are the same as for +Maximum of the elements of an array, with optional `init` and `dims`. The keywords are those of [`reduce`](@ref). -# Examples -Simple maximum of elements in a vector: ```julia import AcceleratedKernels as AK using oneAPI v = oneArray(rand(Int32(1):Int32(100), 100_000)) m = AK.maximum(v) -``` - -Row-wise maximum of a matrix: -```julia -m = oneArray(rand(Int32(1):Int32(100), 10, 100_000)) -m = AK.maximum(m, dims=1) -``` - -If you know the shape of the resulting array (in case of a dimensionwise maximum, i.e. `dims` is not -`nothing`), you can provide the `temp` argument to save results into and avoid allocations: -```julia -m = oneArray(rand(Int32(1):Int32(100), 10, 100_000)) -temp = oneArray(zeros(Int32, 10, 1)) -m = AK.maximum(m, dims=2, temp=temp) +m = AK.maximum(oneArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -function maximum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemin(eltype(src)), - kwargs... -) - reduce( - max, src, backend; - init, - kwargs... - ) -end +maximum(src::AbstractArray; init=typemin(eltype(src)), kwargs...) = reduce(max, src; init, kwargs...) """ - minimum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemax(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) + minimum(src::AbstractArray; init=typemax(eltype(src)), kwargs...) -Minimum of elements of an array, with optional `init` and `dims`. Arguments are the same as for +Minimum of the elements of an array, with optional `init` and `dims`. The keywords are those of [`reduce`](@ref). -# Examples -Simple minimum of elements in a vector: ```julia import AcceleratedKernels as AK using CUDA v = CuArray(rand(Int32(1):Int32(100), 100_000)) m = AK.minimum(v) -``` - -Row-wise minimum of a matrix: -```julia -m = CuArray(rand(Int32(1):Int32(100), 10, 100_000)) -m = AK.minimum(m, dims=1) -``` - -If you know the shape of the resulting array (in case of a dimensionwise minimum, i.e. `dims` is not -`nothing`), you can provide the `temp` argument to save results into and avoid allocations: -```julia -m = CuArray(rand(Int32(1):Int32(100), 10, 100_000)) -temp = CuArray(ones(Int32, 10, 1)) -m = AK.minimum(m, dims=2, temp=temp) +m = AK.minimum(CuArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -function minimum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemax(eltype(src)), - kwargs... -) - reduce( - min, src, backend; - init, - kwargs... - ) -end +minimum(src::AbstractArray; init=typemax(eltype(src)), kwargs...) = reduce(min, src; init, kwargs...) """ - count( - [f=identity], src::AbstractArray, backend::Backend=get_backend(src); - init=0, - dims=nothing, + count([f=identity,] src::AbstractArray; init=0, kwargs...) - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, +Count the elements of `src` for which `f` returns `true`, with optional `init` and `dims`; the +result has the type of `init`. The keywords are those of [`mapreduce`](@ref). - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) - -Count the number of elements in `src` for which the function `f` returns `true`. If `f` is omitted, -count the number of `true` elements in `src`. Arguments are the same as for [`mapreduce`](@ref). - -# Examples -Simple count of `true` elements in a vector: ```julia import AcceleratedKernels as AK -using Metal - -v = MtlArray(rand(Bool, 100_000)) -c = AK.count(v) -``` - -Count of elements greater than 50 in a vector: -```julia -v = MtlArray(rand(Int32(1):Int32(100), 100_000)) -c = AK.count(x -> x > 50, v) -``` - -Row-wise count of `true` elements in a matrix: -```julia -m = MtlArray(rand(Bool, 10, 100_000)) -c = AK.count(m, dims=1) -``` +using CUDA -If you know the shape of the resulting array (in case of a dimensionwise count, i.e. `dims` is not -`nothing`), you can provide the `temp` argument to save results into and avoid allocations: -```julia -m = MtlArray(rand(Bool, 10, 100_000)) -temp = MtlArray(zeros(Int32, 10, 1)) -c = AK.count(m; init=Int32(0), dims=2, temp=temp) +v = CuArray(rand(Float32, 100_000)) +c = AK.count(x -> x > 0.5, v) +c = AK.count(CuArray(rand(Bool, 10, 100_000)); init=Int32(0), dims=2) ``` """ -function count( - src::AbstractArray, backend::Backend=get_backend(src); - init=0, - kwargs... -) - mapreduce( - x -> x ? one(typeof(init)) : zero(typeof(init)), +, src, backend; - init, - neutral=zero(typeof(init)), - kwargs... - ) +function count(src::AbstractArray; init=0, kwargs...) + mapreduce(x -> x ? one(typeof(init)) : zero(typeof(init)), +, src; + init, neutral=zero(typeof(init)), kwargs...) end - -function count( - f, src::AbstractArray, backend::Backend=get_backend(src); - init=0, - kwargs... -) - mapreduce( - x -> f(x) ? one(typeof(init)) : zero(typeof(init)), +, src, backend; - init, - neutral=zero(typeof(init)), - kwargs... - ) +function count(f, src::AbstractArray; init=0, kwargs...) + mapreduce(x -> f(x) ? one(typeof(init)) : zero(typeof(init)), +, src; + init, neutral=zero(typeof(init)), kwargs...) end diff --git a/src/predicates.jl b/src/predicates.jl index cfd63f86..f02e480d 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -134,14 +134,12 @@ function _any_impl( return mapreduce( pred, (x, y) -> x || y, - v, - backend; + v; + backend, init=false, neutral=false, - prefer_threads=true, - block_size, + alg=BlockReduce(; block_size, switch_below=alg.switch_below), temp=alg.temp, - switch_below=alg.switch_below, ) end else @@ -268,14 +266,12 @@ function _all_impl( return mapreduce( pred, (x, y) -> x && y, - v, - backend; + v; + backend, init=true, neutral=true, - prefer_threads=false, - block_size, + alg=BlockReduce(; block_size, switch_below=alg.switch_below), temp=alg.temp, - switch_below=alg.switch_below, ) end else diff --git a/src/reduce/mapreduce_1d_cpu.jl b/src/reduce/mapreduce_1d_cpu.jl index 3a47b603..dc59be3e 100644 --- a/src/reduce/mapreduce_1d_cpu.jl +++ b/src/reduce/mapreduce_1d_cpu.jl @@ -2,15 +2,8 @@ function mapreduce_1d_cpu( f, op, src::MapReduceSource, backend::Backend; init, neutral, - - # CPU settings max_tasks::Int, min_elems::Int, - - # GPU settings - ignored here - block_size::Int, - temp::Union{Nothing, AbstractArray}, - switch_below::Int, ) if src isa Base.Broadcast.Broadcasted return op(init, Base.mapreduce(f, op, src; init=neutral)) diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index cfc81c0b..be90bb7f 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -40,12 +40,6 @@ function mapreduce_1d_gpu( f, op, src::MapReduceSource, backend::Backend; init, neutral, - - # CPU settings - ignored here - max_tasks::Int, - min_elems::Int, - - # GPU settings block_size::Int, items_per_thread::Int, temp::Union{Nothing, AbstractArray}, diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index f2b6c696..cf06c35b 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -15,10 +15,6 @@ # - by_block: one block per output, grid-stride over outputs (few outputs, large reduction) # - multigroup: several blocks per output, two-pass (dst_size==1 or very small dst_size) -# Number of blocks the by_block / multigroup paths aim to launch, so a reduction with -# few output elements can still fill the GPU. A heuristic GPU-occupancy target. -const TARGET_BLOCKS = 256 - # Below this many output elements, splitting a single output's reduction across multiple # blocks (multigroup) is preferred over grid-striding by_block, because grid-stride with # too few blocks cannot fill the GPU on its own. At or above this, by_block grid-strides. @@ -74,22 +70,22 @@ end # Main entry point function mapreduce_nd( - f, op, src::MapReduceSource, backend::Backend; + f, op, src::MapReduceSource, backend::Backend, alg::Union{BlockReduce, CPUThreads.Partitioned}; init, neutral=neutral_element(op, typeof(init)), dims, - - # CPU settings - max_tasks::Int, - min_elems::Int, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int, temp::Union{Nothing, AbstractArray}, ) - @argcheck 1 <= block_size <= 1024 - @argcheck ispow2(block_size) + # Launch settings of the elementwise passes, and of the reduction kernels + launch = alg isa BlockReduce ? (; block_size=alg.block_size) : + (; max_tasks=alg.max_tasks, min_elems=alg.min_elems) + block_size = alg isa BlockReduce ? alg.block_size : 256 + + # Number of blocks the by_block / multigroup paths aim to launch, so a reduction with + # few output elements can still fill the GPU. A heuristic GPU-occupancy target. + target_blocks = reduce_tuning(backend, typeof(init)).target_blocks + target_blocks >= 1 || throw(ArgumentError( + "the reduction tuning's `target_blocks` must be positive, got $target_blocks")) dims_src = dims isa Number ? (dims,) : dims dims_buf = Int[] @@ -115,7 +111,7 @@ function mapreduce_nd( # julia> mapreduce(x -> -x, +, x, dims=3, init=Float32(0)) # 3×5 Matrix{Float32} if isempty(dims_valid) dst = _alloc_or_temp(backend, temp, init, src_sizes) - _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, max_tasks, min_elems, block_size) + _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) return dst end @@ -146,7 +142,7 @@ function mapreduce_nd( # If the reduced extent is 1, just map each element through f (keep init's type) if len == 1 dst = _alloc_or_temp(backend, temp, init, src_sizes) - _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, max_tasks, min_elems, block_size) + _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) return dst end @@ -155,8 +151,8 @@ function mapreduce_nd( dst = _alloc_or_temp(backend, temp, init, dst_sizes) dst_size = length(dst) - if backend == CPU_BACKEND - _mapreduce_nd_cpu_sections!(f, op, dst, src; init, max_tasks, min_elems) + if alg isa CPUThreads.Partitioned + _mapreduce_nd_cpu_sections!(f, op, dst, src; init, launch...) return dst end @@ -200,7 +196,7 @@ function mapreduce_nd( # - otherwise (GS_DST_CUTOFF <= dst_size < reduce_size) -> by_block, grid-striding # over outputs, single pass # - # Rationale: grid-striding by_block launches `min(dst_size, TARGET_BLOCKS)` blocks; + # Rationale: grid-striding by_block launches `min(dst_size, target_blocks)` blocks; # for very small dst_size (e.g. 5 or 9) that under-fills an 84-SM GPU, so splitting # the (large) reduction itself across many blocks via multigroup is still better. # For dst_size==1 there is nothing to grid-stride over, so multigroup is the only @@ -223,7 +219,7 @@ function mapreduce_nd( # coalescing is better than a strided block reduction. use_by_block_for_low_occupancy = dst_size == reduce_size && reduce_size >= block_size && - cld(dst_size, block_size) < TARGET_BLOCKS && + cld(dst_size, block_size) < target_blocks && length(reduce_sizes) == 1 && reduce_sizes[1] != 0 use_by_block = use_by_block_for_coalescing || use_by_block_for_low_occupancy @@ -263,7 +259,7 @@ function mapreduce_nd( # Grid-stride by_block: one full block cooperates on each output, then # grid-strides across remaining outputs if fewer blocks than outputs were # launched. - launch_blocks = min(dst_size, TARGET_BLOCKS) + launch_blocks = min(dst_size, target_blocks) _launch_mapreduce_nd_by_block!( backend, block_size, buffer, dst, f, op, init, neutral, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, @@ -275,7 +271,7 @@ function mapreduce_nd( reduce_groups = min( cld(reduce_size, block_size), block_size, - cld(TARGET_BLOCKS, dst_size), + cld(target_blocks, dst_size), cld(reduce_size, block_size * MIN_ITEMS_PER_THREAD), ) reduce_groups = max(reduce_groups, 1) @@ -312,9 +308,9 @@ function mapreduce_nd( end else # GS_DST_CUTOFF <= dst_size < reduce_size: grid-stride over outputs, one pass. - # Cap launched blocks at TARGET_BLOCKS; each block handles + # Cap launched blocks at target_blocks; each block handles # ceil(dst_size / launch_blocks) outputs sequentially. - launch_blocks = min(dst_size, TARGET_BLOCKS) + launch_blocks = min(dst_size, target_blocks) _launch_mapreduce_nd_by_block!( backend, block_size, buffer, dst, f, op, init, neutral, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, @@ -662,7 +658,7 @@ end # non-striding behavior) at zero extra cost — the while loop runs once. # # Used for GS_DST_CUTOFF <= dst_size < reduce_size (grid-stride, num_blocks == -# min(dst_size, TARGET_BLOCKS) < dst_size in general), and also for the +# min(dst_size, target_blocks) < dst_size in general), and also for the # reduce_groups==1 fallback inside the multigroup branch (num_blocks == dst_size, # so the loop runs exactly once per block — identical to the pre-grid-stride kernel). diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index c18c44e1..df898177 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -1,44 +1,9 @@ # Backend implementations include("utilities.jl") +include("tuning.jl") const MapReduceSource = Union{AbstractArray, Base.Broadcast.Broadcasted} -function _mapreduce_backend(src::AbstractArray) - return _mapreduce_get_backend(src) -end - -function _mapreduce_backend(src::Base.Broadcast.Broadcasted) - backend = _mapreduce_backend_from_args(src.args) - return isnothing(backend) ? CPU_BACKEND : backend -end - -function _mapreduce_get_backend(src::AbstractArray) - try - return get_backend(src) - catch err - err isa ArgumentError || rethrow() - return CPU_BACKEND - end -end - -_mapreduce_backend_from_arg(src::AbstractArray) = _mapreduce_get_backend(src) -_mapreduce_backend_from_arg(src::Base.Broadcast.Broadcasted) = _mapreduce_backend(src) -_mapreduce_backend_from_arg(_) = nothing - -function _mapreduce_backend_from_args(args::Tuple) - backend = nothing - for arg in args - arg_backend = _mapreduce_backend_from_arg(arg) - isnothing(arg_backend) && continue - if isnothing(backend) - backend = arg_backend - else - @argcheck arg_backend == backend - end - end - return backend -end - function _mapreduce_check_map_axes(src::AbstractArray, srcs::AbstractArray...) src_axes = axes(src) for other in srcs @@ -54,58 +19,38 @@ include("mapreduce_nd.jl") """ reduce( - op, src::AbstractArray, backend::Backend=get_backend(src); + op, src::AbstractArray; + backend=nothing, init, neutral=neutral_element(op, typeof(init)), dims=nothing, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - - # GPU settings - block_size::Int=256, - items_per_thread::Int=max(2, default_items_per_thread(backend)), + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, ) -Reduce `src` along dimensions `dims` using the binary operator `op`. If `dims` is `nothing` or -`:`, reduce `src` to a scalar. If `dims` is an integer or a collection of integers, reduce `src` -along those dimension(s). The `init` value is used as the initial value for the reduction; `neutral` -is the neutral element for the operator `op`. +Reduce `src` along dimensions `dims` using the binary operator `op`, which must be associative +and commutative. If `dims` is `nothing` or `:`, reduce `src` to a scalar, returned on the host. +If `dims` is an integer or a collection of integers, reduce `src` along those dimension(s). The +`init` value is used as the initial value for the reduction; `neutral` is the neutral element for +the operator `op`. The returned type is the same as `init` - to control output precision, specify `init` explicitly. -## CPU settings -Use at most `max_tasks` threads with at least `min_elems` elements per task. For N-dimensional -arrays (`dims` is an integer or a collection of integers) multithreading currently only becomes -faster for `max_tasks >= 4`; all other cases are scaling linearly with the number of threads. +`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref +AcceleratedKernels.CPUThreads.Partitioned) on the host and [`BlockReduce`](@ref) on GPUs, with the +device's settings. `backend` is derived from `src`. -Note that multithreading reductions only improves performance for cases with more compute-heavy -operations, which hide the memory latency and thread launch overhead - that includes: -- Reducing more complex types, e.g. reduction of tuples / structs / strings. -- More complex operators, e.g. `op=custom_complex_op_function`. - -For non-memory-bound operations, reductions scale almost linearly with the number of threads. - -## GPU settings -The `block_size` parameter controls the number of threads per block and must be a power of two. - -The GPU-only `items_per_thread` parameter must be at least 1 and controls how many elements each -thread reduces in a block-strided pass. Its default is -`max(2, default_items_per_thread(backend))`; larger values can improve throughput on discrete GPUs. +On the host, multithreading reductions only improves performance for operations that hide the +memory latency and thread launch overhead, e.g. reductions of tuples or structs, or expensive +operators. The `temp` parameter can be used to pass a pre-allocated temporary array. For reduction to a scalar (`dims=nothing` or `dims=:`), `length(temp) >= 2 * cld(length(src), items_per_thread * block_size)` -is required. For reduction along dimensions (`dims` is an integer or a collection of -integers), `temp` is used as the destination array, and thus must have the exact dimensions required - i.e. same -dimensionwise sizes as `src`, except for the reduced dimension(s) which become 1; there are some -corner cases when one dimension is zero, check against `Base.reduce` for CPU arrays for exact -behavior. - -The `switch_below` parameter controls the threshold below which the reduction is performed on the -CPU and is only used for 1D reductions (i.e. `dims=nothing`). +of the `BlockReduce` settings is required. For reduction along dimensions (`dims` is an integer or +a collection of integers), `temp` is used as the destination array, and thus must have the exact +dimensions required - i.e. same dimensionwise sizes as `src`, except for the reduced dimension(s) +which become 1; there are some corner cases when one dimension is zero, check against +`Base.reduce` for CPU arrays for exact behavior. # Examples Computing a sum, reducing down to a scalar that is copied to host: @@ -117,92 +62,42 @@ v = CuArray{Int16}(rand(1:1000, 100_000)) vsum = AK.reduce((x, y) -> x + y, v; init=zero(eltype(v))) ``` -Computing dimensionwise sums in a 2D matrix: +Computing dimensionwise sums in a 2D matrix, with explicit settings: ```julia import AcceleratedKernels as AK using Metal m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) mrowsum = AK.reduce(+, m; init=zero(eltype(m)), dims=1) -mcolsum = AK.reduce(+, m; init=zero(eltype(m)), dims=2) +mcolsum = AK.reduce(+, m; init=zero(eltype(m)), dims=2, alg=AK.BlockReduce(block_size=512)) ``` """ -function reduce( - op, src::AbstractArray, backend::Backend=_mapreduce_backend(src); - init, - kwargs... -) - _mapreduce_impl( - identity, op, src, backend; - init, - kwargs... - ) -end - - +reduce(op, src::AbstractArray; kwargs...) = mapreduce(identity, op, src; kwargs...) """ mapreduce( - f, op, src::AbstractArray, backend::Backend=get_backend(src); + f, op, src::AbstractArray, srcs::AbstractArray...; + backend=nothing, init, neutral=neutral_element(op, typeof(init)), dims=nothing, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - - # GPU settings - block_size::Int=256, - items_per_thread::Int=max(2, default_items_per_thread(backend)), + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, ) - mapreduce(f, op, A::AbstractArray, B::AbstractArray, As::AbstractArray...; init, kwargs...) - mapreduce(f, op, A::AbstractArray, B::AbstractArray, As::AbstractArray..., backend::Backend; init, kwargs...) - -Reduce `src` along dimensions `dims` using the binary operator `op` after applying `f` elementwise. -If `dims` is `nothing` or `:`, reduce `src` to a scalar. If `dims` is an integer or a collection of -integers, reduce `src` along those dimension(s). The `init` value is used as the initial value for -the reduction (i.e. after mapping). - -The `neutral` value is the neutral element for the operator `op`, which is needed for an efficient -GPU implementation that also allows a nonzero `init`. - -The returned type is the same as `init` - to control output precision, specify `init` explicitly. - -Multiple input arrays are supported with the same axes. This follows `Base.mapreduce(f, op, A, B, -...)` semantics: `f` is mapped across corresponding elements of the inputs and the mapped values -are reduced without materializing the intermediate array. Mismatched axes throw -`DimensionMismatch`; singleton-expanding broadcast semantics are reserved for internal -`Broadcasted` sources used by array backends. - -## CPU settings -Use at most `max_tasks` threads with at least `min_elems` elements per task. For N-dimensional -arrays (`dims` is an integer or a collection of integers) multithreading currently only becomes -faster for `max_tasks >= 4`; all other cases are scaling linearly with the number of threads. - -## GPU settings -The `block_size` parameter controls the number of threads per block and must be a power of two. - -The GPU-only `items_per_thread` parameter must be at least 1 and controls how many elements each -thread reduces in a block-strided pass. Its default is -`max(2, default_items_per_thread(backend))`; larger values can improve throughput on discrete GPUs. - -The `temp` parameter can be used to pass a pre-allocated temporary array. For reduction to a scalar -(`dims=nothing` or `dims=:`), `length(temp) >= 2 * cld(length(src), items_per_thread * block_size)` -is required. For reduction along dimensions (`dims` is an integer or a collection of -integers), `temp` is used as the destination array, and thus must have the exact dimensions required - i.e. same -dimensionwise sizes as `src`, except for the reduced dimension(s) which become 1; there are some -corner cases when one dimension is zero, check against `Base.reduce` for CPU arrays for exact -behavior. +Reduce `src` along dimensions `dims` using the binary operator `op` after applying `f` +elementwise. The keywords are those of [`reduce`](@ref); `init` is used as the initial value for +the reduction (i.e. after mapping), and `neutral` is needed for an efficient GPU implementation +that also allows a nonzero `init`. -The `switch_below` parameter controls the threshold below which the reduction is performed on the -CPU and is only used for 1D reductions (i.e. `dims=nothing`). +Multiple input arrays are supported with the same axes. This follows `Base.mapreduce(f, op, A, +B, ...)` semantics: `f` is mapped across corresponding elements of the inputs and the mapped +values are reduced without materializing the intermediate array. Mismatched axes throw +`DimensionMismatch`. `backend` is derived from all inputs. A `Broadcasted` object is also +accepted as the single source, with singleton-expanding broadcast semantics, for array backends. -# Example +# Examples Computing a sum of squares, reducing down to a scalar that is copied to host: ```julia import AcceleratedKernels as AK @@ -220,122 +115,61 @@ using Metal f(x) = x * x m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) mrowsumsq = AK.mapreduce(f, +, m; init=zero(eltype(m)), dims=1) -mcolsumsq = AK.mapreduce(f, +, m; init=zero(eltype(m)), dims=2) ``` Computing a two-input dimensional reduction: ```julia rows = AK.mapreduce((x, y) -> x * y, +, a, b; init=0f0, dims=1) ``` - -An explicit backend may be passed after all input arrays: -```julia -rows = AK.mapreduce((x, y) -> x * y, +, a, b, backend; init=0f0, dims=1) -``` """ function mapreduce( - f, op, src::MapReduceSource, backend::Backend=_mapreduce_backend(src); + f, op, src::MapReduceSource, srcs::AbstractArray...; + backend::Union{Nothing, Backend}=nothing, init, - kwargs... -) - _mapreduce_impl( - f, op, src, backend; - init, - kwargs... - ) -end - -function mapreduce( - f, op, src::AbstractArray, src2::AbstractArray, srcs::AbstractArray...; - init, - kwargs... -) - return _mapreduce_multi( - f, op, nothing, src, src2, srcs...; - init, - kwargs... - ) -end - -function mapreduce( - f, op, src::AbstractArray, src2::AbstractArray, arg, args...; - init, - kwargs... -) - backend = isempty(args) ? arg : args[end] - backend isa Backend || throw(MethodError(mapreduce, (f, op, src, src2, arg, args...))) - srcs = isempty(args) ? () : (arg, args[1:end - 1]...) - return _mapreduce_multi( - f, op, backend, src, src2, srcs...; - init, - kwargs... - ) -end - -function _mapreduce_multi( - f, op, backend::Union{Nothing, Backend}, src::AbstractArray, srcs::AbstractArray...; - init, - kwargs... + neutral=neutral_element(op, typeof(init)), + dims=nothing, + alg::Algorithm=Auto(), + temp::Union{Nothing, AbstractArray}=nothing, ) - _mapreduce_check_map_axes(src, srcs...) - bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) - return mapreduce( - identity, op, bc, isnothing(backend) ? _mapreduce_backend(bc) : backend; - init, - kwargs... - ) + if !isempty(srcs) + src isa AbstractArray || + throw(ArgumentError("a Broadcasted source cannot be combined with more arrays")) + _mapreduce_check_map_axes(src, srcs...) + src = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) + f = identity + end + backend = _resolve_backend(backend, src, temp) + a = _resolve_reduce(alg, backend, typeof(init), dims) + _mapreduce_impl(f, op, src, backend, a; init, neutral, dims, temp) end function _mapreduce_impl( - f, op, src::MapReduceSource, backend::Backend; - init, - neutral=neutral_element(op, typeof(init)), - dims = nothing, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int=256, - items_per_thread::Int=max(2, default_items_per_thread(backend)), - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, + f, op, src::MapReduceSource, backend::Backend, alg; + init, neutral, dims, temp, ) - # scalar *linear* indexing into a multidimensional Broadcasted object is # only available on Julia 1.12; on earlier version, materialize it first. if VERSION < v"1.12-" && src isa Base.Broadcast.Broadcasted src = Base.Broadcast.materialize(src) end - if isnothing(dims) || dims isa Colon - if use_gpu_algorithm(backend, prefer_threads) + if _whole(dims) + if alg isa BlockReduce mapreduce_1d_gpu( f, op, src, backend; init, neutral, - max_tasks, min_elems, - block_size, items_per_thread, temp, - switch_below + block_size=alg.block_size, items_per_thread=alg.items_per_thread, + temp, switch_below=alg.switch_below, ) else mapreduce_1d_cpu( f, op, src, backend; init, neutral, - max_tasks, min_elems, - block_size, temp, - switch_below + max_tasks=alg.max_tasks, min_elems=alg.min_elems, ) end else - return mapreduce_nd( - f, op, src, backend; - init, neutral, dims, - max_tasks, prefer_threads, - min_elems, block_size, - temp, - ) + mapreduce_nd(f, op, src, backend, alg; init, neutral, dims, temp) end end diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl new file mode 100644 index 00000000..94ce27e2 --- /dev/null +++ b/src/reduce/tuning.jl @@ -0,0 +1,116 @@ +# Selection and tuning of reduction algorithms; the steps are those of sorting (src/sort/tuning.jl). + + +""" + BlockReduce(; block_size=nothing, items_per_thread=nothing, switch_below=nothing) + +GPU tree reduction: each block of `block_size` threads (a power of two up to 1024) reduces +`block_size * items_per_thread` elements in local memory, repeatedly, until one value is left. +Reductions to a scalar (`dims=:`) finish on the host once fewer than `switch_below` values +remain; for them, `block_size * items_per_thread` must be between 2 and `typemax(Int32)`. Reductions along `dims` choose among several launch shapes by the sizes and strides of +the input and output; `items_per_thread` and `switch_below` do not apply to them, and setting +either is an `ArgumentError` there. + +Like every GPU reduction in the Julia ecosystem, it requires an associative and commutative +operator: the combination order is fixed for a given setting and shape, but it is not the +element order. +""" +Base.@kwdef struct BlockReduce <: ReduceAlgorithm + block_size::Union{Nothing, Int} = nothing + items_per_thread::Union{Nothing, Int} = nothing + switch_below::Union{Nothing, Int} = nothing +end + + +""" + ReduceTuning(; kwargs...) + +Values that fill unset algorithm fields for reductions on one device, as returned by +[`reduce_tuning`](@ref). + +- `block_size`, `items_per_thread`, `switch_below`: settings for unset `BlockReduce` fields. +- `target_blocks`: the number of blocks a reduction along `dims` aims to launch, to fill the + device when there are few outputs; must be positive. +- `threads_min_elems`: the default `min_elems` of `CPUThreads.Partitioned`. + +The defaults reproduce AK's historical settings. Internal: the fields may change in any release. +""" +Base.@kwdef struct ReduceTuning + block_size::Int = 256 + items_per_thread::Int = 2 + switch_below::Int = 0 + target_blocks::Int = 256 + threads_min_elems::Int = 1 +end + +""" + reduce_tuning(backend, T) -> ReduceTuning + +The reduction tuning for accumulator type `T` on `backend`'s current device; see +[`sort_tuning`](@ref) for the conventions. +""" +reduce_tuning(::Backend, ::Type) = ReduceTuning() + + +_whole(dims) = dims === nothing || dims isa Colon + +""" + _resolve_reduce(alg, backend, T, dims) -> Algorithm + +Resolve `alg` for a reduction with accumulator type `T` along `dims` (`nothing` or `:` for a +whole-array reduction) on `backend`. Returns `BlockReduce` or `CPUThreads.Partitioned` with every +field set, or throws an `ArgumentError`. +""" +function _resolve_reduce(alg::Algorithm, backend::Backend, ::Type{T}, dims) where {T} + _checkdomain(alg) + if alg isa BlockReduce && !_whole(dims) && + (alg.items_per_thread !== nothing || alg.switch_below !== nothing) + throw(ArgumentError("BlockReduce: `items_per_thread` and `switch_below` only apply to " * + "whole-array reductions (`dims=:`)")) + end + t = reduce_tuning(backend, T) + a = alg isa Auto ? _select_reduce(backend) : alg + a = _fill(a, t, T) + _check_reduce(a, backend, dims) + return a +end + +_select_reduce(backend) = _runs_threads(backend) ? CPUThreads.Partitioned() : BlockReduce() + +function _checkdomain(a::BlockReduce) + _check_pow2(a, :block_size) + a.block_size === nothing || a.block_size <= 1024 || throw(ArgumentError( + "BlockReduce: `block_size` must be at most 1024, got $(a.block_size)")) + _check_positive(a, :items_per_thread) + a.switch_below === nothing || a.switch_below >= 0 || throw(ArgumentError( + "BlockReduce: `switch_below` must be non-negative, got $(a.switch_below)")) + nothing +end + +_fill(a::BlockReduce, t::ReduceTuning, T) = + BlockReduce(something(a.block_size, t.block_size), + something(a.items_per_thread, t.items_per_thread), + something(a.switch_below, t.switch_below)) +_fill(a::CPUThreads.Partitioned, t::ReduceTuning, T) = _fill_threads(a, t) +_fill(a::Algorithm, t::ReduceTuning, T) = + throw(ArgumentError("$(_algname(a)) is not a reduction algorithm")) + +function _check_reduce(a::BlockReduce, backend, dims) + _checkdomain(a) + _require_kernels(a, backend) + if _whole(dims) + # Each pass reduces tiles of `block_size * items_per_thread` elements to one value: a tile + # of one element never shrinks the input, and the bound keeps the kernels' index + # arithmetic far from overflow + tile = widemul(a.block_size, a.items_per_thread) + 2 <= tile <= typemax(Int32) || throw(ArgumentError( + "BlockReduce: `block_size * items_per_thread` must be between 2 and $(typemax(Int32)), " * + "got $(a.block_size) * $(a.items_per_thread)")) + end + nothing +end + +function _check_reduce(a::CPUThreads.Partitioned, backend, dims) + _checkdomain(a) + _require_threads(a, backend) +end diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index 452c2b98..a4f55d97 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -32,14 +32,8 @@ end # Apply op(init, f(x)) to each element of src, storing the result in dst. -function _mapreduce_nd_apply_init!( - f, op, dst, src, backend; - init, - max_tasks=Threads.nthreads(), - min_elems=1, - block_size=256, -) - _foreachindex(eachindex(dst), backend; max_tasks, min_elems, block_size) do i +function _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) + _foreachindex(eachindex(dst), backend; launch...) do i dst[i] = op(init, f(src[i])) end end diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index db8c9eed..46bb4844 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -373,7 +373,8 @@ function _rs_key_range(v::AbstractArray{T}, backend::Backend, descending::Bool) min_k, max_k = mapreduce( x -> (k = _to_sort_key(x); (k, k)), (a, b) -> (min(a[1], b[1]), max(a[2], b[2])), - v, backend; + v; + backend, init=ident, neutral=ident, ) diff --git a/src/sort/tuning.jl b/src/sort/tuning.jl index aa5f58dd..4b53bfb4 100644 --- a/src/sort/tuning.jl +++ b/src/sort/tuning.jl @@ -50,11 +50,6 @@ per backend type, which may choose different values per device. sort_tuning(::Backend, ::Type) = SortTuning() -# Algorithm names in error messages -_algname(a) = nameof(typeof(a)) -_algname(::CPUThreads.SampleSort) = "CPUThreads.SampleSort" - - # The operation an algorithm is resolved for, for error messages _sort_opname(perm, pairs) = pairs ? "sort_by_key!" : perm ? "sortperm!" : "sort!" @@ -106,22 +101,6 @@ _rs_ordering(ord) = ord === Base.Order.Forward || ord === Base.Order.Reverse # Domain checks of explicitly set fields, before any arithmetic uses them -_checkdomain(::Algorithm) = nothing - -function _check_positive(a, field) - x = getfield(a, field) - x === nothing || x >= 1 || - throw(ArgumentError("$(_algname(a)): `$field` must be positive, got $x")) - nothing -end - -function _check_pow2(a, field) - x = getfield(a, field) - x === nothing || (x >= 1 && ispow2(x)) || - throw(ArgumentError("$(_algname(a)): `$field` must be a positive power of two, got $x")) - nothing -end - _checkdomain(a::MergeSort) = _check_positive(a, :block_size) function _checkdomain(a::RadixSort) @@ -146,10 +125,7 @@ function _checkdomain(a::BitonicSort) nothing end -function _checkdomain(a::CPUThreads.SampleSort) - _check_positive(a, :max_tasks) - _check_positive(a, :min_elems) -end +_checkdomain(a::CPUThreads.SampleSort) = _check_threads(a) # Fill unset fields from the tuning @@ -162,23 +138,13 @@ _fill(a::RadixSort, t::SortTuning, T) = _fill(a::BitonicSort, t::SortTuning, T) = BitonicSort(something(a.block_size, t.bitonic_block_size), something(a.items_per_thread, t.bitonic_items_per_thread)) -_fill(a::CPUThreads.SampleSort, t::SortTuning, T) = - CPUThreads.SampleSort(something(a.max_tasks, Threads.nthreads()), - something(a.min_elems, t.threads_min_elems)) +_fill(a::CPUThreads.SampleSort, t::SortTuning, T) = _fill_threads(a, t) _fill(a::Algorithm, t::SortTuning, T) = throw(ArgumentError("$(_algname(a)) is not a sorting algorithm")) # Check a complete algorithm against the backend, the operation and the arguments -function _require_kernels(a, backend) - _runs_kernels(backend) || throw(ArgumentError( - "$(_algname(a)) runs AcceleratedKernels' GPU kernels, which " * - "$(_backend_name(backend)) cannot run (on the host, this needs KernelAbstractions 0.10); " * - "use `Auto()` or a `CPUThreads` algorithm")) - nothing -end - function _check(a::MergeSort, backend, T, layout, ord; perm, pairs) _checkdomain(a) _require_kernels(a, backend) @@ -215,7 +181,5 @@ end function _check(a::CPUThreads.SampleSort, backend, T, layout, ord; perm, pairs) _checkdomain(a) - _runs_threads(backend) || throw(ArgumentError( - "CPUThreads.SampleSort only runs on the host backend, not on $(_backend_name(backend))")) - nothing + _require_threads(a, backend) end diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index c004cf03..dd537064 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1,3 +1,12 @@ +# Tests that do not choose an algorithm use `REDUCE_ALG`: `Auto()`, except in the `--cpu-ka` +# configuration, whose point is to run AK's kernels on the host backend. `reduce_alg` builds an +# explicitly tuned algorithm for the configuration from both kinds of settings. +REDUCE_ALG = HOST_KERNELS ? AK.BlockReduce() : AK.Auto() +reduce_alg(; block_size=nothing, items_per_thread=nothing, switch_below=nothing, + max_tasks=nothing, min_elems=nothing) = + TEST_KERNELS ? AK.BlockReduce(; block_size, items_per_thread, switch_below) : + AK.CPUThreads.Partitioned(; max_tasks, min_elems) + struct Point x::Float32 y::Float32 @@ -13,7 +22,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) AK.reduce( (x, y) -> x < y ? x : y, s; - prefer_threads, + alg=REDUCE_ALG, init=typemax(eltype(s)), neutral=typemax(eltype(s)), ) @@ -49,7 +58,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) AK.reduce( (x, y) -> x + y, s; - prefer_threads, + alg=REDUCE_ALG, init=zero(eltype(s)), neutral=zero(eltype(s)), ) @@ -95,7 +104,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Int32(1):Int32(100), num_elems)) - s = AK.reduce(+, v; prefer_threads, init=Int32(10)) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(10)) vh = Array(v) @test s == sum(vh) + 10 end @@ -106,7 +115,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) v = array_from_host(rand(1:100, num_elems), Int32) switch_below = rand(1:100) init = rand(1:100) - s = AK.reduce(+, v; prefer_threads, switch_below=switch_below, init=Int32(init)) + s = AK.reduce(+, v; alg=reduce_alg(switch_below=switch_below), init=Int32(init)) vh = Array(v) @test s == reduce(+, vh; init) end @@ -115,47 +124,41 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) for _ in 1:100 num_elems = rand(1:1000) v = 1:num_elems - s = AK.reduce(+, v, BACKEND; prefer_threads, init=Int32(0)) + s = AK.reduce(+, v; alg=REDUCE_ALG, backend=BACKEND, init=Int32(0)) vh = Array(v) @test s == reduce(+, vh) end # Base-compatible alias: dims=: reduces all dimensions to a scalar. vh_colon = rand(Int32(1):Int32(10), 3, 4, 5) - @test AK.reduce(+, array_from_host(vh_colon); prefer_threads, init=Int32(0), dims=:) == + @test AK.reduce(+, array_from_host(vh_colon); alg=REDUCE_ALG, init=Int32(0), dims=:) == reduce(+, vh_colon; init=Int32(0), dims=:) vh_one = Int32[7] - @test AK.reduce(+, array_from_host(vh_one); prefer_threads, init=Int32(10)) == + @test AK.reduce(+, array_from_host(vh_one); alg=REDUCE_ALG, init=Int32(10)) == reduce(+, vh_one; init=Int32(10)) # Test that undefined kwargs are not accepted @test_throws MethodError AK.reduce(+, array_from_host(rand(Int32, 10)); init=10, bad=:kwarg) - if !prefer_threads - @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 256)); prefer_threads, init=Int32(0), block_size=192) + if TEST_KERNELS + @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 256)); alg=reduce_alg(block_size=192), init=Int32(0)) end # Testing different settings AK.reduce( (x, y) -> x + 1, array_from_host(rand(Int32, 10_000)); - prefer_threads, + alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int64(0), - block_size=64, temp=array_from_host(zeros(Int32, 10_000)), - switch_below=50, - max_tasks=10, - min_elems=100, ) AK.reduce( (x, y) -> x + 1, array_from_host(rand(Int32, 10_000)); - prefer_threads, + alg=reduce_alg(max_tasks=16, min_elems=1000), init=Int32(0), neutral=Int64(0), - max_tasks=16, - min_elems=1000, ) end @@ -170,7 +173,7 @@ end for ksize in 0:3 sh = rand(Int32(1):Int32(100), isize, jsize, ksize) s = array_from_host(sh) - d = AK.reduce(+, s; prefer_threads, init=Int32(10), dims) + d = AK.reduce(+, s; alg=REDUCE_ALG, init=Int32(10), dims) dh = Array(d) @test dh == sum(sh; init=Int32(10), dims) @test eltype(dh) == eltype(sum(sh; init=Int32(10), dims)) @@ -187,7 +190,7 @@ end n3 = rand(1:100) vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.reduce(+, v; prefer_threads, init=Int32(0), dims) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(0), dims) sh = Array(s) @test sh == sum(vh; dims) end @@ -200,7 +203,7 @@ end n3 = rand(1:100) vh = rand(UInt32(1):UInt32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.reduce(+, v; prefer_threads, init=UInt32(0), dims) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=UInt32(0), dims) sh = Array(s) @test sh == sum(vh; dims) end @@ -213,7 +216,7 @@ end n3 = rand(1:100) vh = rand(Float32, n1, n2, n3) v = array_from_host(vh) - s = AK.reduce(+, v; prefer_threads, init=Float32(0), dims) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=Float32(0), dims) sh = Array(s) @test sh ≈ sum(vh; dims) end @@ -228,7 +231,7 @@ end vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) init = rand(1:100) - s = AK.reduce(+, v; prefer_threads, init=Int32(init), dims) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(init), dims) sh = Array(s) @test sh == reduce(+, vh; dims, init) end @@ -236,7 +239,7 @@ end # Duplicate dims match Base semantics and are reduced once. vh_dup = rand(Int32(1):Int32(10), 3, 4, 5) - @test Array(AK.reduce(+, array_from_host(vh_dup); prefer_threads, init=Int32(0), dims=(2,2))) == + @test Array(AK.reduce(+, array_from_host(vh_dup); alg=REDUCE_ALG, init=Int32(0), dims=(2,2))) == sum(vh_dup; init=Int32(0), dims=(2,2)) # min/max with dims: tests correct neutral element in partial reduction @@ -244,8 +247,8 @@ end n1 = rand(1:50); n2 = rand(1:50); n3 = rand(1:50) vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - @test Array(AK.reduce(min, v; prefer_threads, init=typemax(Int32), neutral=typemax(Int32), dims)) == minimum(vh; dims) - @test Array(AK.reduce(max, v; prefer_threads, init=typemin(Int32), neutral=typemin(Int32), dims)) == maximum(vh; dims) + @test Array(AK.reduce(min, v; alg=REDUCE_ALG, init=typemax(Int32), neutral=typemax(Int32), dims)) == minimum(vh; dims) + @test Array(AK.reduce(max, v; alg=REDUCE_ALG, init=typemin(Int32), neutral=typemin(Int32), dims)) == maximum(vh; dims) end # Tuple dims support. Order and duplicates match Base semantics. @@ -253,7 +256,7 @@ end for n1 in [1, 5, 10], n2 in [1, 5, 10], n3 in [1, 5, 10] vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.reduce(+, v; prefer_threads, init=Int32(0), dims) + s = AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(0), dims) sh = Array(s) @test sh == sum(vh; dims) end @@ -263,11 +266,11 @@ end for dims in ([1,2], [1,3], [2,3], [1,2,3], [2,1], [2,1,2], Int[], Any[1,2], Int32[1,2], 1:2) vh = rand(Int32(1):Int32(100), 3, 4, 5) v = array_from_host(vh) - @test Array(AK.reduce(+, v; prefer_threads, init=Int32(0), dims)) == + @test Array(AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(0), dims)) == sum(vh; init=Int32(0), dims) end - @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 3, 4)); prefer_threads, init=Int32(0), dims=[1.0, 2.0]) + @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 3, 4)); alg=REDUCE_ALG, init=Int32(0), dims=[1.0, 2.0]) # Tiled strided GPU path: contiguous kept dimensions, one strided reduce # dimension, and dst_size == reduce_size. The 3D case also exercises a @@ -275,72 +278,72 @@ end for (shape, dims) in (((512, 512), 2), ((20, 13, 260), 3)) vh = rand(Int32(1):Int32(3), shape...) v = array_from_host(vh) - @test Array(AK.reduce(+, v; prefer_threads, init=Int32(0), dims)) == + @test Array(AK.reduce(+, v; alg=REDUCE_ALG, init=Int32(0), dims)) == sum(vh; init=Int32(0), dims) end - if !prefer_threads + if TEST_KERNELS # Aligned stride-1 reductions use 128-bit loads for 4- and 8-byte elements. for T in valid_backend_eltypes(BACKEND, (Float32, Int32, UInt32, Float64, Int64, UInt64)) vh = T <: AbstractFloat ? rand(T, 1024, 1024) : rand(T(1):T(100), 1024, 1024) v = array_from_host(vh) @test AK._contiguous_vector_width(v, 0, (1024,), (1,), 1024, 256) == 16 ÷ sizeof(T) - r = Array(AK.reduce(+, v; prefer_threads, init=zero(T), dims=1)) + r = Array(AK.reduce(+, v; alg=REDUCE_ALG, init=zero(T), dims=1)) @test T <: AbstractFloat ? r ≈ sum(vh; dims=1) : r == sum(vh; dims=1) wh = T <: AbstractFloat ? rand(T, 256, 32, 16) : rand(T(1):T(100), 256, 32, 16) w = array_from_host(wh) - rw = Array(AK.reduce(+, w; prefer_threads, init=zero(T), dims=1)) + rw = Array(AK.reduce(+, w; alg=REDUCE_ALG, init=zero(T), dims=1)) @test T <: AbstractFloat ? rw ≈ sum(wh; dims=1) : rw == sum(wh; dims=1) end # Apply the map lane-wise and preserve non-additive neutral elements. fh = rand(Float32, 256, 1024) fv = array_from_host(fh) - @test Array(AK.mapreduce(abs2, +, fv; prefer_threads, init=0.0f0, dims=1)) ≈ + @test Array(AK.mapreduce(abs2, +, fv; alg=REDUCE_ALG, init=0.0f0, dims=1)) ≈ mapreduce(abs2, +, fh; dims=1) mh = rand(Int32(1):Int32(1000), 256, 1024) mv = array_from_host(mh) - @test Array(AK.reduce(min, mv; prefer_threads, init=typemax(Int32), neutral=typemax(Int32), dims=1)) == + @test Array(AK.reduce(min, mv; alg=REDUCE_ALG, init=typemax(Int32), neutral=typemax(Int32), dims=1)) == minimum(mh; dims=1) - @test Array(AK.reduce(max, mv; prefer_threads, init=typemin(Int32), neutral=typemin(Int32), dims=1)) == + @test Array(AK.reduce(max, mv; alg=REDUCE_ALG, init=typemin(Int32), neutral=typemin(Int32), dims=1)) == maximum(mh; dims=1) # Exercise scalar tails in both by-block dispatch branches. ph = rand(Float32, 1028, 2048) p = array_from_host(ph) for rows in (1026, 1025) - @test Array(AK.reduce(+, @view(p[1:rows, :]); prefer_threads, init=0.0f0, dims=1)) ≈ + @test Array(AK.reduce(+, @view(p[1:rows, :]); alg=REDUCE_ALG, init=0.0f0, dims=1)) ≈ sum(@view(ph[1:rows, :]); dims=1) end gh = rand(Float32, 1028, 64) g = array_from_host(gh) - @test Array(AK.reduce(+, @view(g[1:1025, :]); prefer_threads, init=0.0f0, dims=1)) ≈ + @test Array(AK.reduce(+, @view(g[1:1025, :]); alg=REDUCE_ALG, init=0.0f0, dims=1)) ≈ sum(@view(gh[1:1025, :]); dims=1) qh = rand(Int64(1):Int64(100), 1028, 2048) q = array_from_host(qh) - @test Array(AK.reduce(+, @view(q[1:1027, :]); prefer_threads, init=Int64(0), dims=1)) == + @test Array(AK.reduce(+, @view(q[1:1027, :]); alg=REDUCE_ALG, init=Int64(0), dims=1)) == sum(@view(qh[1:1027, :]); dims=1) # Misaligned rows and offsets fall back to scalar loads. rh = rand(Float32, 1027, 2048) r = array_from_host(rh) @test AK._contiguous_vector_width(r, 0, (1027,), (1,), 1024, 256) == 0 - @test Array(AK.reduce(+, @view(r[1:1024, :]); prefer_threads, init=0.0f0, dims=1)) ≈ + @test Array(AK.reduce(+, @view(r[1:1024, :]); alg=REDUCE_ALG, init=0.0f0, dims=1)) ≈ sum(@view(rh[1:1024, :]); dims=1) bh = rand(Int32(1):Int32(100), 260, 1024) b = array_from_host(bh) @test AK._contiguous_vector_width(b, 2, (260,), (1,), 258, 256) == 0 - @test Array(AK.reduce(+, @view(b[3:end, :]); prefer_threads, init=Int32(0), dims=1)) == + @test Array(AK.reduce(+, @view(b[3:end, :]); alg=REDUCE_ALG, init=Int32(0), dims=1)) == Base.reduce(+, @view(bh[3:end, :]); init=Int32(0), dims=1) end - if prefer_threads + if !TEST_KERNELS storage = Vector{UInt8}(undef, 64) GC.@preserve storage begin offset = Int(mod(-UInt(pointer(storage)), 16)) + 4 @@ -350,7 +353,7 @@ end # The CPU fallback should not require strided storage. vh = reshape(1:12, 1, 3, 4) - @test Array(AK.reduce(+, vh, BACKEND; prefer_threads, init=0, dims=(1,2))) == + @test Array(AK.reduce(+, vh; alg=REDUCE_ALG, backend=BACKEND, init=0, dims=(1,2))) == sum(vh; init=0, dims=(1,2)) else # Strided GPU sources (views, adjoints, permuted dims) take the stride-based @@ -358,48 +361,40 @@ end # base offset. Broadcasted/lazy sources still take the generic fallback. vh = reshape(Int32(1):Int32(40), 5, 8) v = array_from_host(vh) - @test Array(AK.reduce(+, @view(v[:, 1:2:end]); prefer_threads, init=Int32(0), dims=2)) == + @test Array(AK.reduce(+, @view(v[:, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == Base.reduce(+, @view(vh[:, 1:2:end]); init=Int32(0), dims=2) - @test Array(AK.reduce(+, @view(v[2:end, 1:2:end]); prefer_threads, init=Int32(0), dims=2)) == + @test Array(AK.reduce(+, @view(v[2:end, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == Base.reduce(+, @view(vh[2:end, 1:2:end]); init=Int32(0), dims=2) - @test Array(AK.reduce(+, v'; prefer_threads, init=Int32(0), dims=1)) == + @test Array(AK.reduce(+, v'; alg=REDUCE_ALG, init=Int32(0), dims=1)) == Base.reduce(+, vh'; init=Int32(0), dims=1) - @test Array(AK.reduce(+, PermutedDimsArray(v, (2, 1)); prefer_threads, init=Int32(0), dims=1)) == + @test Array(AK.reduce(+, PermutedDimsArray(v, (2, 1)); alg=REDUCE_ALG, init=Int32(0), dims=1)) == Base.reduce(+, PermutedDimsArray(vh, (2, 1)); init=Int32(0), dims=1) end # Test that undefined kwargs are not accepted - @test_throws MethodError AK.reduce(+, array_from_host(rand(Int32, 10, 10)); prefer_threads, init=10, bad=:kwarg) - if !prefer_threads - @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 16, 16)); prefer_threads, init=Int32(0), dims=1, block_size=192) + @test_throws MethodError AK.reduce(+, array_from_host(rand(Int32, 10, 10)); alg=REDUCE_ALG, init=10, bad=:kwarg) + if TEST_KERNELS + @test_throws ArgumentError AK.reduce(+, array_from_host(rand(Int32, 16, 16)); alg=reduce_alg(block_size=192), init=Int32(0), dims=1) end # Testing different settings AK.reduce( (x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, + alg=reduce_alg(block_size=64, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int32(0), dims=2, - block_size=64, temp=array_from_host(zeros(Int32, 3, 1, 5)), - switch_below=50, - max_tasks=10, - min_elems=100, ) AK.reduce( (x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, + alg=reduce_alg(block_size=64, max_tasks=16, min_elems=1000), init=Int32(0), neutral=Int32(0), dims=3, - block_size=64, temp=array_from_host(zeros(Int32, 3, 4, 1)), - switch_below=50, - max_tasks=16, - min_elems=1000, ) end @@ -407,13 +402,13 @@ end @testset "mapreduce_1d" begin Random.seed!(0) - function minbox(s; prefer_threads) + function minbox(s; alg=REDUCE_ALG) # Extract coordinates into tuple and reduce to find dimensionwise minima AK.mapreduce( p -> (p.x, p.y), (a, b) -> (a[1] < b[1] ? a[1] : b[1], a[2] < b[2] ? a[2] : b[2]), s; - prefer_threads, + alg, init=(typemax(Float32), typemax(Float32)), neutral=(typemax(Float32), typemax(Float32)), ) @@ -433,10 +428,10 @@ end for _ in 1:1000 num_elems = rand(1:100_000) v = array_from_host([Point(rand(Float32), rand(Float32)) for _ in 1:num_elems]) - mgpu = minbox(v; prefer_threads) + mgpu = minbox(v; alg=REDUCE_ALG) vh = Array(v) - mcpu = minbox(vh; prefer_threads=true) + mcpu = minbox(vh; alg=AK.Auto()) mbase = minbox_base(vh) @test typeof(mgpu) === typeof(mcpu) === typeof(mbase) @@ -451,10 +446,10 @@ end n3 = rand(1:100) v = array_from_host([Point(rand(Float32), rand(Float32)) for _ in 1:n1, _ in 1:n2, _ in 1:n3]) - mgpu = minbox(v; prefer_threads) + mgpu = minbox(v; alg=REDUCE_ALG) vh = Array(v) - mcpu = minbox(vh; prefer_threads=true) + mcpu = minbox(vh; alg=AK.Auto()) mbase = minbox_base(vh) @test typeof(mgpu) === typeof(mcpu) === typeof(mbase) @@ -466,7 +461,7 @@ end for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Int32(1):Int32(100), num_elems)) - s = AK.mapreduce(abs, +, v; prefer_threads, init=Int32(10)) + s = AK.mapreduce(abs, +, v; alg=REDUCE_ALG, init=Int32(10)) vh = Array(v) @test s == sum(vh) + 10 end @@ -477,7 +472,7 @@ end v = array_from_host(rand(-100:-1, num_elems), Int32) switch_below = rand(1:100) init = rand(1:100) - s = AK.mapreduce(abs, +, v; prefer_threads, switch_below=switch_below, init=Int32(init)) + s = AK.mapreduce(abs, +, v; alg=reduce_alg(switch_below=switch_below), init=Int32(init)) vh = Array(v) @test s == mapreduce(abs, +, vh; init) end @@ -486,48 +481,59 @@ end for _ in 1:100 num_elems = rand(1:1000) v = 1:num_elems - s = AK.mapreduce(abs, +, v, BACKEND; prefer_threads, init=Int32(0)) + s = AK.mapreduce(abs, +, v; alg=REDUCE_ALG, backend=BACKEND, init=Int32(0)) vh = Array(v) @test s == mapreduce(abs, +, vh) end # Base-compatible alias: dims=: reduces all dimensions to a scalar. vh_colon = rand(Int32(-10):Int32(10), 3, 4, 5) - @test AK.mapreduce(abs, +, array_from_host(vh_colon); prefer_threads, init=Int32(0), dims=:) == + @test AK.mapreduce(abs, +, array_from_host(vh_colon); alg=REDUCE_ALG, init=Int32(0), dims=:) == mapreduce(abs, +, vh_colon; init=Int32(0), dims=:) vh_one = Int32[-7] - @test AK.mapreduce(abs, +, array_from_host(vh_one); prefer_threads, init=Int32(10)) == + @test AK.mapreduce(abs, +, array_from_host(vh_one); alg=REDUCE_ALG, init=Int32(10)) == mapreduce(abs, +, vh_one; init=Int32(10)) - if !prefer_threads + if TEST_KERNELS for len in (65, 257, 1025), items_per_thread in (1, 2, 4) vh = Int32.(mod.(1:len, 17) .- 8) v = array_from_host(vh) for (f, op, neutral) in ((x -> 3x - 1, +, Int32(0)), (abs, max, typemin(Int32)), (x -> -x, min, typemax(Int32))) - @test AK.mapreduce(f, op, v; prefer_threads, init=neutral, neutral, - block_size=64, items_per_thread) == + @test AK.mapreduce( + f, + op, + v; + alg=reduce_alg(block_size=64, items_per_thread=items_per_thread), + init=neutral, + neutral, +) == Base.mapreduce(f, op, vh; init=neutral) end end - @test_throws ArgumentError AK.mapreduce(identity, +, array_from_host(Int32[1, 2]); - prefer_threads, init=Int32(0), items_per_thread=0) + @test_throws ArgumentError AK.mapreduce( + identity, + +, + array_from_host(Int32[1, 2]); + alg=reduce_alg(items_per_thread=0), + init=Int32(0), +) end vh_typechange = rand(Int32(-10):Int32(10), 4, 5) f_typechange = x -> Float32(x) / 2 - @test AK.mapreduce(f_typechange, +, array_from_host(vh_typechange); prefer_threads, init=0f0) ≈ + @test AK.mapreduce(f_typechange, +, array_from_host(vh_typechange); alg=REDUCE_ALG, init=0f0) ≈ mapreduce(f_typechange, +, vh_typechange; init=0f0) - @test Array(AK.mapreduce(f_typechange, +, array_from_host(vh_typechange); prefer_threads, init=0f0, dims=2)) ≈ + @test Array(AK.mapreduce(f_typechange, +, array_from_host(vh_typechange); alg=REDUCE_ALG, init=0f0, dims=2)) ≈ mapreduce(f_typechange, +, vh_typechange; init=0f0, dims=2) f_min_typechange = x -> Float32(10_000_000_000 + x) f_max_typechange = x -> Float32(-10_000_000_000 + x) - @test AK.mapreduce(f_min_typechange, min, array_from_host(vh_typechange); prefer_threads, init=Inf32) ≈ + @test AK.mapreduce(f_min_typechange, min, array_from_host(vh_typechange); alg=REDUCE_ALG, init=Inf32) ≈ mapreduce(f_min_typechange, min, vh_typechange; init=Inf32) - @test AK.mapreduce(f_max_typechange, max, array_from_host(vh_typechange); prefer_threads, init=-Inf32) ≈ + @test AK.mapreduce(f_max_typechange, max, array_from_host(vh_typechange); alg=REDUCE_ALG, init=-Inf32) ≈ mapreduce(f_max_typechange, max, vh_typechange; init=-Inf32) # Multi-input mapreduce lowers through a broadcasted source. @@ -537,44 +543,50 @@ end v_a = array_from_host(vh_a) v_b = array_from_host(vh_b) v_c = array_from_host(vh_c) - @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b; prefer_threads, init=Int32(0)) == + @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b; alg=REDUCE_ALG, init=Int32(0)) == mapreduce((x, y) -> x * y, +, vh_a, vh_b; init=Int32(0)) - @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b, BACKEND; prefer_threads, init=Int32(0)) == + @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b; alg=REDUCE_ALG, backend=BACKEND, init=Int32(0)) == mapreduce((x, y) -> x * y, +, vh_a, vh_b; init=Int32(0)) - @test AK.mapreduce((x, y, z) -> x + y * z, +, v_a, v_b, v_c, BACKEND; prefer_threads, init=Int32(0)) == + @test AK.mapreduce((x, y, z) -> x + y * z, +, v_a, v_b, v_c; alg=REDUCE_ALG, backend=BACKEND, init=Int32(0)) == mapreduce((x, y, z) -> x + y * z, +, vh_a, vh_b, vh_c; init=Int32(0)) - @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b; prefer_threads, init=Int32(0), dims=:) == + @test AK.mapreduce((x, y) -> x * y, +, v_a, v_b; alg=REDUCE_ALG, init=Int32(0), dims=:) == mapreduce((x, y) -> x * y, +, vh_a, vh_b; init=Int32(0), dims=:) - @test Array(AK.mapreduce((x, y) -> x * y, +, v_a, v_b; prefer_threads, init=Int32(0), dims=())) == + @test Array(AK.mapreduce((x, y) -> x * y, +, v_a, v_b; alg=REDUCE_ALG, init=Int32(0), dims=())) == mapreduce((x, y) -> x * y, +, vh_a, vh_b; init=Int32(0), dims=()) - @test AK.mapreduce((x, y) -> Float32(x - y) / 3, +, v_a, v_b; prefer_threads, init=0f0) ≈ + @test AK.mapreduce((x, y) -> Float32(x - y) / 3, +, v_a, v_b; alg=REDUCE_ALG, init=0f0) ≈ mapreduce((x, y) -> Float32(x - y) / 3, +, vh_a, vh_b; init=0f0) for (shape, dims) in (((0, 3), 1), ((2, 0), 2), ((0, 0), (1, 2)), ((0, 3), ())) h_empty1 = reshape(Int32[], shape...) h_empty2 = fill(Int32(2), shape...) - @test Array(AK.mapreduce((x, y) -> x + y, +, + @test Array(AK.mapreduce( + (x, y) -> x + y, + +, array_from_host(h_empty1), array_from_host(h_empty2); - prefer_threads, init=Int32(10), dims)) == + alg=REDUCE_ALG, + init=Int32(10), + dims, +)) == mapreduce((x, y) -> x + y, +, h_empty1, h_empty2; init=Int32(10), dims) end @test_throws DimensionMismatch AK.mapreduce( - (x, y) -> x + y, +, + (x, y) -> x + y, + +, array_from_host(rand(Int32, 2, 3)), array_from_host(rand(Int32, 1, 3)); - prefer_threads, + alg=REDUCE_ALG, init=Int32(0), ) - if prefer_threads + if !TEST_KERNELS bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(+, reshape(1:6, 2, 3), reshape(10:15, 2, 3))) - @test AK.mapreduce(identity, +, bc; prefer_threads, init=0) == + @test AK.mapreduce(identity, +, bc; alg=REDUCE_ALG, init=0) == mapreduce(identity, +, bc; init=0) - @test Array(AK.mapreduce(identity, +, bc; prefer_threads, init=0, dims=2)) == + @test Array(AK.mapreduce(identity, +, bc; alg=REDUCE_ALG, init=0, dims=2)) == mapreduce(identity, +, bc; init=0, dims=2) - @test Array(AK.mapreduce(identity, +, bc; prefer_threads, init=0, dims=())) == + @test Array(AK.mapreduce(identity, +, bc; alg=REDUCE_ALG, init=0, dims=())) == mapreduce(identity, +, bc; init=0, dims=()) end @@ -583,23 +595,19 @@ end p -> (p.x, p.y), (a, b) -> (a[1] < b[1] ? a[1] : b[1], a[2] < b[2] ? a[2] : b[2]), s; - prefer_threads, + alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), init=(typemax(Float32), typemax(Float32)), neutral=(typemax(Float32), typemax(Float32)), - block_size=64, temp=temp, - switch_below=50, - max_tasks=10, - min_elems=100, ) v = array_from_host([Point(rand(Float32), rand(Float32)) for _ in 1:10_042]) temp = similar(v, Tuple{Float32, Float32}) f(v, temp) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.mapreduce(-, +, v; prefer_threads, init=10, bad=:kwarg) - if !prefer_threads - @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 256)); prefer_threads, init=Int32(0), block_size=192) + @test_throws MethodError AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=10, bad=:kwarg) + if TEST_KERNELS + @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 256)); alg=reduce_alg(block_size=192), init=Int32(0)) end end @@ -614,7 +622,7 @@ end for ksize in 0:3 sh = rand(Int32(-100):Int32(100), isize, jsize, ksize) s = array_from_host(sh) - d = AK.mapreduce(-, +, s; prefer_threads, init=Int32(-10), dims) + d = AK.mapreduce(-, +, s; alg=REDUCE_ALG, init=Int32(-10), dims) dh = Array(d) @test dh == mapreduce(-, +, sh; init=Int32(-10), dims) @test eltype(dh) == eltype(mapreduce(-, +, sh; init=Int32(-10), dims)) @@ -631,19 +639,19 @@ end n3 = rand(1:100) vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.mapreduce(-, +, v; prefer_threads, init=Int32(0), dims) + s = AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=Int32(0), dims) sh = Array(s) @test sh == mapreduce(-, +, vh; init=Int32(0), dims) end end - function minbox(s, dims; prefer_threads) + function minbox(s, dims; alg=REDUCE_ALG) # Extract coordinates into tuple and reduce to find dimensionwise minima AK.mapreduce( p -> (p.x, p.y), (a, b) -> (a[1] < b[1] ? a[1] : b[1], a[2] < b[2] ? a[2] : b[2]), s; - prefer_threads, + alg, init=(typemax(Float32), typemax(Float32)), neutral=(typemax(Float32), typemax(Float32)), dims, @@ -668,10 +676,10 @@ end n2 = rand(1:100) n3 = rand(1:100) v = array_from_host([Point(rand(Float32), rand(Float32)) for _ in 1:n1, _ in 1:n2, _ in 1:n3]) - mgpu = minbox(v, dims; prefer_threads) + mgpu = minbox(v, dims; alg=REDUCE_ALG) vh = Array(v) - mcpu = minbox(vh, dims; prefer_threads=true) + mcpu = minbox(vh, dims; alg=AK.Auto()) mbase = minbox_base(vh, dims) @test eltype(mgpu) === eltype(mcpu) === eltype(mbase) @@ -691,7 +699,7 @@ end vh = rand(Int32(-100):Int32(100), n1, n2, n3) v = array_from_host(vh) init = rand(1:100) - s = AK.mapreduce(-, +, v; prefer_threads, init=Int32(init), dims) + s = AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=Int32(init), dims) sh = Array(s) @test sh == mapreduce(-, +, vh; dims, init) end @@ -699,7 +707,7 @@ end # Duplicate dims match Base semantics and are reduced once. vh_dup = rand(Int32(1):Int32(10), 3, 4, 5) - @test Array(AK.mapreduce(-, +, array_from_host(vh_dup); prefer_threads, init=Int32(0), dims=(2,2))) == + @test Array(AK.mapreduce(-, +, array_from_host(vh_dup); alg=REDUCE_ALG, init=Int32(0), dims=(2,2))) == mapreduce(-, +, vh_dup; init=Int32(0), dims=(2,2)) # Multi-input mapreduce with dimensional reductions. @@ -708,21 +716,21 @@ end v_ma = array_from_host(vh_ma) v_mb = array_from_host(vh_mb) for dims in (1, 2, (1, 2), (1, 3), (1, 2, 3), (2, 2)) - @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb; prefer_threads, init=Int32(0), dims)) == + @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb; alg=REDUCE_ALG, init=Int32(0), dims)) == mapreduce((x, y) -> x * y, +, vh_ma, vh_mb; init=Int32(0), dims) end - @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb, BACKEND; prefer_threads, init=Int32(0), dims=(1, 2))) == + @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb; alg=REDUCE_ALG, backend=BACKEND, init=Int32(0), dims=(1, 2))) == mapreduce((x, y) -> x * y, +, vh_ma, vh_mb; init=Int32(0), dims=(1, 2)) - @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb; prefer_threads, init=Int32(0), dims=())) == + @test Array(AK.mapreduce((x, y) -> x * y, +, v_ma, v_mb; alg=REDUCE_ALG, init=Int32(0), dims=())) == mapreduce((x, y) -> x * y, +, vh_ma, vh_mb; init=Int32(0), dims=()) - @test Array(AK.mapreduce((x, y) -> Float32(x - y) / 3, +, v_ma, v_mb; prefer_threads, init=0f0, dims=(1, 2))) ≈ + @test Array(AK.mapreduce((x, y) -> Float32(x - y) / 3, +, v_ma, v_mb; alg=REDUCE_ALG, init=0f0, dims=(1, 2))) ≈ mapreduce((x, y) -> Float32(x - y) / 3, +, vh_ma, vh_mb; init=0f0, dims=(1, 2)) vh_typechange_nd = rand(Int32(-10):Int32(10), 4, 5) f_min_typechange_nd = x -> Float32(10_000_000_000 + x) f_max_typechange_nd = x -> Float32(-10_000_000_000 + x) - @test Array(AK.mapreduce(f_min_typechange_nd, min, array_from_host(vh_typechange_nd); prefer_threads, init=Inf32, dims=2)) ≈ + @test Array(AK.mapreduce(f_min_typechange_nd, min, array_from_host(vh_typechange_nd); alg=REDUCE_ALG, init=Inf32, dims=2)) ≈ mapreduce(f_min_typechange_nd, min, vh_typechange_nd; init=Inf32, dims=2) - @test Array(AK.mapreduce(f_max_typechange_nd, max, array_from_host(vh_typechange_nd); prefer_threads, init=-Inf32, dims=2)) ≈ + @test Array(AK.mapreduce(f_max_typechange_nd, max, array_from_host(vh_typechange_nd); alg=REDUCE_ALG, init=-Inf32, dims=2)) ≈ mapreduce(f_max_typechange_nd, max, vh_typechange_nd; init=-Inf32, dims=2) # min/max with dims: tests correct neutral element in partial reduction @@ -730,8 +738,8 @@ end n1 = rand(1:50); n2 = rand(1:50); n3 = rand(1:50) vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - @test Array(AK.reduce(min, v; prefer_threads, init=typemax(Int32), neutral=typemax(Int32), dims)) == minimum(vh; dims) - @test Array(AK.reduce(max, v; prefer_threads, init=typemin(Int32), neutral=typemin(Int32), dims)) == maximum(vh; dims) + @test Array(AK.reduce(min, v; alg=REDUCE_ALG, init=typemax(Int32), neutral=typemax(Int32), dims)) == minimum(vh; dims) + @test Array(AK.reduce(max, v; alg=REDUCE_ALG, init=typemin(Int32), neutral=typemin(Int32), dims)) == maximum(vh; dims) end # Tuple dims support. Order and duplicates match Base semantics. @@ -739,7 +747,7 @@ end for n1 in [1, 5, 10], n2 in [1, 5, 10], n3 in [1, 5, 10] vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.mapreduce(-, +, v; prefer_threads, init=Int32(0), dims) + s = AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=Int32(0), dims) sh = Array(s) @test sh == mapreduce(-, +, vh; init=Int32(0), dims) end @@ -749,25 +757,25 @@ end for dims in ([1,2], [1,3], [2,3], [1,2,3], [2,1], [2,1,2], Int[], Any[1,2], Int32[1,2], 1:2) vh = rand(Int32(1):Int32(100), 3, 4, 5) v = array_from_host(vh) - @test Array(AK.mapreduce(-, +, v; prefer_threads, init=Int32(0), dims)) == + @test Array(AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=Int32(0), dims)) == mapreduce(-, +, vh; init=Int32(0), dims) end - @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 3, 4)); prefer_threads, init=Int32(0), dims=[1.0, 2.0]) + @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 3, 4)); alg=REDUCE_ALG, init=Int32(0), dims=[1.0, 2.0]) # Tiled strided GPU path coverage for mapreduce, including a 3D case with # a partial output tile. for (shape, dims) in (((512, 512), 2), ((20, 13, 260), 3)) vh = rand(Int32(1):Int32(3), shape...) v = array_from_host(vh) - @test Array(AK.mapreduce(x -> x - Int32(1), +, v; prefer_threads, init=Int32(0), dims)) == + @test Array(AK.mapreduce(x -> x - Int32(1), +, v; alg=REDUCE_ALG, init=Int32(0), dims)) == mapreduce(x -> x - Int32(1), +, vh; init=Int32(0), dims) end - if prefer_threads + if !TEST_KERNELS # The CPU fallback should not require strided storage. vh = reshape(1:12, 1, 3, 4) - @test Array(AK.mapreduce(x -> 2x, +, vh, BACKEND; prefer_threads, init=0, dims=(1,2))) == + @test Array(AK.mapreduce(x -> 2x, +, vh; alg=REDUCE_ALG, backend=BACKEND, init=0, dims=(1,2))) == mapreduce(x -> 2x, +, vh; init=0, dims=(1,2)) else # Strided GPU sources (views, adjoints, permuted dims) take the stride-based @@ -775,18 +783,18 @@ end # base offset. Broadcasted/lazy sources still take the generic fallback. vh = reshape(Int32(1):Int32(40), 5, 8) v = array_from_host(vh) - @test Array(AK.mapreduce(x -> x - Int32(1), +, @view(v[:, 1:2:end]); prefer_threads, init=Int32(0), dims=2)) == + @test Array(AK.mapreduce(x -> x - Int32(1), +, @view(v[:, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == mapreduce(x -> x - Int32(1), +, @view(vh[:, 1:2:end]); init=Int32(0), dims=2) - @test Array(AK.mapreduce(x -> x - Int32(1), +, @view(v[2:end, 1:2:end]); prefer_threads, init=Int32(0), dims=2)) == + @test Array(AK.mapreduce(x -> x - Int32(1), +, @view(v[2:end, 1:2:end]); alg=REDUCE_ALG, init=Int32(0), dims=2)) == mapreduce(x -> x - Int32(1), +, @view(vh[2:end, 1:2:end]); init=Int32(0), dims=2) - @test Array(AK.mapreduce(x -> x - Int32(1), +, PermutedDimsArray(v, (2, 1)); prefer_threads, init=Int32(0), dims=1)) == + @test Array(AK.mapreduce(x -> x - Int32(1), +, PermutedDimsArray(v, (2, 1)); alg=REDUCE_ALG, init=Int32(0), dims=1)) == mapreduce(x -> x - Int32(1), +, PermutedDimsArray(vh, (2, 1)); init=Int32(0), dims=1) end # Test that undefined kwargs are not accepted - @test_throws MethodError AK.mapreduce(-, +, array_from_host(rand(Int32, 3, 4, 5)); prefer_threads, init=10, bad=:kwarg) - if !prefer_threads - @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 16, 16)); prefer_threads, init=Int32(0), dims=1, block_size=192) + @test_throws MethodError AK.mapreduce(-, +, array_from_host(rand(Int32, 3, 4, 5)); alg=REDUCE_ALG, init=10, bad=:kwarg) + if TEST_KERNELS + @test_throws ArgumentError AK.mapreduce(-, +, array_from_host(rand(Int32, 16, 16)); alg=reduce_alg(block_size=192), init=Int32(0), dims=1) end # Testing different settings @@ -794,29 +802,21 @@ end -, (x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, + alg=reduce_alg(block_size=64, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int32(0), dims=2, - block_size=64, temp=array_from_host(zeros(Int32, 3, 1, 5)), - switch_below=50, - max_tasks=10, - min_elems=100, ) AK.mapreduce( -, (x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, + alg=reduce_alg(block_size=64, max_tasks=16, min_elems=1000), init=Int32(0), neutral=Int32(0), dims=3, - block_size=64, temp=array_from_host(zeros(Int32, 3, 4, 1)), - switch_below=50, - max_tasks=16, - min_elems=1000, ) end @testset "sum" begin @@ -825,13 +825,13 @@ end # Simple correctness tests v = array_from_host(1:100) - @test AK.sum(v; prefer_threads) == sum(Array(v)) + @test AK.sum(v; alg=REDUCE_ALG) == sum(Array(v)) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.sum(v; prefer_threads) ≈ sum(Array(v)) + @test AK.sum(v; alg=REDUCE_ALG) ≈ sum(Array(v)) end for _ in 1:100 @@ -843,10 +843,10 @@ end v = array_from_host(vh) # Indexing into array as if linear - @test AK.sum(v; prefer_threads) == sum(vh) + @test AK.sum(v; alg=REDUCE_ALG) == sum(vh) # Along dimensions - r = Array(AK.sum(v; prefer_threads, dims)) + r = Array(AK.sum(v; alg=REDUCE_ALG, dims)) rh = sum(vh; dims) @test r == rh @@ -855,11 +855,11 @@ end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.sum(v; prefer_threads, block_size=64) - @test AK.sum(v; prefer_threads, dims=:) == sum(Array(v); dims=:) + AK.sum(v; alg=reduce_alg(block_size=64)) + @test AK.sum(v; alg=REDUCE_ALG, dims=:) == sum(Array(v); dims=:) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.sum(v; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.sum(v; alg=REDUCE_ALG, bad=:kwarg) # The other settings are stress-tested in reduce end @@ -871,13 +871,13 @@ end # Simple correctness tests v = array_from_host(1:100) - @test AK.prod(v; prefer_threads) == prod(Array(v)) + @test AK.prod(v; alg=REDUCE_ALG) == prod(Array(v)) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.prod(v; prefer_threads) ≈ prod(Array(v)) + @test AK.prod(v; alg=REDUCE_ALG) ≈ prod(Array(v)) end for _ in 1:100 @@ -889,10 +889,10 @@ end v = array_from_host(vh) # Indexing into array as if linear - @test AK.sum(v; prefer_threads) == sum(vh) + @test AK.sum(v; alg=REDUCE_ALG) == sum(vh) # Along dimensions - r = Array(AK.sum(v; prefer_threads, dims)) + r = Array(AK.sum(v; alg=REDUCE_ALG, dims)) rh = sum(vh; dims) @test r == rh @@ -901,11 +901,11 @@ end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.prod(v; prefer_threads, block_size=64) - @test AK.prod(v; prefer_threads, dims=:) == prod(Array(v); dims=:) + AK.prod(v; alg=reduce_alg(block_size=64)) + @test AK.prod(v; alg=REDUCE_ALG, dims=:) == prod(Array(v); dims=:) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.prod(v; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.prod(v; alg=REDUCE_ALG, bad=:kwarg) # The other settings are stress-tested in reduce end @@ -917,13 +917,13 @@ end # Simple correctness tests v = array_from_host(1:100) - @test AK.minimum(v; prefer_threads) == minimum(Array(v)) + @test AK.minimum(v; alg=REDUCE_ALG) == minimum(Array(v)) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.minimum(v; prefer_threads) == minimum(Array(v)) + @test AK.minimum(v; alg=REDUCE_ALG) == minimum(Array(v)) end for _ in 1:100 @@ -935,10 +935,10 @@ end v = array_from_host(vh) # Indexing into array as if linear - @test AK.minimum(v; prefer_threads) == minimum(vh) + @test AK.minimum(v; alg=REDUCE_ALG) == minimum(vh) # Along dimensions - r = Array(AK.minimum(v; prefer_threads, dims)) + r = Array(AK.minimum(v; alg=REDUCE_ALG, dims)) rh = minimum(vh; dims) @test r == rh @@ -947,11 +947,11 @@ end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.minimum(v; prefer_threads, block_size=64) - @test AK.minimum(v; prefer_threads, dims=:) == minimum(Array(v); dims=:) + AK.minimum(v; alg=reduce_alg(block_size=64)) + @test AK.minimum(v; alg=REDUCE_ALG, dims=:) == minimum(Array(v); dims=:) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.minimum(v; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.minimum(v; alg=REDUCE_ALG, bad=:kwarg) # The other settings are stress-tested in reduce end @@ -963,13 +963,13 @@ end # Simple correctness tests v = array_from_host(1:100) - @test AK.maximum(v; prefer_threads) == maximum(Array(v)) + @test AK.maximum(v; alg=REDUCE_ALG) == maximum(Array(v)) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.maximum(v; prefer_threads) == maximum(Array(v)) + @test AK.maximum(v; alg=REDUCE_ALG) == maximum(Array(v)) end for _ in 1:100 @@ -981,10 +981,10 @@ end v = array_from_host(vh) # Indexing into array as if linear - @test AK.maximum(v; prefer_threads) == maximum(vh) + @test AK.maximum(v; alg=REDUCE_ALG) == maximum(vh) # Along dimensions - r = Array(AK.maximum(v; prefer_threads, dims)) + r = Array(AK.maximum(v; alg=REDUCE_ALG, dims)) rh = maximum(vh; dims) @test r == rh @@ -993,11 +993,11 @@ end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.maximum(v; prefer_threads, block_size=64) - @test AK.maximum(v; prefer_threads, dims=:) == maximum(Array(v); dims=:) + AK.maximum(v; alg=reduce_alg(block_size=64)) + @test AK.maximum(v; alg=REDUCE_ALG, dims=:) == maximum(Array(v); dims=:) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.maximum(v; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.maximum(v; alg=REDUCE_ALG, bad=:kwarg) # The other settings are stress-tested in reduce end @@ -1009,13 +1009,13 @@ end # Simple correctness tests v = array_from_host(1:100) - @test AK.count(x->x>50, v; prefer_threads) == count(x->x>50, Array(v)) + @test AK.count(x->x>50, v; alg=REDUCE_ALG) == count(x->x>50, Array(v)) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.count(x->x>0.5, v; prefer_threads) == count(x->x>0.5, Array(v)) + @test AK.count(x->x>0.5, v; alg=REDUCE_ALG) == count(x->x>0.5, Array(v)) end for _ in 1:100 @@ -1027,10 +1027,10 @@ end v = array_from_host(vh) # Indexing into array as if linear - @test AK.count(x->x>0.5, v; prefer_threads) == count(x->x>0.5, vh) + @test AK.count(x->x>0.5, v; alg=REDUCE_ALG) == count(x->x>0.5, vh) # Along dimensions - r = Array(AK.count(x->x>0.5, v; prefer_threads, dims)) + r = Array(AK.count(x->x>0.5, v; alg=REDUCE_ALG, dims)) rh = count(x->x>0.5, vh; dims) @test r == rh @@ -1041,16 +1041,16 @@ end for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Bool, num_elems)) - @test AK.count(v; prefer_threads) == count(Array(v)) + @test AK.count(v; alg=REDUCE_ALG) == count(Array(v)) end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.count(x->x>0, v; prefer_threads, block_size=64) - @test AK.count(x->x>0, v; prefer_threads, dims=:) == count(x->x>0, Array(v); dims=:) + AK.count(x->x>0, v; alg=reduce_alg(block_size=64)) + @test AK.count(x->x>0, v; alg=REDUCE_ALG, dims=:) == count(x->x>0, Array(v); dims=:) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.count(v; prefer_threads, bad=:kwarg) + @test_throws MethodError AK.count(v; alg=REDUCE_ALG, bad=:kwarg) # The other settings are stress-tested in reduce end @@ -1078,11 +1078,12 @@ end v = array_from_host(vh) shapes = (((3000, 40), 1), ((40, 3000), 2), ((20_000, 4), 1)) mats = [(array_from_host(rand(Int32(1):Int32(100), sz)), dims) for (sz, dims) in shapes] - for block_size in filter(<=(MAX_BLOCK_SIZE), 2 .^ (0:10)) - @test AK.reduce(+, v; prefer_threads, init=Int32(0), block_size) == sum(vh) + block_sizes = TEST_KERNELS ? filter(<=(MAX_BLOCK_SIZE), 2 .^ (0:10)) : [nothing] + for block_size in block_sizes + alg = reduce_alg(; block_size) + @test AK.reduce(+, v; init=Int32(0), alg) == sum(vh) for (m, dims) in mats - @test Array(AK.reduce(+, m; prefer_threads, init=Int32(0), dims, block_size)) == - sum(Array(m); dims) + @test Array(AK.reduce(+, m; init=Int32(0), dims, alg)) == sum(Array(m); dims) end end end diff --git a/test/generic/reduce_resolve.jl b/test/generic/reduce_resolve.jl new file mode 100644 index 00000000..474287c5 --- /dev/null +++ b/test/generic/reduce_resolve.jl @@ -0,0 +1,158 @@ +# Resolution of reduction algorithms, without launching kernels: resolution only looks at types +# and `dims`, so host arrays stand in for device arrays where an entry point is called. + +# A GPU backend with a tuning the tests can change +struct ReduceResolveTestBackend <: KernelAbstractions.Backend end +const REDUCE_RESOLVE_TUNING = Ref(AK.ReduceTuning()) +AK.reduce_tuning(::ReduceResolveTestBackend, ::Type) = REDUCE_RESOLVE_TUNING[] + +# A backend that cannot run AK's kernels, like KernelAbstractions 0.9's `CPU` +struct ReduceNoKernelsTestBackend <: KernelAbstractions.Backend end +AK._runs_kernels(::ReduceNoKernelsTestBackend) = false + +const RRB = ReduceResolveTestBackend() + +resolve_reduce(alg; T=Float32, dims=:, backend=RRB) = AK._resolve_reduce(alg, backend, T, dims) + +function with_reduce_tuning(f; kwargs...) + old = REDUCE_RESOLVE_TUNING[] + REDUCE_RESOLVE_TUNING[] = AK.ReduceTuning(; kwargs...) + try + f() + finally + REDUCE_RESOLVE_TUNING[] = old + end +end + + +@testset "reduce resolution: Auto" begin + # The default tuning reproduces AK's historical settings, for whole arrays and `dims` + @test resolve_reduce(AK.Auto()) === AK.BlockReduce(256, 2, 0) + @test resolve_reduce(AK.Auto(); dims=nothing) === AK.BlockReduce(256, 2, 0) + @test resolve_reduce(AK.Auto(); dims=2) === AK.BlockReduce(256, 2, 0) + @test resolve_reduce(AK.Auto(); dims=(1, 3)) === AK.BlockReduce(256, 2, 0) + + # On the host backend: the threaded algorithm, filled + host = AK.HOST_BACKEND + @test resolve_reduce(AK.Auto(); backend=host) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 1) + @test resolve_reduce(AK.Auto(); backend=host, dims=1) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 1) + + # The tuning fills every field; `stable` means nothing to a reduction + with_reduce_tuning(; block_size=512, items_per_thread=4, switch_below=100) do + @test resolve_reduce(AK.Auto()) === AK.BlockReduce(512, 4, 100) + @test resolve_reduce(AK.Auto(stable=false)) === AK.BlockReduce(512, 4, 100) + end +end + + +@testset "reduce resolution: explicit algorithms" begin + # Explicit fields win over the tuning, unset ones come from it + with_reduce_tuning(; block_size=512, items_per_thread=4, switch_below=100, + threads_min_elems=7) do + @test resolve_reduce(AK.BlockReduce()) === AK.BlockReduce(512, 4, 100) + @test resolve_reduce(AK.BlockReduce(block_size=64)) === AK.BlockReduce(64, 4, 100) + @test resolve_reduce(AK.BlockReduce(items_per_thread=1, switch_below=0)) === + AK.BlockReduce(512, 1, 0) + @test resolve_reduce(AK.BlockReduce(block_size=128); dims=1) === AK.BlockReduce(128, 4, 100) + end + host = AK.HOST_BACKEND + @test resolve_reduce(AK.CPUThreads.Partitioned(max_tasks=3); backend=host) === + AK.CPUThreads.Partitioned(3, 1) + @test resolve_reduce(AK.CPUThreads.Partitioned(min_elems=1000); backend=host) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 1000) + + # Domains of explicit fields + for bad in (AK.BlockReduce(block_size=0), AK.BlockReduce(block_size=192), + AK.BlockReduce(block_size=2048), AK.BlockReduce(items_per_thread=0), + AK.BlockReduce(switch_below=-1)) + @test_throws ArgumentError resolve_reduce(bad) + end + for bad in (AK.CPUThreads.Partitioned(max_tasks=0), AK.CPUThreads.Partitioned(min_elems=0)) + @test_throws ArgumentError resolve_reduce(bad; backend=host) + end + + # `items_per_thread` and `switch_below` only apply to whole-array reductions; they used to be + # ignored silently along `dims` + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(items_per_thread=4); dims=1) + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(switch_below=0); dims=(1, 2)) + @test resolve_reduce(AK.BlockReduce(items_per_thread=4); dims=nothing) isa AK.BlockReduce + + # Whole-array reductions need tiles of at least two elements (one never shrinks the input) + # and at most typemax(Int32), whether the settings are explicit or from the tuning + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(block_size=1, items_per_thread=1)) + @test resolve_reduce(AK.BlockReduce(block_size=1, items_per_thread=2)) === AK.BlockReduce(1, 2, 0) + @test resolve_reduce(AK.BlockReduce(block_size=1); dims=1) === AK.BlockReduce(1, 2, 0) + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(items_per_thread=1 << 56)) + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(block_size=1024, items_per_thread=1 << 22)) + with_reduce_tuning(; block_size=1, items_per_thread=1) do + @test_throws ArgumentError resolve_reduce(AK.Auto()) + @test resolve_reduce(AK.Auto(); dims=2) === AK.BlockReduce(1, 1, 0) + end + + # Capabilities: the threaded algorithm only on the host, kernels only where they run + @test_throws ArgumentError resolve_reduce(AK.CPUThreads.Partitioned()) + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); backend=ReduceNoKernelsTestBackend()) + if AK._runs_kernels(host) + @test resolve_reduce(AK.BlockReduce(); backend=host) === AK.BlockReduce(256, 2, 0) + else + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); backend=host) + end + + # Algorithms of other families + @test_throws ArgumentError resolve_reduce(AK.MergeSort()) + @test_throws ArgumentError resolve_reduce(AK.CPUThreads.SampleSort(); backend=host) + @test_throws ArgumentError AK._resolve_sort(AK.BlockReduce(), RRB, zeros(Float32, 10), :, + Base.Order.Forward; perm=false) +end + + +@testset "reduce resolution: tuning values" begin + # A tuning cannot make a reduction along `dims` launch nothing + with_reduce_tuning(; target_blocks=0) do + @test_throws ArgumentError AK.mapreduce_nd(identity, +, zeros(Int32, 4, 4), RRB, + AK.BlockReduce(256, 2, 0); + init=Int32(0), dims=1, temp=nothing) + end +end + + +@testset "reduce resolution: inference" begin + host = typeof(AK.HOST_BACKEND) + for (B, dims) in ((ReduceResolveTestBackend, Colon), (ReduceResolveTestBackend, Int), + (host, Nothing), (host, Tuple{Int, Int})) + rt = only(Base.return_types(AK._resolve_reduce, (AK.Auto, B, Type{Float32}, dims))) + @test rt <: Union{AK.BlockReduce, AK.CPUThreads.Partitioned} + end +end + + +@testset "reduce resolution: entry points" begin + v = rand(Int32(1):Int32(100), 1000) + m = rand(Int32(1):Int32(100), 10, 100) + + # Invalid explicit algorithms are rejected before any data is touched, whatever the length + for x in (v, Int32[], m) + @test_throws ArgumentError AK.reduce(+, x; init=Int32(0), alg=AK.BlockReduce(block_size=3)) + @test_throws ArgumentError AK.mapreduce(abs, +, x; init=Int32(0), alg=AK.MergeSort()) + @test_throws ArgumentError AK.sum(x; alg=AK.CPUThreads.Partitioned(max_tasks=0)) + end + @test_throws ArgumentError AK.reduce(+, m; init=Int32(0), dims=1, + alg=AK.BlockReduce(items_per_thread=4)) + + # The convenience reductions forward `alg` + alg = AK.CPUThreads.Partitioned(max_tasks=4, min_elems=10) + @test AK.sum(v; alg) == sum(v) + @test AK.prod(Int64.(v[1:5]); alg) == prod(Int64.(v[1:5])) + @test AK.maximum(v; alg) == maximum(v) + @test AK.minimum(v; alg) == minimum(v) + @test AK.count(>(50), v; alg) == count(>(50), v) + @test AK.sum(m; dims=2, alg) == sum(m; dims=2) + + # The loose settings are gone + @test_throws MethodError AK.reduce(+, v; init=Int32(0), block_size=256) + @test_throws MethodError AK.reduce(+, v; init=Int32(0), max_tasks=2) + @test_throws MethodError AK.reduce(+, v; init=Int32(0), prefer_threads=false) + @test_throws MethodError AK.sum(v; switch_below=10) +end diff --git a/test/generic/sort_resolve.jl b/test/generic/sort_resolve.jl index db6edf43..dc139f54 100644 --- a/test/generic/sort_resolve.jl +++ b/test/generic/sort_resolve.jl @@ -261,6 +261,11 @@ end 1.0, Ref(2), (3, 1:2), nothing) === RB @test AK._resolve_backend(nothing, 1:4, d) === RB @test AK._resolve_backend(nothing, h, 1:4) == host + # ... nor do Base's views, reshapes and permutations of them, while those of arrays vote + @test AK._resolve_backend(nothing, d, reshape(1:6, 2, 3), view(1:10, 2:4), + PermutedDimsArray(reshape(1:6, 2, 3), (2, 1))) === RB + @test AK._resolve_backend(nothing, view(d, 1:2)) === RB + @test AK._resolve_backend(nothing, reshape(1:6, 2, 3)) == host # All votes must agree, including the non-destination inputs @test_throws ArgumentError AK._resolve_backend(nothing, d, h) From b8082d588ade0d5c02f95fa71d63da48f64f210b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 07:19:15 +0200 Subject: [PATCH 4/8] Add `mapreducedim!` and define the reduction contract 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. --- Project.toml | 2 +- docs/make.jl | 1 + docs/src/api/differences.md | 19 ++ docs/src/api/mapreduce.md | 9 + docs/src/api/reduce.md | 20 ++- src/AcceleratedKernels.jl | 2 +- src/arithmetics.jl | 83 ++++++--- src/reduce/mapreduce_1d_cpu.jl | 21 +-- src/reduce/mapreduce_1d_gpu.jl | 44 +++-- src/reduce/mapreduce_nd.jl | 213 ++++++++++------------ src/reduce/reduce.jl | 314 ++++++++++++++++++++++++--------- src/reduce/tuning.jl | 10 +- src/reduce/utilities.jl | 183 +++++++++++++++++-- src/utils.jl | 4 + test/generic/reduce.jl | 263 +++++++++++++++++++++++++-- test/generic/reduce_resolve.jl | 54 +++++- 16 files changed, 942 insertions(+), 300 deletions(-) create mode 100644 docs/src/api/differences.md diff --git a/Project.toml b/Project.toml index 25db4364..637ccdcd 100644 --- a/Project.toml +++ b/Project.toml @@ -33,7 +33,7 @@ ArgCheck = "2" Atomix = "0.1, 1" CUDA = "6" CUDACore = "6.4.1" -GPUArraysCore = "0.2.0" +GPUArraysCore = "0.2.1" KernelAbstractions = "0.9.43, 0.10" Markdown = "1" Metal = "1.10.1" diff --git a/docs/make.jl b/docs/make.jl index 9c340a93..c332e76f 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -36,6 +36,7 @@ makedocs(; "Custom Structs" => "api/custom_structs.md", "Task Partitioning" => "api/task_partition.md", "Utilities" => "api/utilities.md", + "Differences from Base" => "api/differences.md", ], "Testing" => "testing.md", "Tuning and Capabilities" => "tuning.md", diff --git a/docs/src/api/differences.md b/docs/src/api/differences.md new file mode 100644 index 00000000..a6938666 --- /dev/null +++ b/docs/src/api/differences.md @@ -0,0 +1,19 @@ +### Differences from Base + +AcceleratedKernels' functions with Base's names take the arguments Base's do, but follow +contracts of their own, modelled on GPU libraries such as CUB: `init` is applied once, empty +inputs have no implicit result, and results have one documented type. Base's rules that exist for +Base's own reasons (empty results chosen per operator, `typeof(init)` as the result type along +`dims`, and so on) are left to front-ends: GPUArrays.jl, for example, implements Base's API for +GPU arrays on top of AcceleratedKernels and reproduces Base's results. The differences: + +| Call | Base | AcceleratedKernels | +|---|---|---| +| `reduce(op, A)`, `mapreduce(f, op, A)` of an empty `A`, no `init` | `Base.mapreduce_empty`: e.g. `sum(Int[]) == 0`, an error for `maximum` and for most mapped reductions | `ArgumentError` (`sum` and `prod` give zero and one) | +| `mapreduce(f, op, A; dims)` with an empty reduced dimension, no `init` | `Base.reducedim_init`: e.g. `[0 0]` for `x -> x + 1` with `+`, an error for `max` | `ArgumentError` where there are outputs (`sum`, `prod` and `count` give zero or one) | +| a one-element reduction, e.g. `reduce((a, b) -> a + b, [true])` | `true` (`Base.mapreduce_first`) | `1`, the accumulator type | +| the element type of `mapreduce(f, op, A; dims, init)` | `typeof(init)` | the accumulator type: `sum(Int8[1 2]; dims=1, init=Int16(0))` is a `Matrix{Int}` | +| the accumulator type of reductions into an array (`sum!`, reductions along `dims`) | depends on the code path, e.g. on the reduced dimension | one rule, the fold type from `eltype(R)` (see [`mapreducedim!`](@ref AcceleratedKernels.mapreducedim!)) | +| an `init` that is not a neutral element of `op` | outside the contract: `init` must be neutral, and it is unspecified whether it is used for non-empty collections | any value, applied exactly once | +| `reduce(op, A; dims)` for an `op` without `Base.reducedim_init`, such as a closure | `MethodError` | works | +| a non-commutative `op` in a reduction | works (elements keep their order) | unsupported | diff --git a/docs/src/api/mapreduce.md b/docs/src/api/mapreduce.md index 79b05eb4..8291709f 100644 --- a/docs/src/api/mapreduce.md +++ b/docs/src/api/mapreduce.md @@ -8,3 +8,12 @@ Equivalent to `reduce(op, map(f, iterable))`, without saving the intermediate ma ```@docs AcceleratedKernels.mapreduce ``` + +To reduce into an existing array, for example to accumulate into it across calls or to avoid an +allocation, use `mapreducedim!`. Its docstring states the contract every reduction follows: +the operator algebra, neutral elements, the accumulator type, how `init` and the destination's +values combine with the result, and empty slices. + +```@docs +AcceleratedKernels.mapreducedim! +``` diff --git a/docs/src/api/reduce.md b/docs/src/api/reduce.md index 59304d54..20f549b8 100644 --- a/docs/src/api/reduce.md +++ b/docs/src/api/reduce.md @@ -11,7 +11,19 @@ AcceleratedKernels.reduce The operator must be associative and commutative: like every GPU reduction, AcceleratedKernels combines elements in an order that depends on the algorithm and its settings, not in element -order. +order. `init` is optional and applied exactly once; without it, an empty reduction is an error, +except for `sum`, `prod` and `count`, and operators without a known neutral element need none. +Results have the accumulator type, which `acctype` can set (see +[`mapreducedim!`](@ref AcceleratedKernels.mapreducedim!) for the contract, and +[Differences from Base](@ref) for how it differs from Base's). + +```julia +AK.sum(CuArray(Int32[])) # 0, an Int +AK.minimum(CuArray(Int32[])) # an ArgumentError +AK.reduce((a, b) -> a + b, v) # no neutral element needed +AK.reduce(+, v; init=10) # init is added exactly once +AK.sum(CuArray(rand(Float32, 10^6)); acctype=Float64) # summed in Float64 +``` [`Auto()`](@ref AcceleratedKernels.Auto), the default `alg`, uses [`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host and @@ -26,7 +38,7 @@ AcceleratedKernels.CPUThreads.Partitioned ```julia m = ROCArray(rand(Float32, 1000, 1000)) -AK.reduce(+, m; init=0f0, alg=AK.BlockReduce(block_size=512)) # 512 threads per block -AK.reduce(+, m; init=0f0, alg=AK.BlockReduce(switch_below=1024)) # finish on the host below 1024 values -AK.reduce(+, Array(m); init=0f0, alg=AK.CPUThreads.Partitioned(max_tasks=4)) +AK.reduce(+, m; alg=AK.BlockReduce(block_size=512)) # 512 threads per block +AK.reduce(+, m; alg=AK.BlockReduce(switch_below=1024)) # finish on the host below 1024 values +AK.reduce(+, Array(m); alg=AK.CPUThreads.Partitioned(max_tasks=4)) ``` diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index fc091af8..e92b7918 100644 --- a/src/AcceleratedKernels.jl +++ b/src/AcceleratedKernels.jl @@ -12,7 +12,7 @@ module AcceleratedKernels # Internal dependencies using ArgCheck: @argcheck -using GPUArraysCore: AnyGPUArray, @allowscalar +using GPUArraysCore: AnyGPUArray, @allowscalar, neutral_element using KernelAbstractions using KernelAbstractions: @context using SIMD: Vec, vloada diff --git a/src/arithmetics.jl b/src/arithmetics.jl index d9211845..b2f6c485 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -1,29 +1,32 @@ """ - sum(src::AbstractArray; init=zero(eltype(src)), kwargs...) + sum(src::AbstractArray; kwargs...) -Sum of the elements of an array, with optional `init` and `dims`. The keywords are those of -[`reduce`](@ref). +Sum of the elements of an array, with Base's `add_sum`, so that small integers are summed as +`Int`. The keywords are those of [`mapreduce`](@ref). Without `init`, an empty array, or along +`dims` an empty slice, sums to zero of the accumulator type. ```julia import AcceleratedKernels as AK using Metal v = MtlArray(rand(Int32(1):Int32(100), 100_000)) -s = AK.sum(v) +s = AK.sum(v) # an Int +s = AK.sum(v; init=Int32(0)) # an Int as well +s = AK.sum(v; acctype=Int32) # an Int32, summed as Int32 m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) s = AK.sum(m; dims=1) # row-wise -s = AK.sum(m; dims=2, temp=MtlArray(zeros(Int32, 10, 1))) # into a preallocated result ``` """ -sum(src::AbstractArray; init=zero(eltype(src)), kwargs...) = reduce(+, src; init, kwargs...) +sum(src::AbstractArray; kwargs...) = _reduce_or_empty(zero, Base.add_sum, src; kwargs...) """ - prod(src::AbstractArray; init=one(eltype(src)), kwargs...) + prod(src::AbstractArray; kwargs...) -Product of the elements of an array, with optional `init` and `dims`. The keywords are those of -[`reduce`](@ref). +Product of the elements of an array, with Base's `mul_prod`. The keywords are those of +[`mapreduce`](@ref). Without `init`, an empty array, or along `dims` an empty slice, has the +product one of the accumulator type. ```julia import AcceleratedKernels as AK @@ -34,14 +37,39 @@ p = AK.prod(v) p = AK.prod(ROCArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -prod(src::AbstractArray; init=one(eltype(src)), kwargs...) = reduce(*, src; init, kwargs...) +prod(src::AbstractArray; kwargs...) = _reduce_or_empty(one, Base.mul_prod, src; kwargs...) + +# `reduce(op, src; kwargs...)`, except that without `init` an empty reduction gives `empty(A)` of +# the accumulator type `A` (for a whole array, and for every output along `dims`), instead of +# applying it as an `init`, which would change results such as the sign of a sum of `-0.0`s +function _reduce_or_empty(empty, op, src; init=_NoInit(), dims=:, acctype=nothing, + backend::Union{Nothing, Backend}=nothing, alg::Algorithm=Auto(), + kwargs...) + init isa _NoInit || return reduce(op, src; init, dims, acctype, backend, alg, kwargs...) + A = _acctype(op, Union{}, _mapped_eltype(identity, src), acctype) + if A !== Union{} && isempty(src) + # (the algorithm is checked as for any other input) + b = _resolve_backend(backend, src) + if _whole(dims) + _resolve_reduce(alg, b, A, dims) + return empty(A) + end + dims_valid = _reduced_dims(dims, ndims(src)) + if Base.any(d -> size(src, d) == 0, dims_valid) + _resolve_reduce(alg, b, A, dims_valid) + dst_sizes = ntuple(d -> d in dims_valid ? 1 : size(src, d), ndims(src)) + return fill!(KernelAbstractions.allocate(b, A, dst_sizes), empty(A)) + end + end + return reduce(op, src; dims, acctype, backend, alg, kwargs...) +end """ - maximum(src::AbstractArray; init=typemin(eltype(src)), kwargs...) + maximum(src::AbstractArray; kwargs...) -Maximum of the elements of an array, with optional `init` and `dims`. The keywords are those of -[`reduce`](@ref). +Maximum of the elements of an array; the maximum of an empty array, or along `dims` of an empty +slice, is an error unless `init` is given. The keywords are those of [`mapreduce`](@ref). ```julia import AcceleratedKernels as AK @@ -52,14 +80,14 @@ m = AK.maximum(v) m = AK.maximum(oneArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -maximum(src::AbstractArray; init=typemin(eltype(src)), kwargs...) = reduce(max, src; init, kwargs...) +maximum(src::AbstractArray; kwargs...) = reduce(max, src; kwargs...) """ - minimum(src::AbstractArray; init=typemax(eltype(src)), kwargs...) + minimum(src::AbstractArray; kwargs...) -Minimum of the elements of an array, with optional `init` and `dims`. The keywords are those of -[`reduce`](@ref). +Minimum of the elements of an array; the minimum of an empty array, or along `dims` of an empty +slice, is an error unless `init` is given. The keywords are those of [`mapreduce`](@ref). ```julia import AcceleratedKernels as AK @@ -70,14 +98,15 @@ m = AK.minimum(v) m = AK.minimum(CuArray(rand(Int32(1):Int32(100), 10, 100_000)); dims=1) ``` """ -minimum(src::AbstractArray; init=typemax(eltype(src)), kwargs...) = reduce(min, src; init, kwargs...) +minimum(src::AbstractArray; kwargs...) = reduce(min, src; kwargs...) """ count([f=identity,] src::AbstractArray; init=0, kwargs...) -Count the elements of `src` for which `f` returns `true`, with optional `init` and `dims`; the -result has the type of `init`. The keywords are those of [`mapreduce`](@ref). +Count the elements of `src` for which `f` returns `true`: `f` must return a `Bool`, and the +count is added to `init` (so an empty array counts `init`). The keywords are those of +[`mapreduce`](@ref). ```julia import AcceleratedKernels as AK @@ -88,15 +117,15 @@ c = AK.count(x -> x > 0.5, v) c = AK.count(CuArray(rand(Bool, 10, 100_000)); init=Int32(0), dims=2) ``` """ -function count(src::AbstractArray; init=0, kwargs...) - mapreduce(x -> x ? one(typeof(init)) : zero(typeof(init)), +, src; - init, neutral=zero(typeof(init)), kwargs...) -end +count(src::AbstractArray; kwargs...) = count(identity, src; kwargs...) +count(f, src::AbstractArray; init=0, kwargs...) = + mapreduce(_CountBool(f), Base.add_sum, src; init, kwargs...) -function count(f, src::AbstractArray; init=0, kwargs...) - mapreduce(x -> f(x) ? one(typeof(init)) : zero(typeof(init)), +, src; - init, neutral=zero(typeof(init)), kwargs...) +# `f`, required to return a `Bool` (Base's `_bool`) +struct _CountBool{F} <: Function + f::F end +@inline (c::_CountBool)(x) = c.f(x)::Bool """ diff --git a/src/reduce/mapreduce_1d_cpu.jl b/src/reduce/mapreduce_1d_cpu.jl index dc59be3e..3d31ed10 100644 --- a/src/reduce/mapreduce_1d_cpu.jl +++ b/src/reduce/mapreduce_1d_cpu.jl @@ -1,3 +1,5 @@ +# Reduce the non-empty `src` to a value on Julia threads; `neutral` and `init` as for +# `mapreduce_1d_gpu`. function mapreduce_1d_cpu( f, op, src::MapReduceSource, backend::Backend; init, @@ -5,26 +7,19 @@ function mapreduce_1d_cpu( max_tasks::Int, min_elems::Int, ) - if src isa Base.Broadcast.Broadcasted - return op(init, Base.mapreduce(f, op, src; init=neutral)) - end - - if max_tasks == 1 - return op(init, Base.mapreduce(f, op, src; init=neutral)) - end - + f, op_lanes = _lanefuncs(f, op, neutral) tp = TaskPartitioner(length(src), max_tasks, min_elems) - if tp.num_tasks == 1 - return op(init, Base.mapreduce(f, op, src; init=neutral)) + if src isa Base.Broadcast.Broadcasted || tp.num_tasks == 1 + return _finish(op, init, nothing, 0, Base.mapreduce(f, op_lanes, src; init=neutral)) end # Each task reduces an independent chunk of the array - shared = Vector{typeof(init)}(undef, tp.num_tasks) + shared = Vector{typeof(neutral)}(undef, tp.num_tasks) itask_partition(tp) do itask, irange @inbounds begin # This shared buffer is only modified once per task, so false sharing is not a problem - shared[itask] = Base.mapreduce(f, op, @view(src[irange]); init=neutral) + shared[itask] = Base.mapreduce(f, op_lanes, @view(src[irange]); init=neutral) end end - return op(init, Base.reduce(op, shared; init=neutral)) + return _finish(op, init, nothing, 0, Base.reduce(op_lanes, shared; init=neutral)) end diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index be90bb7f..09f5bc90 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -1,9 +1,12 @@ # NI and K are compile-time values so the local-memory size is static and the load loop unrolls. +# `neutral` seeds each thread's partial result (an empty `_Lane` when `op` has no known neutral +# element); `f === _Partials()` when `src` holds partial results of an earlier pass. @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_block!( @Const(src), dst, f, op, neutral, ::Val{NI}, ::Val{K}, ) where {NI, K} - sdata = @localmem eltype(dst) (NI,) + sdata = @localmem typeof(neutral) (NI,) + f, op = _lanefuncs(f, op, neutral) len = length(src) # NOTE: for many index calculations in this library, computation using zero-indexing leads to @@ -15,8 +18,7 @@ iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 - # Consecutive threads load consecutive elements, while each thread advances by NI. Starting - # from neutral preserves the previous `op(neutral, x) == x` bootstrap. + # Consecutive threads load consecutive elements, while each thread advances by NI. acc = neutral for s in 0x0:(K - 0x1) idx = iblock * (NI * K) + s * NI + ithread @@ -36,6 +38,8 @@ end +# Reduce the non-empty `src` to a host value; `neutral` is the partial-result seed of +# `_reduce_seed`, and `init` is a value or `_NoInit()`. function mapreduce_1d_gpu( f, op, src::MapReduceSource, backend::Backend; init, @@ -50,13 +54,17 @@ function mapreduce_1d_gpu( @argcheck items_per_thread >= 1 @argcheck switch_below >= 0 + P = typeof(neutral) + f_host, op_host = _lanefuncs(f, op, neutral) + # Degenerate cases len = length(src) - len == 0 && return init - len == 1 && return op(init, @allowscalar f(src[1])) + # `f` may index device arrays too + len == 1 && return _finish(op, init, nothing, 0, + @allowscalar(op_host(neutral, f_host(src[1])))) if len < switch_below - h_src = Vector(src) - return Base.mapreduce(f, op, h_src; init) + h_src = _host_copy(src) + return _finish(op, init, nothing, 0, Base.mapreduce(f_host, op_host, h_src; init=neutral)) end # Each block handles `items_per_thread * block_size` elements. @@ -64,14 +72,15 @@ function mapreduce_1d_gpu( blocks = (len + num_per_block - 1) ÷ num_per_block if !isnothing(temp) + neutral isa _Lane && throw(ArgumentError( + "`temp` needs a neutral element of the operator; pass `neutral`")) @argcheck get_backend(temp) === backend - @argcheck eltype(temp) === typeof(init) + eltype(temp) === P || throw(ArgumentError( + "`temp` must have the reduction's accumulator element type $P, got $(eltype(temp))")) @argcheck length(temp) >= blocks * 2 dst = temp else - # Figure out type for destination - dst_type = typeof(init) - dst = KernelAbstractions.allocate(backend, dst_type, blocks * 2) + dst = KernelAbstractions.allocate(backend, P, blocks * 2) end # Later the kernel will be compiled for views anyways, so use same types for arrays. @@ -87,7 +96,7 @@ function mapreduce_1d_gpu( len = blocks if len < switch_below h_src = Vector(@view(dst[1:len])) - return Base.reduce(op, h_src; init) + return _finish(op, init, nothing, 0, Base.reduce(op_host, h_src; init=neutral)) end # Now all src elements have been passed through f; just do final reduction, no map needed @@ -98,22 +107,25 @@ function mapreduce_1d_gpu( blocks = (len + num_per_block - 1) ÷ num_per_block # Each block produces one reduced value - kernel!(p1, p2, identity, op, neutral, Val(block_size), Val(items_per_thread); + kernel!(p1, p2, _Partials(), op, neutral, Val(block_size), Val(items_per_thread); ndrange=(block_size * blocks,)) len = blocks if len < switch_below h_src = Vector(@view(p2[1:len])) - return Base.reduce(op, h_src; init) + return _finish(op, init, nothing, 0, Base.reduce(op_host, h_src; init=neutral)) end p1, p2 = p2, p1 p1 = @view p1[1:len] end - # The GPU kernel reduced all elements to one, but without the init value - return op(init, @allowscalar(p1[1])) + # The GPU kernel reduced all elements to one; apply init + return _finish(op, init, nothing, 0, @allowscalar(p1[1])) end +_host_copy(src::AbstractArray) = Array(src) +_host_copy(src::Base.Broadcast.Broadcasted) = Array(Base.Broadcast.materialize(src)) + _mapreduce_1d_src_view(src::AbstractArray) = @view src[1:end] _mapreduce_1d_src_view(src::Base.Broadcast.Broadcasted) = src diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index cf06c35b..165b9da4 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -69,13 +69,30 @@ end # Main entry point -function mapreduce_nd( - f, op, src::MapReduceSource, backend::Backend, alg::Union{BlockReduce, CPUThreads.Partitioned}; +# The reduced dimensions `dims` (an integer or a collection of integers), as Base validates them: +# duplicates are ignored, and dimensions beyond `ndims` reduce nothing. +function _reduced_dims(dims, ndim::Int) + dims_src = dims isa Number ? (dims,) : dims + dims_buf = Int[] + for d in dims_src + d isa Integer || throw(ArgumentError("reduced dimension(s) must be integers")) + dim = Int(d) + dim < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $d")) + push!(dims_buf, dim) + end + return Tuple(d for d in Base.unique(dims_buf) if d <= ndim) +end + +# Reduce `src` into `dst`, which has `src`'s number of dimensions and size 1 along the reduced +# dimensions `dims_valid`; `A` is the accumulator type, and `init` is a value, `_NoInit()` or +# `_Fold()` (see `_finish`). +function mapreduce_nd!( + f, op, dst::AbstractArray, src::MapReduceSource, backend::Backend, + alg::Union{BlockReduce, CPUThreads.Partitioned}, ::Type{A}; init, - neutral=neutral_element(op, typeof(init)), - dims, - temp::Union{Nothing, AbstractArray}, -) + neutral, + dims_valid::Tuple, +) where {A} # Launch settings of the elementwise passes, and of the reduction kernels launch = alg isa BlockReduce ? (; block_size=alg.block_size) : (; max_tasks=alg.max_tasks, min_elems=alg.min_elems) @@ -83,76 +100,41 @@ function mapreduce_nd( # Number of blocks the by_block / multigroup paths aim to launch, so a reduction with # few output elements can still fill the GPU. A heuristic GPU-occupancy target. - target_blocks = reduce_tuning(backend, typeof(init)).target_blocks + target_blocks = reduce_tuning(backend, A).target_blocks target_blocks >= 1 || throw(ArgumentError( "the reduction tuning's `target_blocks` must be positive, got $target_blocks")) - dims_src = dims isa Number ? (dims,) : dims - dims_buf = Int[] - for d in dims_src - d isa Integer || throw(ArgumentError("reduced dimension(s) must be integers")) - dim = Int(d) - dim < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $d")) - push!(dims_buf, dim) - end - - # Match Base: duplicate dims are ignored, e.g. dims=(2,2) behaves like dims=2. - dims_all = Tuple(Base.unique(dims_buf)) - src_sizes = size(src) - ndim = length(src_sizes) - - dims_valid = Tuple(d for d in dims_all if d <= ndim) # Degenerate cases begin; order of priority matters - # All reduced dims are beyond ndims: just map each element through f and add init, e.g.: - # julia> x = rand(Float64, 3, 5); - # julia> mapreduce(x -> -x, +, x, dims=3, init=Float32(0)) # 3×5 Matrix{Float32} - if isempty(dims_valid) - dst = _alloc_or_temp(backend, temp, init, src_sizes) - _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) - return dst - end - - # The per-dimension sizes of the destination array; construct tuple without allocations - dst_sizes = unrolled_map_index(src_sizes) do i - i in dims_valid ? 1 : src_sizes[i] - end - - # If any kept dimension is zero, return empty array (reduced dims become 1), e.g.: - # julia> x = rand(3, 0, 5); reduce(+, x, dims=3) # 3×0×1 Array{Float64, 3} - for isize in eachindex(src_sizes) - isize in dims_valid && continue - if src_sizes[isize] == 0 - return _alloc_or_temp(backend, temp, init, dst_sizes) - end - end + len = Base.prod(src_sizes[d] for d in dims_valid; init=1) - len = Base.prod(src_sizes[d] for d in dims_valid) - - # If a reduced dimension is zero, return array filled with init, e.g.: - # julia> x = rand(3, 0, 5); mapreduce(+, x, dims=2) # 3×1×5 of zeros + # Empty reductions, e.g. reduce(+, rand(3, 0, 5), dims=2): `init`, else nothing is written if len == 0 - dst = _alloc_or_temp(backend, temp, init, dst_sizes) - fill!(dst, init) + init isa Union{_NoInit, _Fold} || fill!(dst, init) return dst end - # If the reduced extent is 1, just map each element through f (keep init's type) + # Nothing to write, e.g. reduce(+, rand(3, 0, 5), dims=3), a 3×0×1 array + isempty(dst) && return dst + # (overwriting with single elements combines nothing, even when `op` always throws) + len == 1 && init isa _NoInit || _check_acctype(op, f, A) + + # Every output reduces one element, e.g. reduce(+, rand(3, 5), dims=3): no reduced + # dimensions, or only ones of size 1 if len == 1 - dst = _alloc_or_temp(backend, temp, init, src_sizes) - _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) + _mapreduce_nd_single!(f, op, dst, src, backend, A; init, launch...) return dst end # Degenerate cases end - dst = _alloc_or_temp(backend, temp, init, dst_sizes) dst_size = length(dst) + neutral = _reduce_seed(op, A, neutral) if alg isa CPUThreads.Partitioned - _mapreduce_nd_cpu_sections!(f, op, dst, src; init, launch...) + _mapreduce_nd_cpu_sections!(f, op, dst, src; init, neutral, launch...) return dst end @@ -168,7 +150,7 @@ function mapreduce_nd( blocks = cld(dst_size, block_size) kernel! = _mapreduce_nd_generic!(backend, block_size) kernel!( - src, dst, f, op, init, + src, dst, f, op, neutral, init, CartesianIndices(dst), _mapreduce_reduce_indices(src, dst), dst_size, ndrange=(block_size * blocks,), ) @@ -241,7 +223,7 @@ function mapreduce_nd( blocks = cld(dst_size, rows_per_block) kernel! = _mapreduce_nd_by_thread_tiled_strided!(backend, block_size) kernel!( - buffer, dst, f, op, init, neutral, + buffer, dst, f, op, neutral, init, base_offset, reduce_strides[1], dst_size, reduce_size, Val(rows_per_block), ndrange=(block_size * blocks,), ) @@ -250,7 +232,7 @@ function mapreduce_nd( blocks = cld(dst_size, block_size) kernel! = _mapreduce_nd_by_thread!(backend, block_size) kernel!( - buffer, dst, f, op, init, + buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, dst_size, reduce_size, ndrange=(block_size * blocks,), @@ -261,7 +243,7 @@ function mapreduce_nd( # launched. launch_blocks = min(dst_size, target_blocks) _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, init, neutral, + backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, dst_size, reduce_size, launch_blocks, ) @@ -277,7 +259,7 @@ function mapreduce_nd( reduce_groups = max(reduce_groups, 1) if reduce_groups > 1 - partial = KernelAbstractions.allocate(backend, typeof(init), (dst_size, reduce_groups)) + partial = KernelAbstractions.allocate(backend, typeof(neutral), (dst_size, reduce_groups)) kernel! = _mapreduce_nd_multigroup!(backend, block_size) kernel!( @@ -293,7 +275,7 @@ function mapreduce_nd( pass2_block_size = _pass2_block_size(reduce_groups) kernel2! = _mapreduce_partial_to_dst!(backend, pass2_block_size) kernel2!( - partial, dst, op, init, neutral, + partial, dst, op, neutral, init, dst_size, reduce_groups, ndrange=(pass2_block_size * dst_size,), ) @@ -301,7 +283,7 @@ function mapreduce_nd( # reduce_groups collapsed to 1 (e.g. reduce_size <= block_size): just do a # single-block-per-output reduction directly, no partial array needed. _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, init, neutral, + backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, dst_size, reduce_size, dst_size, ) @@ -312,7 +294,7 @@ function mapreduce_nd( # ceil(dst_size / launch_blocks) outputs sequentially. launch_blocks = min(dst_size, target_blocks) _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, init, neutral, + backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, dst_size, reduce_size, launch_blocks, ) @@ -341,7 +323,8 @@ function _mapreduce_strided_layout(src::AbstractArray) s = try strides(src) catch err - err isa MethodError || rethrow() + # (an `ArgumentError` for a reshape of a non-contiguous view, on some Julia versions) + err isa Union{MethodError, ArgumentError} || rethrow() return nothing end @@ -364,7 +347,8 @@ function _mapreduce_is_dense_buffer(p::AbstractArray) try return strides(p) == _mapreduce_dense_strides(size(p)) catch err - err isa MethodError || rethrow() + # (an `ArgumentError` for a reshape of a non-contiguous view, on some Julia versions) + err isa Union{MethodError, ArgumentError} || rethrow() return false end end @@ -413,7 +397,7 @@ function _contiguous_vector_width( end function _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, init, neutral, + backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, num_blocks, ) @@ -423,7 +407,7 @@ function _launch_mapreduce_nd_by_block!( if W == 4 kernel! = _mapreduce_nd_by_block_contiguous!(backend, block_size) kernel!( - buffer, dst, f, op, init, neutral, + buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, output_size, reduce_size, num_blocks, Val(4), ndrange=(block_size * num_blocks,), @@ -431,7 +415,7 @@ function _launch_mapreduce_nd_by_block!( elseif W == 2 kernel! = _mapreduce_nd_by_block_contiguous!(backend, block_size) kernel!( - buffer, dst, f, op, init, neutral, + buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, output_size, reduce_size, num_blocks, Val(2), ndrange=(block_size * num_blocks,), @@ -439,7 +423,7 @@ function _launch_mapreduce_nd_by_block!( else kernel! = _mapreduce_nd_by_block!(backend, block_size) kernel!( - buffer, dst, f, op, init, neutral, + buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, num_blocks, ndrange=(block_size * num_blocks,), @@ -465,33 +449,25 @@ _mapreduce_reduce_indices(src, dst) = end -# Allocate a destination array, or validate and reuse a user-provided `temp`. -function _alloc_or_temp(backend, temp, init, sizes) - isnothing(temp) && return KernelAbstractions.allocate(backend, typeof(init), sizes) - @argcheck get_backend(temp) == backend - @argcheck size(temp) == sizes - @argcheck eltype(temp) == typeof(init) - temp -end - # CPU path function _mapreduce_nd_cpu_sections!( f, op, dst, src; - init, max_tasks, min_elems, + init, neutral, max_tasks, min_elems, ) Rother = CartesianIndices(dst) Rreduce = _mapreduce_reduce_indices(src, dst) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) _foreachindex(eachindex(dst), HOST_BACKEND; max_tasks, min_elems) do idst @inbounds begin Iother = Rother[idst] - res = init + res = neutral for Ireduce in Rreduce J = max(Iother, Ireduce) - res = op(res, f(src[J])) + res = op_lanes(res, f_lanes(src[J])) end - dst[idst] = res + dst[idst] = _finish(op, init, dst, idst, res) end end dst @@ -548,9 +524,10 @@ end # being elided, leaving a `throw` that the GPU backends cannot compile. @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_generic!( src, dst, - f, op, init, + f, op, neutral, init, Rother, Rreduce, output_size, ) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) N = @groupsize()[1] iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -558,25 +535,26 @@ end if tid < output_size Iother = Rother[tid + 0x1] - res = init + res = neutral for Ireduce in Rreduce J = max(Iother, Ireduce) - res = op(res, f(src[J])) + res = op_lanes(res, f_lanes(src[J])) end - dst[Iother] = res + dst[Iother] = _finish(op, init, dst, Iother, res) end end @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_by_thread_tiled_strided!( @Const(src), dst, - f, op, init, neutral, + f, op, neutral, init, base_offset, reduce_stride, output_size, reduce_size, ::Val{rows}, ) where {rows} @uniform N = @groupsize()[1] @uniform reduce_threads = N ÷ rows - sdata = @localmem eltype(dst) (N,) + sdata = @localmem typeof(neutral) (N,) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -594,7 +572,7 @@ end sbase = base_offset + iout j = lane while j < reduce_size - acc = op(acc, f(src[sbase + j * reduce_stride + 0x1])) + acc = op_lanes(acc, f_lanes(src[sbase + j * reduce_stride + 0x1])) j += reduce_threads end end @@ -607,14 +585,14 @@ end if lane < step dst_lane = row + lane * rows src_lane = row + (lane + step) * rows - sdata[dst_lane + 0x1] = op(sdata[dst_lane + 0x1], sdata[src_lane + 0x1]) + sdata[dst_lane + 0x1] = op_lanes(sdata[dst_lane + 0x1], sdata[src_lane + 0x1]) end @synchronize() step ÷= 2 end if lane == 0x0 && iout < output_size - dst[iout + 0x1] = op(init, sdata[row + 0x1]) + dst[iout + 0x1] = _finish(op, init, dst, iout + 0x1, sdata[row + 0x1]) end end @@ -624,12 +602,13 @@ end @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_by_thread!( @Const(src), dst, - f, op, init, + f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, ) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) # NOTE: index calculations use zero-indexing (fewer ops, matches the CUDA / ROCm / oneAPI / # Metal code this is transpiled to), converting to one-indexing only at memory accesses. # `base_offset` (the source's element offset within its dense buffer; 0 for a dense @@ -642,13 +621,13 @@ end if tid < output_size input_base = base_offset + _outer_decode(tid, outer_strides, outer_sizes) - res = init + res = neutral for j in 0x0:reduce_size - 0x1 off = _reduce_offset(j, reduce_strides, reduce_sizes) - res = op(res, f(src[input_base + off + 0x1])) + res = op_lanes(res, f_lanes(src[input_base + off + 0x1])) end - dst[tid + 0x1] = res + dst[tid + 0x1] = _finish(op, init, dst, tid + 0x1, res) end end @@ -664,14 +643,15 @@ end @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_by_block!( @Const(src), dst, - f, op, init, neutral, + f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, output_size, reduce_size, num_blocks, ) @uniform N = @groupsize()[1] - sdata = @localmem eltype(dst) (N,) + sdata = @localmem typeof(neutral) (N,) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -683,17 +663,17 @@ end j = ithread while j < reduce_size off = _reduce_offset(j, reduce_strides, reduce_sizes) - acc = op(acc, f(src[input_base + off + 0x1])) + acc = op_lanes(acc, f_lanes(src[input_base + off + 0x1])) j += N end sdata[ithread + 0x1] = acc @synchronize() - @inline reduce_group!(@context, op, sdata, ithread) + @inline reduce_group!(@context, op_lanes, sdata, ithread) if ithread == 0x0 - dst[iout + 0x1] = op(init, sdata[0x1]) + dst[iout + 0x1] = _finish(op, init, dst, iout + 0x1, sdata[0x1]) end next_iout = iout + num_blocks @@ -714,7 +694,7 @@ end # The host launcher verifies the pointer, row alignment, and UInt32 index bound. @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_nd_by_block_contiguous!( src, dst, - f, op, init, neutral, + f, op, neutral, init, base_offset, outer_strides, outer_sizes, output_size, reduce_size, num_blocks, @@ -723,7 +703,8 @@ end @uniform N = @groupsize()[1] @uniform T = eltype(src) @uniform W = _unval(valW) - sdata = @localmem eltype(dst) (N,) + sdata = @localmem typeof(neutral) (N,) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -737,21 +718,21 @@ end stridew = UInt32(N) * UInt32(W) while jw + UInt32(W - 1) < rs v = vloada(Vec{W, T}, pointer(src, input_base + jw + 0x1)) - acc = _acc_lanes(f, op, acc, v) + acc = _acc_lanes(f_lanes, op_lanes, acc, v) jw += stridew end while jw < rs - acc = op(acc, f(src[input_base + jw + 0x1])) + acc = op_lanes(acc, f_lanes(src[input_base + jw + 0x1])) jw += UInt32(1) end sdata[ithread + 0x1] = acc @synchronize() - @inline reduce_group!(@context, op, sdata, ithread) + @inline reduce_group!(@context, op_lanes, sdata, ithread) if ithread == 0x0 - dst[iout + 0x1] = op(init, sdata[0x1]) + dst[iout + 0x1] = _finish(op, init, dst, iout + 0x1, sdata[0x1]) end next_iout = iout + num_blocks @@ -775,7 +756,8 @@ end ::Val{output_size}, reduce_size, reduce_groups, ) where {output_size} @uniform N = @groupsize()[1] - sdata = @localmem eltype(partial) (N,) + sdata = @localmem typeof(neutral) (N,) + f_lanes, op_lanes = _lanefuncs(f, op, neutral) iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -790,14 +772,14 @@ end j = ithread + igroup * N while j < reduce_size off = _reduce_offset(j, reduce_strides, reduce_sizes) - acc = op(acc, f(src[input_base + off + 0x1])) + acc = op_lanes(acc, f_lanes(src[input_base + off + 0x1])) j += N * reduce_groups end sdata[ithread + 0x1] = acc @synchronize() - @inline reduce_group!(@context, op, sdata, ithread) + @inline reduce_group!(@context, op_lanes, sdata, ithread) if ithread == 0x0 partial[iout + igroup * output_size + 0x1] = sdata[0x1] @@ -805,17 +787,18 @@ end end # GPU kernel: multi-group second pass — one block per output reduces over the `reduce_groups` -# partials and folds in `init`. Launched with a block size sized to `reduce_groups` +# partials and applies `init` (see `_finish`). Launched with a block size sized to `reduce_groups` # (see _pass2_block_size), avoiding an oversubscribed tree-reduce when reduce_groups # is much smaller than 256. @kernel inbounds=true cpu=false unsafe_indices=true function _mapreduce_partial_to_dst!( @Const(partial), dst, - op, init, neutral, + op, neutral, init, output_size, reduce_groups, ) @uniform N = @groupsize()[1] - sdata = @localmem eltype(dst) (N,) + sdata = @localmem typeof(neutral) (N,) + _, op_lanes = _lanefuncs(_Partials(), op, neutral) iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 @@ -824,17 +807,17 @@ end acc = neutral g = ithread while g < reduce_groups - acc = op(acc, partial[iblock + g * output_size + 0x1]) + acc = op_lanes(acc, partial[iblock + g * output_size + 0x1]) g += N end sdata[ithread + 0x1] = acc @synchronize() - @inline reduce_group!(@context, op, sdata, ithread) + @inline reduce_group!(@context, op_lanes, sdata, ithread) if ithread == 0x0 - dst[iblock + 0x1] = op(init, sdata[0x1]) + dst[iblock + 0x1] = _finish(op, init, dst, iblock + 0x1, sdata[0x1]) end end end diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index df898177..cc3908dd 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -4,6 +4,9 @@ include("tuning.jl") const MapReduceSource = Union{AbstractArray, Base.Broadcast.Broadcasted} +# The size of a reduction source along `d`; `Broadcasted` objects only have `size` without `d` +_srcsize(src, d::Int) = d <= ndims(src) ? size(src)[d] : 1 + function _mapreduce_check_map_axes(src::AbstractArray, srcs::AbstractArray...) src_axes = axes(src) for other in srcs @@ -18,58 +21,20 @@ include("mapreduce_nd.jl") """ - reduce( - op, src::AbstractArray; - backend=nothing, - init, - neutral=neutral_element(op, typeof(init)), - dims=nothing, - alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, - ) - -Reduce `src` along dimensions `dims` using the binary operator `op`, which must be associative -and commutative. If `dims` is `nothing` or `:`, reduce `src` to a scalar, returned on the host. -If `dims` is an integer or a collection of integers, reduce `src` along those dimension(s). The -`init` value is used as the initial value for the reduction; `neutral` is the neutral element for -the operator `op`. - -The returned type is the same as `init` - to control output precision, specify `init` explicitly. - -`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref -AcceleratedKernels.CPUThreads.Partitioned) on the host and [`BlockReduce`](@ref) on GPUs, with the -device's settings. `backend` is derived from `src`. + reduce(op, src::AbstractArray; kwargs...) -On the host, multithreading reductions only improves performance for operations that hide the -memory latency and thread launch overhead, e.g. reductions of tuples or structs, or expensive -operators. - -The `temp` parameter can be used to pass a pre-allocated temporary array. For reduction to a scalar -(`dims=nothing` or `dims=:`), `length(temp) >= 2 * cld(length(src), items_per_thread * block_size)` -of the `BlockReduce` settings is required. For reduction along dimensions (`dims` is an integer or -a collection of integers), `temp` is used as the destination array, and thus must have the exact -dimensions required - i.e. same dimensionwise sizes as `src`, except for the reduced dimension(s) -which become 1; there are some corner cases when one dimension is zero, check against -`Base.reduce` for CPU arrays for exact behavior. +Reduce `src` with the binary operator `op`, which must be associative and commutative; the +keywords are those of [`mapreduce`](@ref). Equivalent to `mapreduce(identity, op, src; kwargs...)`. # Examples -Computing a sum, reducing down to a scalar that is copied to host: -```julia -import AcceleratedKernels as AK -using CUDA - -v = CuArray{Int16}(rand(1:1000, 100_000)) -vsum = AK.reduce((x, y) -> x + y, v; init=zero(eltype(v))) -``` - -Computing dimensionwise sums in a 2D matrix, with explicit settings: ```julia import AcceleratedKernels as AK using Metal m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) -mrowsum = AK.reduce(+, m; init=zero(eltype(m)), dims=1) -mcolsum = AK.reduce(+, m; init=zero(eltype(m)), dims=2, alg=AK.BlockReduce(block_size=512)) +AK.reduce(+, m) # a host scalar +AK.reduce(max, m; dims=1) # a 1×100_000 MtlArray +AK.reduce(+, m; dims=2, alg=AK.BlockReduce(block_size=512)) ``` """ reduce(op, src::AbstractArray; kwargs...) = mapreduce(identity, op, src; kwargs...) @@ -77,25 +42,52 @@ reduce(op, src::AbstractArray; kwargs...) = mapreduce(identity, op, src; kwargs. """ mapreduce( - f, op, src::AbstractArray, srcs::AbstractArray...; + f, op, src, srcs::AbstractArray...; backend=nothing, - init, - neutral=neutral_element(op, typeof(init)), - dims=nothing, + init=, + neutral=nothing, + acctype=nothing, + dims=:, alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, ) -Reduce `src` along dimensions `dims` using the binary operator `op` after applying `f` -elementwise. The keywords are those of [`reduce`](@ref); `init` is used as the initial value for -the reduction (i.e. after mapping), and `neutral` is needed for an efficient GPU implementation -that also allows a nonzero `init`. +Apply `f` to each element of `src` and reduce the results with the binary operator `op`. With +`dims=:` (the default; `nothing` is accepted too) the result is a scalar, returned on the host. +With `dims` an integer or a collection of integers, the result is a new array on `src`'s backend +with size 1 along those dimensions; [`mapreducedim!`](@ref) reduces into an existing array +instead, and states the contract all reductions follow. In short: + +- `op` must be associative and commutative. +- `init`, when given, is applied exactly once, as `op(init, partial)`, and is the result of an + empty reduction. Without `init`, an empty reduction is an `ArgumentError`: of a whole array, and + along `dims` of any output whose slice is empty. ([`sum`](@ref), [`prod`](@ref) and + [`count`](@ref) give zero or one instead.) +- The result has the accumulator type: `acctype` when it is given, else the type the fold of `op` + settles on from `init`'s type (when given) and the mapped elements; `op(init, partial)` is + converted to it. So `AK.sum(Int8[1, 2])` is an `Int`, and `AK.sum(Int8[1 2]; dims=1, + init=Int16(0))` an array of `Int`s. A single element has that type too: + `AK.reduce((a, b) -> a + b, [true]) === 1`. (An empty reduction returns `init` as it is.) +- `neutral`, a two-sided identity of `op`, only seeds partial results and never appears in the + result. It defaults to `GPUArraysCore.neutral_element(op, T)` where that is defined; for other + operators, each partial result starts from its first element instead. + +These rules are AcceleratedKernels' own, and differ from Base's in a few places, e.g. for empty +reductions without `init`; see [Differences from Base](@ref). + +Several source arrays with equal axes are reduced without materializing the mapped array +(`f` takes one argument per array); mismatched axes throw a `DimensionMismatch`. A `Broadcasted` +object is also accepted as the single source. Before Julia 1.12, both are materialized when the +reduction runs. `backend` is derived from all sources. + +`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref +AcceleratedKernels.CPUThreads.Partitioned) on the host and [`BlockReduce`](@ref) on GPUs, with the +device's settings. On the host, multithreading reductions only improves performance for +operations that hide the memory latency and thread launch overhead, e.g. reductions of tuples or +structs, or expensive operators. -Multiple input arrays are supported with the same axes. This follows `Base.mapreduce(f, op, A, -B, ...)` semantics: `f` is mapped across corresponding elements of the inputs and the mapped -values are reduced without materializing the intermediate array. Mismatched axes throw -`DimensionMismatch`. `backend` is derived from all inputs. A `Broadcasted` object is also -accepted as the single source, with singleton-expanding broadcast semantics, for array backends. +`temp` is scratch space for whole-array `BlockReduce` reductions: a vector with the accumulator +element type and at least `2 * cld(length(src), block_size * items_per_thread)` elements. # Examples Computing a sum of squares, reducing down to a scalar that is copied to host: @@ -104,7 +96,7 @@ import AcceleratedKernels as AK using CUDA v = CuArray{Int16}(rand(1:1000, 100_000)) -vsumsq = AK.mapreduce(x -> x * x, (x, y) -> x + y, v; init=zero(eltype(v))) +vsumsq = AK.mapreduce(x -> Int(x) * x, +, v) ``` Computing dimensionwise sums of squares in a 2D matrix: @@ -114,20 +106,21 @@ using Metal f(x) = x * x m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000)) -mrowsumsq = AK.mapreduce(f, +, m; init=zero(eltype(m)), dims=1) +mrowsumsq = AK.mapreduce(f, +, m; dims=1) ``` Computing a two-input dimensional reduction: ```julia -rows = AK.mapreduce((x, y) -> x * y, +, a, b; init=0f0, dims=1) +rows = AK.mapreduce((x, y) -> x * y, +, a, b; dims=1) ``` """ function mapreduce( f, op, src::MapReduceSource, srcs::AbstractArray...; backend::Union{Nothing, Backend}=nothing, - init, - neutral=neutral_element(op, typeof(init)), - dims=nothing, + init=_NoInit(), + neutral=nothing, + acctype=nothing, + dims=:, alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, ) @@ -135,41 +128,188 @@ function mapreduce( src isa AbstractArray || throw(ArgumentError("a Broadcasted source cannot be combined with more arrays")) _mapreduce_check_map_axes(src, srcs...) - src = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) - f = identity + bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) + return _mapreduce(identity, op, bc, backend, init, neutral, acctype, dims, alg, temp) end + return _mapreduce(f, op, src, backend, init, neutral, acctype, dims, alg, temp) +end + +function _mapreduce(f, op, src, backend, init, neutral, acctype, dims, alg, temp) backend = _resolve_backend(backend, src, temp) - a = _resolve_reduce(alg, backend, typeof(init), dims) - _mapreduce_impl(f, op, src, backend, a; init, neutral, dims, temp) + M = _mapped_eltype(f, src) + + if _whole(dims) + # Scalar reductions accumulate from `init`'s type + A = _acctype(op, init isa _NoInit ? Union{} : typeof(init), M, acctype) + a = _resolve_reduce(alg, backend, A, dims) + return _mapreduce_whole(f, op, src, backend, a, A; init, neutral, temp) + end + + temp === nothing || throw(ArgumentError( + "`temp` only applies to whole-array reductions; to reduce into an existing array, " * + "use `mapreducedim!`")) + dims_valid = _reduced_dims(dims, ndims(src)) + dst_sizes = ntuple(d -> d in dims_valid ? 1 : _srcsize(src, d), ndims(src)) + if init isa _NoInit && Base.any(d -> _srcsize(src, d) == 0, dims_valid) && + Base.prod(dst_sizes) > 0 + throw(ArgumentError( + "reducing over an empty dimension is not allowed without `init`; pass `init`")) + end + # The result has the accumulator type; where there is none (`op` always throws), only an + # empty reduction can succeed + A = _acctype(op, init isa _NoInit ? Union{} : typeof(init), M, acctype) + R_type = A !== Union{} ? A : !(init isa _NoInit) ? typeof(init) : M === Union{} ? Nothing : M + dst = KernelAbstractions.allocate(backend, R_type, dst_sizes) + return _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, A; + init, neutral, overwrite=true) end -function _mapreduce_impl( - f, op, src::MapReduceSource, backend::Backend, alg; - init, neutral, dims, temp, +""" + mapreducedim!( + f, op, R::AbstractArray, src; + backend=get_backend(R), + init=, + neutral=nothing, + overwrite::Bool=false, + acctype=nothing, + alg::Algorithm=Auto(), + ) -> R + +Reduce `src` (an array or a `Broadcasted` object) into `R`: the dimensions along which `R` has +size 1 and `src` does not are reduced, and the others must match (`R` may leave off trailing +dimensions of size 1). `R` must not alias `src`. + +This is the contract of every AcceleratedKernels reduction, modelled on CUB's rather than on +Base's (see [Differences from Base](@ref)): + +- **Algebra.** `op` must be associative and commutative. Like every GPU reduction, AK combines + elements in an order that depends on the algorithm, its settings and the shape, not in element + order; for a given setting and shape the order is fixed, so results are reproducible on one + device. Non-commutative operators such as string concatenation or matrix products are not + supported. +- **Partial results** start from `neutral`, which must be a two-sided identity of `op` and never + appears in the result. It defaults to `GPUArraysCore.neutral_element(op, T)` where that is + defined; otherwise each partial result starts from its first element, + `Base.mapreduce_first(f, op, x)`, so no neutral element is needed. +- **Accumulator type.** Partial results have one type, and are converted to `eltype(R)` only when + stored. It is `acctype` when that is given; otherwise the type the fold of `op` settles on, + starting from `eltype(R)` and the mapped elements. Every element is also a one-element partial + result, and partial results are combined with each other, so those types are joined in as well + (`R`'s values and `init` are applied once, and are not partial results). A sum of `Int8`s + accumulates in `Int`, and `Float32`s reduced into a `Float64` array accumulate in `Float64`. + An `acctype` that cannot hold the partial results at all is an `ArgumentError`; whether their + values fit is the caller's obligation (a value that does not fit throws an `InexactError`, + where the backend reports errors thrown in kernels). +- **Result.** For each output whose slice is not empty, with `partial` its reduction: + `R[i] = op(init, partial)` when `init` is given (applied once); else `partial` with + `overwrite=true`; else `op(R[i], partial)`, folding into `R`'s previous value as + `Base.mapreducedim!` does. +- **Empty slices.** An output whose slice is empty is set to `init` when that is given, and is + otherwise not written, with or without `overwrite`. +- **Errors.** Mismatched shapes are a `DimensionMismatch`. An `R` that aliases `src`, an `op` + that inference shows always throws for these types (when there is something to combine; an + explicit `acctype` that `op` cannot combine at all is rejected whatever the input), a + non-bits accumulator type for a kernel algorithm, and an algorithm that cannot run on the + backend are `ArgumentError`s. + +```julia +import AcceleratedKernels as AK +using CUDA + +A = CuArray([1 3; 2 4]) +R = CuArray([10 20]) +AK.mapreducedim!(identity, +, R, A) # [13 27] +AK.mapreducedim!(identity, +, R, A; overwrite=true) # [3 7] +AK.mapreducedim!(identity, +, R, A; init=100) # [103 107] +``` +""" +function mapreducedim!( + f, op, R::AbstractArray, src::MapReduceSource; + backend::Union{Nothing, Backend}=nothing, + init=_NoInit(), + neutral=nothing, + overwrite::Bool=false, + acctype=nothing, + alg::Algorithm=Auto(), ) + backend = _resolve_backend(backend, R, src) + nd = ndims(src) + for d in 1:max(nd, ndims(R)) + sR, sA = size(R, d), _srcsize(src, d) + sR == 1 || sR == sA || throw(DimensionMismatch( + "cannot reduce an array of size $(size(src)) into one of size $(size(R))")) + end + _check_noalias(R, src) + dst_sizes = ntuple(d -> size(R, d), nd) + dst = size(R) == dst_sizes ? R : reshape(R, dst_sizes) + dims_valid = Tuple(d for d in 1:nd if dst_sizes[d] == 1 && _srcsize(src, d) != 1) + A = _acctype(op, eltype(dst), _mapped_eltype(f, src), acctype) + _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, A; init, neutral, overwrite) + return R +end + +function _check_noalias(R, src::AbstractArray) + Base.mightalias(R, src) && + throw(ArgumentError("the destination of a reduction must not alias its source")) + nothing +end +_check_noalias(R, src::Base.Broadcast.Broadcasted) = foreach(a -> _check_noalias(R, a), src.args) +_check_noalias(R, src::Base.Broadcast.Extruded) = _check_noalias(R, src.x) +_check_noalias(R, src) = nothing + +# Reduce `src` into `dst`, which has `ndims(src)` dimensions and size 1 along `dims_valid`, with +# accumulator type `A`. +function _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, ::Type{A}; + init, neutral, overwrite) where {A} # scalar *linear* indexing into a multidimensional Broadcasted object is # only available on Julia 1.12; on earlier version, materialize it first. if VERSION < v"1.12-" && src isa Base.Broadcast.Broadcasted - src = Base.Broadcast.materialize(src) + src = _materialize_source(src) + end + a = _resolve_reduce(alg, backend, A, dims_valid) + init = init isa _NoInit && !overwrite ? _Fold() : init + mapreduce_nd!(f, op, dst, src, backend, a, A; init, neutral, dims_valid) + return dst +end + +# Reduce all of `src` to a host value, with accumulator type `A`. +function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, temp) where {A} + if length(src) == 0 + init isa _NoInit || return init + throw(ArgumentError( + "reducing over an empty collection is not allowed without `init`; pass `init`")) end + # (a single element without `init` has nothing to combine, even when `op` always throws) + if A === Union{} && length(src) == 1 && init isa _NoInit + return @allowscalar Base.mapreduce_first(f, op, src[first(eachindex(src))]) + end + _check_acctype(op, f, A) + neutral = _reduce_seed(op, A, neutral) - if _whole(dims) - if alg isa BlockReduce - mapreduce_1d_gpu( - f, op, src, backend; - init, neutral, - block_size=alg.block_size, items_per_thread=alg.items_per_thread, - temp, switch_below=alg.switch_below, - ) - else - mapreduce_1d_cpu( - f, op, src, backend; - init, neutral, - max_tasks=alg.max_tasks, min_elems=alg.min_elems, - ) - end + # scalar *linear* indexing into a multidimensional Broadcasted object is + # only available on Julia 1.12; on earlier version, materialize it first. + if VERSION < v"1.12-" && src isa Base.Broadcast.Broadcasted + src = _materialize_source(src) + end + + # The result, `op(init, partial)` or `partial`, has the accumulator type + return convert(A, _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, temp)) +end + +function _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, temp) + if alg isa BlockReduce + mapreduce_1d_gpu( + f, op, src, backend; + init, neutral, + block_size=alg.block_size, items_per_thread=alg.items_per_thread, + temp, switch_below=alg.switch_below, + ) else - mapreduce_nd(f, op, src, backend, alg; init, neutral, dims, temp) + mapreduce_1d_cpu( + f, op, src, backend; + init, neutral, + max_tasks=alg.max_tasks, min_elems=alg.min_elems, + ) end end diff --git a/src/reduce/tuning.jl b/src/reduce/tuning.jl index 94ce27e2..25f081f8 100644 --- a/src/reduce/tuning.jl +++ b/src/reduce/tuning.jl @@ -71,7 +71,7 @@ function _resolve_reduce(alg::Algorithm, backend::Backend, ::Type{T}, dims) wher t = reduce_tuning(backend, T) a = alg isa Auto ? _select_reduce(backend) : alg a = _fill(a, t, T) - _check_reduce(a, backend, dims) + _check_reduce(a, backend, T, dims) return a end @@ -95,9 +95,13 @@ _fill(a::CPUThreads.Partitioned, t::ReduceTuning, T) = _fill_threads(a, t) _fill(a::Algorithm, t::ReduceTuning, T) = throw(ArgumentError("$(_algname(a)) is not a reduction algorithm")) -function _check_reduce(a::BlockReduce, backend, dims) +function _check_reduce(a::BlockReduce, backend, ::Type{T}, dims) where {T} _checkdomain(a) _require_kernels(a, backend) + # `Union{}`: no accumulator type, which only an empty reduction gets past + T === Union{} || isbitstype(T) || throw(ArgumentError( + "BlockReduce: the accumulator type $T is not a bits type; pass an `init` of the type " * + "to accumulate in, or make `f` and `op` inferable")) if _whole(dims) # Each pass reduces tiles of `block_size * items_per_thread` elements to one value: a tile # of one element never shrinks the input, and the bound keeps the kernels' index @@ -110,7 +114,7 @@ function _check_reduce(a::BlockReduce, backend, dims) nothing end -function _check_reduce(a::CPUThreads.Partitioned, backend, dims) +function _check_reduce(a::CPUThreads.Partitioned, backend, T, dims) _checkdomain(a) _require_threads(a, backend) end diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index a4f55d97..f89b2229 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -1,17 +1,164 @@ -# neutral_element moved over from GPUArrays.jl -neutral_element(op, T) = - error("""AcceleratedKernels.jl needs to know the neutral element for your operator `$op`. - Please pass it as an explicit keyword argument `neutral`.""") -neutral_element(::typeof(Base.:(|)), T) = zero(T) -neutral_element(::typeof(Base.:(⊻)), T) = zero(T) -neutral_element(::typeof(Base.:(&)), T) = ~zero(T) -neutral_element(::typeof(Base.:(+)), T) = zero(T) -neutral_element(::typeof(Base.add_sum), T) = zero(T) -neutral_element(::typeof(Base.:(*)), T) = one(T) -neutral_element(::typeof(Base.mul_prod), T) = one(T) -neutral_element(::typeof(Base.min), T) = typemax(T) -neutral_element(::typeof(Base.max), T) = typemin(T) -neutral_element(::typeof(Base._extrema_rf), ::Type{<:NTuple{2,T}}) where {T} = typemax(T), typemin(T) +# The reduction contract's machinery (see `mapreducedim!`): the accumulator type, the value that +# seeds partial results, and how a finished partial becomes an output value. + +# The mapped element type of a reduction source. Base's `combine_eltypes` does not see through +# the `Extruded` wrappers of a preprocessed `Broadcasted`, so its element types are derived here. +_mapped_eltype(f, src) = Base.promote_op(f, _source_eltype(src)) +_source_eltype(src::AbstractArray) = eltype(src) +_source_eltype(bc::Base.Broadcast.Broadcasted) = Base.promote_op(bc.f, Base.map(_source_eltype, bc.args)...) +_source_eltype(x::Base.Broadcast.Extruded) = eltype(x.x) +_source_eltype(x) = Base.Broadcast._broadcast_getindex_eltype(x) +# A `Broadcasted` source as a dense array of its element type (`materialize` would take the +# element type from `combine_eltypes`, which gives `Any` for `Extruded` arguments) +_materialize_source(bc::Base.Broadcast.Broadcasted) = copyto!(similar(bc, _source_eltype(bc)), bc) + +# The type of a reduction's first partial value when no `init` is given. +_first_type(op, ::Type{M}) where {M} = Base.promote_op(Base.reduce_first, Core.Typeof(op), M) + +# The accumulator type of a reduction of mapped elements of type `M` into `S` (the type of `init` +# or of the destination, `Union{}` for neither): the types Base's fold `op(op(S, M), M)...` goes +# through, joined with `promote_type`. In a parallel reduction every element is also a +# one-element partial result (`Base.reduce_first(op, x)`), and partial results are combined with +# each other, so those types are joined as well; `S` itself is not a partial result (`init` and +# the destination's values are applied once, at the end). So `sum(Int8[...])` accumulates in +# `Int`, as Base's `add_sum` does, and partial results have one type however many elements they +# combine. `Union{}` when inference shows that the fold always throws; only a reduction with +# elements to combine is an error then (`_check_acctype`). +function _reduce_acctype(op, ::Type{S}, ::Type{M}) where {S, M} + A = _first_type(op, M) + if S !== Union{} + B = Base.promote_op(op, S, M) + B === Union{} && return Union{} + A = A === Union{} ? B : promote_type(A, B) + end + A === Union{} && return Union{} + for _ in 1:8 + B = Base.promote_op(op, A, M) + B === Union{} && return Union{} + B = promote_type(A, B) + C = Base.promote_op(op, B, B) + C === Union{} || (B = promote_type(B, C)) + B == A && return A + A = B + end + return A +end + +_check_acctype(op, f, ::Type{A}) where {A} = A === Union{} && throw(ArgumentError( + "`$op` cannot combine the mapped elements (`$f`): calling it always throws for their types")) + +# The accumulator type of a reduction: the caller's `acctype`, else the rule of +# `_reduce_acctype`. An `acctype` is rejected only where the types rule it out, when a value that +# a partial result takes (`Base.reduce_first`'s of an element, or `op`'s of a partial result and +# an element or of two partial results) has no conversion to it at all; whether the values fit is +# the caller's obligation. +_acctype(op, ::Type{S}, ::Type{M}, ::Nothing) where {S, M} = _reduce_acctype(op, S, M) +function _acctype(op, ::Type{S}, ::Type{M}, ::Type{A}) where {S, M, A} + M === Union{} && return A # (no elements to hold) + _check_holds(A, _first_type(op, M), "a one-element partial result (`Base.reduce_first`)") + for (T, what) in ((Base.promote_op(op, A, M), "a partial result and an element"), + (Base.promote_op(op, A, A), "two partial results")) + T === Union{} && throw(ArgumentError("`acctype=$A`: `$op` of $what always throws")) + _check_holds(A, T, "`$op` of $what") + end + return A +end +_check_holds(::Type{A}, ::Type{T}, what) where {A, T} = + T === Union{} || Base.promote_op(convert, Type{A}, T) !== Union{} || + throw(ArgumentError("`acctype=$A` cannot hold $what, of type `$T`")) +_acctype(op, ::Type, ::Type, acctype) = + throw(ArgumentError("`acctype` must be a type or `nothing`, got $(repr(acctype))")) + +# A partial result of a reduction whose operator has no known neutral element. It starts out +# empty and takes its first value from `Base.mapreduce_first`; combining with an empty lane +# returns the other operand. For associative and commutative operators this gives the same result +# as seeding with a neutral element. +# +# An empty lane of a bits type (every lane kernels see) holds zero bits; one of another type leaves +# the value undefined, which is how it is recognised. +struct _Lane{T} + value::T + valid::Bool + _Lane{T}() where {T} = _inline_value(T) ? new{T}(_zero_bits(T), false) : new{T}() + _Lane{T}(value) where {T} = new{T}(value, true) + _Lane{T}(value, valid::Bool) where {T} = new{T}(value, valid) +end + +_inline_value(::Type{T}) where {T} = isbitstype(T) || Base.isbitsunion(T) +@inline _valid(l::_Lane{T}) where {T} = _inline_value(T) ? l.valid : isdefined(l, :value) + +# A value of the bits type or bits union `T` with all-zero bits +function _zero_bits(::Type{T}) where {T} + isbitstype(T) || return _zero_bits(first(Base.uniontypes(T))) + Base.issingletontype(T) && return T.instance + bytes = zeros(UInt8, sizeof(T)) + return GC.@preserve bytes unsafe_load(Ptr{T}(pointer(bytes))) +end + +struct _LaneMap{F, OP, T} + f::F + op::OP +end +_LaneMap{T}(f::F, op::OP) where {T, F, OP} = _LaneMap{F, OP, T}(f, op) +@inline (m::_LaneMap{F, OP, T})(x) where {F, OP, T} = _Lane{T}(Base.mapreduce_first(m.f, m.op, x)) + +struct _LaneOp{OP} + op::OP +end +@inline function (o::_LaneOp)(a::_Lane{T}, b::_Lane{T}) where {T} + _valid(a) || return b + _valid(b) || return a + return _Lane{T}(o.op(a.value, b.value)) +end + +# The map `f` of a reduction pass whose inputs are already partial results. +struct _Partials end + +# `op`, with its result converted to the accumulator type `A`: partial results keep that type +# (the accumulator rule makes this a no-op; a narrower `acctype` needs it) +struct _AccOp{A, OP} + op::OP +end +_AccOp{A}(op::OP) where {A, OP} = _AccOp{A, OP}(op) +@inline (o::_AccOp{A})(a, b) where {A} = convert(A, o.op(a, b)) + +# The map and operator a reduction applies for the given seed: `f` and `op` (converting to the +# seed's type, the accumulator type) after a neutral element, lane-wrapped after an empty `_Lane`. +# Kernels call this on the device, so the wrappers never need to be converted for a backend. +@inline _lanefuncs(f, op, neutral::A) where {A} = (f, _AccOp{A}(op)) +@inline _lanefuncs(::_Partials, op, neutral::A) where {A} = (identity, _AccOp{A}(op)) +@inline _lanefuncs(f, op, ::_Lane{T}) where {T} = (_LaneMap{T}(f, op), _LaneOp(op)) +@inline _lanefuncs(::_Partials, op, ::_Lane) = (identity, _LaneOp(op)) + +@inline _unlane(x) = x +@inline _unlane(x::_Lane) = x.value + +# The seed of every partial result with accumulator type `A`: the caller's `neutral`, else +# GPUArraysCore's neutral element for `op` when it has one, else an empty lane. A seed that is not +# a lane stands for the accumulator type, so an abstract one (only on the host) is a lane: holding +# the caller's `neutral`, or empty. +function _reduce_seed(op, ::Type{A}, neutral) where {A} + isconcretetype(A) || return neutral === nothing ? _Lane{A}() : _Lane{A}(neutral) + neutral === nothing || return convert(A, neutral) + Base.promote_op(neutral_element, Core.Typeof(op), Type{A}) === Union{} && return _Lane{A}() + return _exact_neutral(op, convert(A, neutral_element(op, A))) +end + +# GPUArraysCore's neutral element of `+` is `zero(T)`, which turns a floating-point sum of +# negative zeros into a positive zero; the identity of floating-point addition is `-0.0` +_exact_neutral(op, n) = n +_exact_neutral(::Union{typeof(+), typeof(Base.add_sum)}, + n::Union{AbstractFloat, Complex{<:AbstractFloat}}) = iszero(n) ? -n : n + +# `init` of the kernels that write a destination: a value, `_NoInit()`, or `_Fold()` to fold in +# the destination's previous value as `Base.mapreducedim!` does. +struct _Fold end + +# The value stored for an output whose reduction produced `partial`: `init` is applied exactly +# once, here. +@inline _finish(op, init, dst, i, partial) = op(init, _unlane(partial)) +@inline _finish(op, ::_NoInit, dst, i, partial) = _unlane(partial) +@inline _finish(op, ::_Fold, dst, i, partial) = op(dst[i], _unlane(partial)) # Unrolled map constructing a tuple @@ -31,10 +178,12 @@ end end -# Apply op(init, f(x)) to each element of src, storing the result in dst. -function _mapreduce_nd_apply_init!(f, op, dst, src, backend; init, launch...) +# Reductions whose every output reduces a single element, a partial result of the accumulator +# type `A`: `dst` and `src` have the same length. +function _mapreduce_nd_single!(f, op, dst, src, backend, ::Type{A}; init, launch...) where {A} _foreachindex(eachindex(dst), backend; launch...) do i - dst[i] = op(init, f(src[i])) + x = Base.mapreduce_first(f, op, src[i]) + dst[i] = _finish(op, init, dst, i, A === Union{} ? x : convert(A, x)) end end diff --git a/src/utils.jl b/src/utils.jl index 46fd96cc..bd4c446c 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -18,6 +18,10 @@ end # provide (`CL_DEVICE_LOCAL_MEM_SIZE`, full profile), and Metal's threadgroup memory limit. const LOCAL_MEMORY_BUDGET = 32 * 1024 +# The default of every `init` keyword: no initial value given. As in Base, this differs from +# `init=nothing`, which is an explicit initial value. +struct _NoInit end + # Keep the default shared-memory use within the local-memory budget. @inline function default_scan_items_per_thread(backend, ::Type{T}, block_size) where T block_size > 0 || return default_scan_items_per_thread(backend) diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index dd537064..7e01bd5d 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -151,7 +151,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int64(0), - temp=array_from_host(zeros(Int32, 10_000)), + temp=array_from_host(zeros(Int64, 10_000)), ) AK.reduce( (x, y) -> x + 1, @@ -378,23 +378,23 @@ end end # Testing different settings - AK.reduce( + AK.mapreducedim!( + identity, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 1, 5)), array_from_host(rand(Int32, 3, 4, 5)); alg=reduce_alg(block_size=64, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int32(0), - dims=2, - temp=array_from_host(zeros(Int32, 3, 1, 5)), ) - AK.reduce( + AK.mapreducedim!( + identity, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 4, 1)), array_from_host(rand(Int32, 3, 4, 5)); alg=reduce_alg(block_size=64, max_tasks=16, min_elems=1000), init=Int32(0), neutral=Int32(0), - dims=3, - temp=array_from_host(zeros(Int32, 3, 4, 1)), ) end @@ -798,25 +798,23 @@ end end # Testing different settings - AK.mapreduce( + AK.mapreducedim!( -, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 1, 5)), array_from_host(rand(Int32, 3, 4, 5)); alg=reduce_alg(block_size=64, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int32(0), - dims=2, - temp=array_from_host(zeros(Int32, 3, 1, 5)), ) - AK.mapreduce( + AK.mapreducedim!( -, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 4, 1)), array_from_host(rand(Int32, 3, 4, 5)); alg=reduce_alg(block_size=64, max_tasks=16, min_elems=1000), init=Int32(0), neutral=Int32(0), - dims=3, - temp=array_from_host(zeros(Int32, 3, 4, 1)), ) end @testset "sum" begin @@ -1087,3 +1085,242 @@ end end end end + + +@testset "reduction contract" begin + Random.seed!(0) + alg = REDUCE_ALG + + # Omitted `init`: an empty reduction is an error, except for `sum`, `prod` and `count`, which + # give zero or one of the accumulator type + @test AK.sum(array_from_host(Int32[]); alg) === 0 + @test AK.sum(array_from_host(Float32[]); alg) === 0.0f0 + @test AK.sum(array_from_host(Int32[]); acctype=Int32, alg) === Int32(0) + @test AK.prod(array_from_host(Int32[]); alg) === 1 + @test AK.count(array_from_host(Bool[]); alg) === 0 + @test_throws ArgumentError AK.minimum(array_from_host(Int32[]); alg) + @test_throws ArgumentError AK.reduce(max, array_from_host(Int32[]); alg) + @test_throws ArgumentError AK.reduce(+, array_from_host(Int32[]); alg) + @test_throws ArgumentError AK.reduce((a, b) -> a + b, array_from_host(Int32[]); alg) + # An explicit `init` is returned as it is for an empty reduction, and applied once otherwise, + # also when it is not a neutral element + @test AK.sum(array_from_host(Float32[]); init=1, alg) === 1 + @test AK.maximum(array_from_host(Int32[]); init=Int32(-1), alg) === Int32(-1) + @test AK.sum(array_from_host(ones(Int32, 1000)); init=Int32(10), alg) === 1010 + @test AK.reduce(*, array_from_host(fill(Int32(2), 20)); init=Int32(3), alg) === Int32(3 << 20) + # `init=nothing` is an initial value, not an omitted one; `+` cannot fold it, which only + # matters when there is something to fold + @test_throws Exception AK.reduce(+, array_from_host(Int32[1, 2]); init=nothing, alg) + @test AK.reduce(+, array_from_host(Int32[]); init=nothing, alg) === nothing + @test AK.mapreduce(x -> error("never called"), +, array_from_host(Int32[]); init=7, alg) === 7 + + # The accumulator type is the fold type of `op`, from `init`'s type and the mapped elements + h8 = Int8[100, 100, 27] + v8 = array_from_host(h8) + @test AK.sum(v8; alg) === 227 + @test AK.reduce(+, v8; alg) === Int8(-29) + @test AK.reduce(+, v8; init=0, alg) === 227 + @test AK.sum(array_from_host([1, 2]); init=Int8(0), alg) === 3 + @test AK.reduce((a, b) -> Base.add_sum(a, b), v8; alg) === 227 + @test AK.maximum(v8; alg) === Int8(100) + # ... joined with the one-element partial results': here the element type, not the narrower + # one the fold produces + @test AK.reduce(Returns(Int8(0)), array_from_host(Int32[5, 6, 7]); alg) === Int32(0) + if !TEST_KERNELS + @test AK.reduce(coalesce, [missing, 1, 2]; init=0, alg) === 0 + end + # Partial results of small integers are `Int`s, also for mixed signedness (where Julia 1.10's + # `add_sum` of the `init` and an element would stay a `UInt8`) + @test AK.sum(array_from_host(Int8[-1, -2]); init=UInt8(0), alg) === -3 + @test AK.count(array_from_host([true, true]); init=UInt8(0), alg) === 2 + # A single element has the accumulator type too; `f` may index device arrays + @test AK.reduce((a, b) -> a + b, array_from_host([true]); alg) === 1 + @test AK.sum(array_from_host([true]); alg) === 1 + w = array_from_host(Int32[11, 22]) + @test AK.mapreduce(let w = w; x -> w[x]; end, +, array_from_host(Int32[2]); alg) === Int32(22) + # Seeds are exact identities, and `sum`'s empty rule is no `init`: a sum of negative zeros is + # a negative zero + for n in (1, 2, 1000) + @test AK.sum(array_from_host(fill(-0.0f0, n)); alg) === -0.0f0 + @test Base.all(signbit, Array(AK.sum(array_from_host(fill(-0.0f0, 2, n)); dims=2, alg))) + end + + # `acctype` sets the accumulator type, and the result's + @test AK.reduce(+, array_from_host(Int8[100, 100, 56]); acctype=Int8, alg) === Int8(0) + @test AK.sum(array_from_host(Int8[100, 100, 56]); acctype=Int16, alg) === Int16(256) + @test AK.reduce(max, v8; init=Int8(0), acctype=Int32, alg) === Int32(100) + xf = Float32[1.0f8, 1, -1.0f8] + if Float64 in valid_backend_eltypes(BACKEND, (Float64,)) + @test AK.sum(array_from_host(xf); acctype=Float64, alg) === 1.0 + R = array_from_host(zeros(Float32, 1, 2)) + AK.mapreducedim!(identity, +, R, array_from_host([xf xf]); overwrite=true, + acctype=Float64, alg) + @test Array(R) == [1 1] + end + m8 = rand(Int8(-9):Int8(9), 40, 30) + dm8 = array_from_host(m8) + for dims in (1, 2) + r = AK.sum(dm8; dims, acctype=Int16, alg) + @test eltype(r) === Int16 && Array(r) == sum(m8; dims) + end + # ... which `init` does not change (only an empty reduction returns it as it is) + @test AK.sum(array_from_host(Int32[1, 2]); init=0, acctype=Int32, alg) === Int32(3) + @test AK.count(array_from_host([true, true]); acctype=Int32, alg) === Int32(2) + @test AK.sum(array_from_host(Int32[]); init=0, acctype=Int32, alg) === 0 + if !TEST_KERNELS + @test_throws InexactError AK.reduce(+, Int8[1, 2]; init=0.5, acctype=Int16, alg) + end + # ... and must be able to hold the partial results + @test_throws ArgumentError AK.sum(v8; acctype=String, alg) + @test_throws ArgumentError AK.reduce((a, b) -> "", v8; acctype=Int, alg) + @test_throws ArgumentError AK.sum(v8; acctype=1, alg) + # (whether their values fit is the caller's obligation) + if !TEST_KERNELS + @test_throws InexactError AK.sum([100, 100]; acctype=Int8, alg) + end + + # Reductions along `dims` return the accumulator type, also with `init` + for dims in (1, 2, (1, 2), 3) + r = AK.sum(dm8; dims, alg) + @test eltype(r) === Int && Array(r) == sum(m8; dims) + r = AK.sum(dm8; dims, init=Int16(1), alg) + @test eltype(r) === Int && Array(r) == sum(m8; dims, init=Int16(1)) + r = AK.maximum(dm8; dims, alg) + @test eltype(r) === Int8 && Array(r) == maximum(m8; dims) + r = AK.count(x -> x > 0, dm8; dims, alg) + @test eltype(r) === Int && Array(r) == count(x -> x > 0, m8; dims) + end + + # Empty reduced dimensions: `init`, else an error, except for `sum`, `prod` and `count` + e8 = array_from_host(zeros(Int8, 0, 3)) + for (r, value) in ((AK.sum(e8; dims=1, alg), 0), (AK.prod(e8; dims=1, alg), 1), + (AK.count(x -> x > 0, e8; dims=1, alg), 0)) + @test eltype(r) === Int && Array(r) == fill(value, 1, 3) + end + @test eltype(AK.sum(e8; dims=1, acctype=Int8, alg)) === Int8 + @test Array(AK.minimum(e8; dims=1, init=Int8(7), alg)) == fill(Int8(7), 1, 3) + @test_throws ArgumentError AK.minimum(e8; dims=1, alg) + @test_throws ArgumentError AK.mapreduce(x -> x + 1, +, array_from_host(zeros(Int32, 0, 2)); + dims=1, alg) + @test size(AK.maximum(array_from_host(zeros(Int8, 3, 0)); dims=1, alg)) == (1, 0) + # ... an error only where an output has an empty slice + @test size(AK.minimum(array_from_host(zeros(Int32, 0, 0)); dims=1, alg)) == (1, 0) + + # An operator that always throws for the element types: an error only where something is + # combined + throws(a, b) = throw(ArgumentError("never called")) + @test AK.reduce(throws, array_from_host(Int32[5]); alg) === Int32(5) + x5 = array_from_host(Int32[5]) + @test AK.mapreduce(+, throws, x5, x5; alg) === Int32(10) + @test Array(AK.reduce(throws, array_from_host(Int32[5 6]); dims=1, alg)) == [5 6] + @test_throws ArgumentError AK.reduce(throws, array_from_host(Int32[5, 6]); alg) + @test_throws ArgumentError AK.reduce(throws, array_from_host(Int32[5, 6]); dims=1, alg) + # On the host, an abstract accumulator type with the caller's neutral element + if !TEST_KERNELS + @test AK.reduce(+, Real[1, 2.5, 3]; neutral=0, alg=AK.CPUThreads.Partitioned(max_tasks=2, min_elems=1)) == 6.5 + end + + # Operators without a known neutral element: partial results start from their first element + x = rand(Int32(-100):Int32(100), 10_001) + dx = array_from_host(x) + @test AK.reduce((a, b) -> a + b, dx; alg) == sum(x) + @test AK.reduce((a, b) -> min(a, b), dx; alg) == minimum(x) + @test AK.reduce((a, b) -> a + b, dx; init=Int32(10), alg) == sum(x) + 10 + @test AK.reduce((a, b) -> a + b, array_from_host(Int32[7]); alg) === Int32(7) + m = rand(Int32(-100):Int32(100), 37, 300) + for dims in (1, 2, (1, 2), 3) + @test Array(AK.reduce((a, b) -> a + b, array_from_host(m); dims, alg)) == sum(m; dims) + @test Array(AK.mapreduce(abs, (a, b) -> max(a, b), array_from_host(m); dims, alg)) == + maximum(abs, m; dims) + end + + # Tuple and named-tuple accumulators, as GPUArrays uses for `findmin` and Missing-aware `==` + xs = rand(Float32, 5000) + ix = collect(Int32(1):Int32(5000)) + findmin_op(a, b) = (a[1] < b[1] || (a[1] == b[1] && a[2] < b[2])) ? a : b + @test AK.mapreduce(tuple, findmin_op, array_from_host(xs), array_from_host(ix); alg) == + (minimum(xs), Int32(argmin(xs))) + ms = rand(Float32, 20, 300) + im = reshape(collect(Int32(1):Int32(length(ms))), size(ms)) + r = Array(AK.mapreduce(tuple, findmin_op, array_from_host(ms), array_from_host(im); dims=2, alg)) + @test r == reshape([(minimum(ms[i, :]), im[i, argmin(ms[i, :])]) for i in 1:20], 20, 1) + eq_op(a, b) = (is_missing = a.is_missing | b.is_missing, is_equal = a.is_equal & b.is_equal) + @test AK.mapreduce(x -> (is_missing = false, is_equal = x > -101), eq_op, dx; alg) == + (is_missing = false, is_equal = true) + + # `mapreducedim!`: fold into the destination, overwrite it, or apply `init` once + A = array_from_host(Int32[1 3; 2 4]) + R = array_from_host(Int32[10 20]) + @test AK.mapreducedim!(identity, +, R, A; alg) === R + @test Array(R) == [13 27] + R = array_from_host(Int32[10 20]) + @test Array(AK.mapreducedim!(identity, +, R, A; overwrite=true, alg)) == [3 7] + R = array_from_host(Int32[10 20]) + @test Array(AK.mapreducedim!(identity, +, R, A; init=Int32(100), alg)) == [103 107] + R = array_from_host(Int32[10 20]) + @test Array(AK.mapreducedim!(identity, (a, b) -> a + b, R, A; alg)) == [13 27] + R = array_from_host(Int32[10 20]) + @test Array(AK.mapreducedim!(x -> 2x, +, R, A; alg)) == [16 34] + # A destination may leave off trailing singleton dimensions, and have extra ones + @test Array(AK.mapreducedim!(identity, +, array_from_host(zeros(Int32, 2)), A; + overwrite=true, alg)) == [4, 6] + @test Array(AK.mapreducedim!(identity, +, array_from_host(zeros(Int32, 1, 2, 1)), A; + overwrite=true, alg)) == reshape([3, 7], 1, 2, 1) + # No reduced dimension: every output reduces one element + R = array_from_host(Int32[1 1; 1 1]) + @test Array(AK.mapreducedim!(identity, +, R, A; alg)) == [2 4; 3 5] + # Empty slices: `init`, else not written, when folding and overwriting alike + E = array_from_host(zeros(Int32, 0, 2)) + R = array_from_host(Int32[10 20]) + @test Array(AK.mapreducedim!(identity, +, R, E; alg)) == [10 20] + @test Array(AK.mapreducedim!(identity, +, R, E; overwrite=true, alg)) == [10 20] + @test Array(AK.mapreducedim!(identity, min, R, E; overwrite=true, alg)) == [10 20] + @test Array(AK.mapreducedim!(identity, min, R, E; init=Int32(5), alg)) == [5 5] + # The accumulator type starts from the destination's element type + R = array_from_host(zeros(Int16, 1, 3)) + AK.mapreducedim!(identity, +, R, array_from_host(fill(Int8(100), 4, 3)); overwrite=true, alg) + @test Array(R) == fill(Int16(400), 1, 3) + # A Broadcasted source + bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(*, A, Int32(2))) + R = array_from_host(zeros(Int32, 1, 2)) + @test Array(AK.mapreducedim!(identity, +, R, bc; overwrite=true, alg)) == [6 14] + # Shapes and aliasing + @test_throws DimensionMismatch AK.mapreducedim!(identity, +, array_from_host(zeros(Int32, 3, 1)), A; alg) + @test_throws DimensionMismatch AK.mapreducedim!(identity, +, array_from_host(zeros(Int32, 1, 1, 2)), A; alg) + @test_throws ArgumentError AK.mapreducedim!(identity, +, view(A, 1:1, :), A; alg) + bc = Base.Broadcast.preprocess(nothing, Base.Broadcast.instantiate( + Base.Broadcast.broadcasted(identity, A))) + @test_throws ArgumentError AK.mapreducedim!(identity, +, view(A, 1:1, :), bc; + overwrite=true, alg) + # A preprocessed Broadcasted source (with `Extruded` arguments) keeps its element type + R = array_from_host(zeros(Int32, 1, 2)) + @test Array(AK.mapreducedim!(identity, +, R, bc; overwrite=true, alg)) == [3 7] + + # Kernels need a bits-type accumulator + if TEST_KERNELS + @test_throws ArgumentError AK.mapreduce(x -> x > 0 ? x : missing, +, dx; alg) + end + + # A source whose wrappers `@Const` rebuilds on the device: a reshaped view + hr = rand(Int32(0):Int32(9), 50, 40) + dr = array_from_host(hr) + vr, hv = vec(view(dr, 1:40, 1:30)), vec(view(hr, 1:40, 1:30)) + @test AK.reduce(+, vr; alg) == sum(hv) + @test AK.count(!iszero, vr; alg) == count(!iszero, hv) + rr, hrr = reshape(view(dr, 1:40, 1:30), 30, 40), reshape(view(hr, 1:40, 1:30), 30, 40) + for dims in (1, 2) + @test Array(AK.reduce(+, rr; dims, alg)) == sum(hrr; dims) + end + + # A source of a bits-union element type, with every kernel shape (where the backend's arrays + # can hold one) + code(x) = x === missing ? 0x01 : x ? 0x02 : 0x00 + TEST_BITS_UNIONS && for (shape, dims) in (((1000, 300), 1), ((1000, 300), 2), ((256, 256), 2), + ((20_000, 2), (1, 2)), ((5, 3), 3)) + hm = rand([true, false, missing], shape...) + dm = array_from_host(hm) + @test AK.mapreduce(code, max, dm; init=0x00, alg) === mapreduce(code, max, hm) + @test Array(AK.mapreduce(code, max, dm; dims, init=0x00, alg)) == + mapreduce(code, max, hm; dims) + end +end diff --git a/test/generic/reduce_resolve.jl b/test/generic/reduce_resolve.jl index 474287c5..0798a098 100644 --- a/test/generic/reduce_resolve.jl +++ b/test/generic/reduce_resolve.jl @@ -111,9 +111,11 @@ end @testset "reduce resolution: tuning values" begin # A tuning cannot make a reduction along `dims` launch nothing with_reduce_tuning(; target_blocks=0) do - @test_throws ArgumentError AK.mapreduce_nd(identity, +, zeros(Int32, 4, 4), RRB, - AK.BlockReduce(256, 2, 0); - init=Int32(0), dims=1, temp=nothing) + @test_throws ArgumentError AK.mapreduce_nd!(identity, +, zeros(Int32, 1, 4), + zeros(Int32, 4, 4), RRB, + AK.BlockReduce(256, 2, 0), Int32; + init=AK._NoInit(), neutral=nothing, + dims_valid=(1,)) end end @@ -156,3 +158,49 @@ end @test_throws MethodError AK.reduce(+, v; init=Int32(0), prefer_threads=false) @test_throws MethodError AK.sum(v; switch_below=10) end + + +@testset "reduce resolution: accumulator and seed" begin + # The accumulator type follows Base's promotion to a fixed point + @test AK._reduce_acctype(Base.add_sum, Int8, Int8) === Int + @test AK._reduce_acctype(+, Int8, Int8) === Int8 + @test AK._reduce_acctype(+, Bool, Bool) === Int + @test AK._reduce_acctype(+, Int, Float32) === Float32 + @test AK._reduce_acctype(max, Tuple{Int32, Int32}, Tuple{Int32, Int32}) === Tuple{Int32, Int32} + @test AK._first_type(Base.add_sum, Int8) === Int + if VERSION >= v"1.13-" # older `add_sum`s keep mixed-signedness small integers small + @test AK._reduce_acctype(Base.add_sum, UInt8, Int8) === Int + @test AK._reduce_acctype(Base.add_sum, UInt8, Bool) === Int + end + @test AK._reduce_acctype(+, Union{}, Char) === Union{} + @test AK._reduce_acctype(Returns(false), Union{}, Int) === Int + @test AK._reduce_acctype(coalesce, Int, Union{Missing, Int}) === Union{Missing, Int} + # `init` is applied once, not a partial result: only `op(init, x)`'s type counts + @test AK._reduce_acctype((a, b) -> something(a, 0) + something(b, 0), Nothing, Int) === Int + + # Seeds: the caller's neutral, else GPUArraysCore's, else an empty lane + @test AK._reduce_seed(+, Int32, nothing) === Int32(0) + @test AK._reduce_seed(min, Float32, nothing) === Inf32 + @test AK._reduce_seed(+, Float32, nothing) === -0.0f0 + @test AK._reduce_seed(Base.add_sum, ComplexF32, nothing) === complex(-0.0f0, -0.0f0) + @test AK._reduce_seed(+, Int32, 0) === Int32(0) + @test AK._reduce_seed((a, b) -> a + b, Int32, nothing) isa AK._Lane{Int32} + @test !AK._valid(AK._reduce_seed((a, b) -> a + b, Int32, nothing)) + # Empty lanes of every kind of type + @test !AK._valid(AK._Lane{Tuple{Int8, Int64}}()) + @test !AK._valid(AK._Lane{Union{Missing, Int}}()) + @test !AK._valid(AK._Lane{String}()) + @test AK._valid(AK._Lane{String}("x")) + + # Results infer, whether or not the operator has a known neutral element + @test only(Base.return_types(v -> AK.sum(v), (Vector{Int8},))) === Int + @test only(Base.return_types(v -> AK.reduce((a, b) -> a + b, v), (Vector{Int32},))) === Int32 + @test only(Base.return_types(m -> AK.sum(m; dims=1), (Matrix{Int8},))) === Matrix{Int} + @test only(Base.return_types(v -> AK.mapreduce(tuple, (a, b) -> a, v, v), + (Vector{Int32},))) === Tuple{Int32, Int32} + + # BlockReduce needs a bits-type accumulator; the host algorithm does not + @test_throws ArgumentError resolve_reduce(AK.Auto(); T=Union{Missing, Int32}) + @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); T=String) + @test resolve_reduce(AK.Auto(); T=String, backend=AK.HOST_BACKEND) isa AK.CPUThreads.Partitioned +end From 489bbfc85b76cd601facb335b52758ed10848e47 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 07:29:47 +0200 Subject: [PATCH 5/8] Move scans onto algorithm selection `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. --- docs/src/api/accumulate.md | 21 +++ docs/src/api/algorithms.md | 7 +- docs/src/tuning.md | 6 +- ext/AMDGPUExt.jl | 4 + ext/CUDACoreExt.jl | 4 + src/accumulate/accumulate.jl | 224 ++++++++-------------------- src/accumulate/accumulate_1d_cpu.jl | 19 +-- src/accumulate/accumulate_1d_gpu.jl | 30 +--- src/accumulate/accumulate_nd.jl | 25 +--- src/accumulate/tuning.jl | 163 ++++++++++++++++++++ src/algorithms.jl | 19 +++ src/arithmetics.jl | 66 ++------ src/findall.jl | 2 +- src/sort/cpu_sample_sort.jl | 5 +- src/sort/radix_sort.jl | 2 +- src/utils.jl | 13 -- test/generic/accumulate.jl | 155 ++++++++++--------- test/generic/scan_resolve.jl | 145 ++++++++++++++++++ 18 files changed, 537 insertions(+), 373 deletions(-) create mode 100644 src/accumulate/tuning.jl create mode 100644 test/generic/scan_resolve.jl diff --git a/docs/src/api/accumulate.md b/docs/src/api/accumulate.md index 3704620c..0f7af000 100644 --- a/docs/src/api/accumulate.md +++ b/docs/src/api/accumulate.md @@ -5,3 +5,24 @@ AcceleratedKernels.accumulate! AcceleratedKernels.accumulate ``` + +[`Auto()`](@ref AcceleratedKernels.Auto), the default `alg`, uses +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host; on GPUs, +[`ScanPrefixes`](@ref AcceleratedKernels.ScanPrefixes) for whole arrays and +[`SliceScan`](@ref AcceleratedKernels.SliceScan) along `dims`, with the device's settings. Pass an +algorithm to choose it and its settings yourself: + +```@docs +AcceleratedKernels.ScanAlgorithm +AcceleratedKernels.ScanPrefixes +AcceleratedKernels.DecoupledLookback +AcceleratedKernels.SliceScan +``` + +```julia +v = CuArray(rand(Int32(1):Int32(100), 1_000_000)) +AK.accumulate!(+, v; init=Int32(0), alg=AK.ScanPrefixes(block_size=512)) # 512 threads, 8 items per thread +AK.accumulate!(+, v; init=Int32(0), alg=AK.DecoupledLookback()) # CUDA and AMDGPU only +m = CuArray(rand(Int32(1):Int32(100), 100, 10_000)) +AK.accumulate(+, m; init=Int32(0), dims=2, alg=AK.SliceScan(block_size=128)) +``` diff --git a/docs/src/api/algorithms.md b/docs/src/api/algorithms.md index 09bdc288..79a58d76 100644 --- a/docs/src/api/algorithms.md +++ b/docs/src/api/algorithms.md @@ -11,8 +11,8 @@ AcceleratedKernels.SortAlgorithm AcceleratedKernels.CPUThreads ``` -The algorithm types of each operation are documented with it, for [sorting](sort.md) and for -[reductions](reduce.md). +The algorithm types of each operation are documented with it, for [sorting](sort.md), +[reductions](reduce.md) and [scans](accumulate.md). #### How `Auto` chooses @@ -21,7 +21,8 @@ slices along `dims`, and their length) and the ordering, never at the array's co host backend it chooses the family's `CPUThreads` algorithm, which runs on Julia threads. On a GPU it follows per-device tuning values: for sorting, `BitonicSort` for short arrays and slices, `RadixSort` for long whole arrays of the element types it supports, and `MergeSort` otherwise, -subject to `Auto`'s requirements (`stable=true` by default); for reductions, `BlockReduce`. +subject to `Auto`'s requirements (`stable=true` by default); for reductions, `BlockReduce`; for +scans, `ScanPrefixes` for whole arrays and `SliceScan` along `dims`. Algorithms carry their settings as fields, and a field left at `nothing` takes the device's tuned value: diff --git a/docs/src/tuning.md b/docs/src/tuning.md index 3da43486..48ee2acb 100644 --- a/docs/src/tuning.md +++ b/docs/src/tuning.md @@ -21,10 +21,11 @@ AcceleratedKernels._resolve_sort ``` The other families follow the same steps, with a one-line selection: reductions resolve with -`_resolve_reduce`. +`_resolve_reduce`, scans with `_resolve_scan`. ```@docs AcceleratedKernels._resolve_reduce +AcceleratedKernels._resolve_scan ``` `Auto` and explicit algorithms share steps 3 and 4, so a tuning cannot make an invalid algorithm @@ -40,6 +41,8 @@ AcceleratedKernels.SortTuning AcceleratedKernels.sort_tuning AcceleratedKernels.ReduceTuning AcceleratedKernels.reduce_tuning +AcceleratedKernels.ScanTuning +AcceleratedKernels.scan_tuning ``` AcceleratedKernels defines the hook's generic method, whose values reproduce the library's @@ -56,5 +59,6 @@ alike; tunings cannot change them. ```@docs AcceleratedKernels._runs_threads AcceleratedKernels._runs_kernels +AcceleratedKernels._supports_lookback AcceleratedKernels._resolve_backend ``` diff --git a/ext/AMDGPUExt.jl b/ext/AMDGPUExt.jl index 92758ef2..c43045d5 100644 --- a/ext/AMDGPUExt.jl +++ b/ext/AMDGPUExt.jl @@ -8,4 +8,8 @@ import AcceleratedKernels as AK AMDGPU.Device.@device_override AK._decoupled_fence() = UnsafeAtomics.fence(UnsafeAtomics.seq_cst, AMDGPU.syncscope_agent) +# The fence above orders device-scope memory, and blocks make forward progress while others +# wait on them: DecoupledLookback is correct here. +AK._supports_lookback(::ROCBackend) = true + end diff --git a/ext/CUDACoreExt.jl b/ext/CUDACoreExt.jl index 9f0a3c25..352a7e47 100644 --- a/ext/CUDACoreExt.jl +++ b/ext/CUDACoreExt.jl @@ -6,4 +6,8 @@ import AcceleratedKernels as AK # Device-scope fence for the DecoupledLookback scan. CUDACore.@device_override AK._decoupled_fence() = CUDACore.threadfence() +# The fence above orders device-scope memory, and blocks make forward progress while others +# wait on them: DecoupledLookback is correct here. +AK._supports_lookback(::CUDABackend) = true + end diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index fdd5acd0..0e7ba211 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -1,7 +1,4 @@ -# Available accumulation algorithms -abstract type AccumulateAlgorithm end -struct DecoupledLookback <: AccumulateAlgorithm end -struct ScanPrefixes <: AccumulateAlgorithm end +include("tuning.jl") # Helpers @@ -28,51 +25,25 @@ include("accumulate_nd.jl") """ - accumulate!( - op, v::AbstractArray, backend::Backend=get_backend(v); - init, - neutral=neutral_element(op, eltype(v)), - dims::Union{Nothing, Int}=nothing, - inclusive::Bool=true, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, - - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), - - # GPU settings - block_size::Int=256, - items_per_thread::Union{Nothing, Int}=nothing, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, - ) + accumulate!(op, v::AbstractArray; kwargs...) -> v + accumulate!(op, dst::AbstractArray, src::AbstractArray; kwargs...) -> dst accumulate!( - op, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(dst); + op, dst, src; + backend=nothing, init, neutral=neutral_element(op, eltype(dst)), - dims::Union{Nothing, Int}=nothing, + dims::Union{Nothing, Integer}=nothing, inclusive::Bool=true, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, - - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), - - # GPU settings - block_size::Int=256, - items_per_thread::Union{Nothing, Int}=nothing, + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_flags::Union{Nothing, AbstractArray}=nothing, ) Compute accumulated running totals along a sequence by applying a binary operator to all elements up to the current one; often used in GPU programming as a first step in finding / extracting -subsets of data. +subsets of data. The first form scans `v` in place; the second writes the scan of `src` to `dst` +(which may be `src` itself). **Other names**: prefix sum, `thrust::scan`, cumulative sum; inclusive (or exclusive) if the first element is included in the accumulation (or not). @@ -80,42 +51,25 @@ element is included in the accumulation (or not). The operator `op` must be associative, as elements are combined in parallel; it does not need to be commutative, as every combination keeps the elements in order (e.g. matrix products are fine). -For compatibility with the `Base.accumulate!` function, we provide the two-array interface too, but -we do not need the constraint of `dst` and `src` being different; to minimise memory use, we -recommend using the single-array interface (the first one above). - -## CPU -Use at most `max_tasks` threads with at least `min_elems` elements per task. - -Note that accumulation is typically a memory-bound operation, so multithreaded accumulation only -becomes faster if it is a more compute-heavy operation to hide memory latency - that includes: -- Accumulating more complex types, e.g. accumulation of tuples / structs / strings. -- More complex operators, e.g. `op=custom_complex_function`. - -## GPU -For the 1D case (`dims=nothing`), the `alg` can be one of the following: -- `ScanPrefixes()`: the default algorithm that scans the prefixes of each block, with no lookback; it - has better performance than `DecoupledLookback()` for large block sizes, and small to medium arrays, - but poorer scaling for many blocks; there is no performance degradation below `block_size^2` - elements, but it remains fast well into millions of elements. -- `DecoupledLookback()`: a more complex algorithm using opportunistic lookback to reuse earlier - blocks' results; requires device-level memory consistency guarantees (which Apple Metal does not - provide) and atomic orderings; theoretically more scalable for many blocks. - -A different, unique algorithm is used for the multi-dimensional case (`dims` is an integer). - -The `block_size` should be a power of 2 and greater than 0. `items_per_thread` controls how many -elements each thread processes per block and does not need to be a power of 2. Its default is at -most 8, reduced for wide element types to limit shared-memory use. - -The temporaries are only used for the 1D case (`dims=nothing`): `temp` stores per-block aggregates; -`temp_flags` is only used for the `DecoupledLookback()` algorithm for flagging if blocks are ready; -they should both have at least `cld(length(v), block_size * items_per_thread)` elements, using the -effective default described above when `items_per_thread` is omitted. -Also, `eltype(v) === eltype(temp)` is required; the elements in `temp_flags` can be any integers, -but `UInt8` is used by default to reduce memory usage. Multi-block exclusive scans with -`DecoupledLookback()` use two epilogue kernels to shift the result in place; they reuse `temp` for -tile-boundary values and do not allocate a full-array copy. +With `dims=nothing` the array is scanned in linear order, whatever its number of dimensions; with +an integer `dims`, each slice along that dimension is scanned independently. + +`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref +AcceleratedKernels.CPUThreads.Partitioned) on the host, and on GPUs [`ScanPrefixes`](@ref) for +whole arrays and [`SliceScan`](@ref) along `dims`, with the device's settings. +[`DecoupledLookback`](@ref) is available on backends that support it. `backend` is derived from +`dst` and `src`. + +On the host, accumulation is typically a memory-bound operation, so multithreaded accumulation +only becomes faster for more compute-heavy operations that hide memory latency, e.g. accumulating +tuples or structs, or expensive operators. + +The temporaries only apply to whole-array GPU scans: `temp` stores per-block aggregates, with +`eltype(temp) === eltype(dst)`; `temp_flags` stores `DecoupledLookback`'s block flags (any +integer type). Both need at least `cld(length(dst), block_size * items_per_thread)` elements +of the resolved algorithm. Multi-block exclusive scans with `DecoupledLookback` use two epilogue +kernels to shift the result in place; they reuse `temp` for tile-boundary values and do not +allocate a full-array copy. # Examples Example computing an inclusive prefix sum (the typical GPU "scan"): @@ -124,126 +78,66 @@ import AcceleratedKernels as AK using oneAPI v = oneAPI.ones(Int32, 100_000) -AK.accumulate!(+, v, init=0) +AK.accumulate!(+, v; init=0) -# Use a different algorithm -AK.accumulate!(+, v, alg=AK.DecoupledLookback()) +# Choose the algorithm and its settings +AK.accumulate!(+, v; init=0, alg=AK.ScanPrefixes(block_size=512)) ``` """ -function accumulate!( - op, v::AbstractArray, backend::Backend=get_backend(v); - init, - kwargs... -) - _accumulate_impl!( - op, v, backend; - init, - kwargs... - ) +function accumulate!(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) + _accumulate_impl!(op, v, v, _resolve_backend(backend, v); kwargs...) end - function accumulate!( - op, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(dst); - init, + op, dst::AbstractArray, src::AbstractArray; + backend::Union{Nothing, Backend}=nothing, kwargs... ) - copyto!(dst, src) - _accumulate_impl!( - op, dst, backend; - init, - kwargs... - ) + _accumulate_impl!(op, dst, src, _resolve_backend(backend, dst, src); kwargs...) end +# Scan `src` into `v` (which may be `src`), after checking the algorithm function _accumulate_impl!( - op, v::AbstractArray, backend::Backend; + op, v::AbstractArray, src::AbstractArray, backend::Backend; init, neutral=neutral_element(op, eltype(v)), - dims::Union{Nothing, Int}=nothing, + dims::Union{Nothing, Integer}=nothing, inclusive::Bool=true, - - alg::AccumulateAlgorithm=ScanPrefixes(), - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int=256, - items_per_thread::Union{Nothing, Int}=nothing, + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_flags::Union{Nothing, AbstractArray}=nothing, ) - if isnothing(dims) - return if use_gpu_algorithm(backend, prefer_threads) - items_per_thread = something( - items_per_thread, - default_scan_items_per_thread(backend, eltype(v), block_size), - ) - accumulate_1d_gpu!( - op, v, backend, alg; - init, neutral, inclusive, - max_tasks, min_elems, - block_size, items_per_thread, temp, temp_flags, - ) + dims isa Integer && dims < 1 && + throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) + a = _resolve_scan(alg, backend, eltype(v), dims) + if dims !== nothing && (temp !== nothing || temp_flags !== nothing) + throw(ArgumentError( + "`temp` and `temp_flags` only apply to whole-array scans (`dims=nothing`)")) + end + v === src || copyto!(v, src) + if dims === nothing + if a isa CPUThreads.Partitioned + accumulate_1d_cpu!(op, v, backend, a; init, neutral, inclusive) else - accumulate_1d_cpu!( - op, v, backend, alg; - init, neutral, inclusive, - max_tasks, min_elems, - block_size, temp, temp_flags, - ) + accumulate_1d_gpu!(op, v, backend, a; init, neutral, inclusive, temp, temp_flags) end else - return accumulate_nd!( - op, v, backend; - init, neutral, dims, inclusive, - max_tasks, min_elems, prefer_threads, - block_size, - ) + accumulate_nd!(op, v, backend, a; init, neutral, dims=Int(dims), inclusive) end + return v end """ - accumulate( - op, v::AbstractArray, backend::Backend=get_backend(v); - init, - neutral=neutral_element(op, eltype(v)), - dims::Union{Nothing, Int}=nothing, - inclusive::Bool=true, - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, - - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), + accumulate(op, v::AbstractArray; init, kwargs...) - # GPU settings - block_size::Int=256, - items_per_thread::Union{Nothing, Int}=nothing, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, - ) - -Out-of-place version of [`accumulate!`](@ref). +Out-of-place version of [`accumulate!`](@ref), with the same keywords; the result's element type +is `Base.promote_op(op, eltype(v), typeof(init))`. """ -function accumulate( - op, v::AbstractArray, backend::Backend=get_backend(v); - init, - kwargs... -) +function accumulate(op, v::AbstractArray; init, kwargs...) dst_type = Base.promote_op(op, eltype(v), typeof(init)) vcopy = similar(v, dst_type) copyto!(vcopy, v) - accumulate!( - op, vcopy, backend; - init, - kwargs... - ) - vcopy + accumulate!(op, vcopy; init, kwargs...) end diff --git a/src/accumulate/accumulate_1d_cpu.jl b/src/accumulate/accumulate_1d_cpu.jl index 8b64de72..0165313d 100644 --- a/src/accumulate/accumulate_1d_cpu.jl +++ b/src/accumulate/accumulate_1d_cpu.jl @@ -1,27 +1,18 @@ function accumulate_1d_cpu!( - op, v::AbstractArray, backend::Backend, alg; + op, v::AbstractArray, backend::Backend, alg::CPUThreads.Partitioned; init, neutral, inclusive::Bool, - - # CPU settings - max_tasks::Int, - min_elems::Int, - - # GPU settings - not used - block_size::Int, - temp::Union{Nothing, AbstractArray}, - temp_flags::Union{Nothing, AbstractArray}, ) # Trivial case if length(v) == 0 return v end - # Sanity checks - for exclusive accumulation, each task section / chunk must have at least 2 - # elements to be correct (otherwise we have to include more complicated logic in the threaded - # code); it makes no sense to have each task accumulate only 1 element anyways - @argcheck min_elems >= 2 + # For exclusive accumulation, each task section / chunk must have at least 2 elements to be + # correct (otherwise we have to include more complicated logic in the threaded code); the + # resolution checks `min_elems` + max_tasks, min_elems = alg.max_tasks, alg.min_elems # First accumulate chunks independently tp = TaskPartitioner(length(v), max_tasks, min_elems) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index fb9da460..469ed176 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -273,25 +273,14 @@ end # DecoupledLookback algorithm function accumulate_1d_gpu!( - op, v::AbstractArray, backend::Backend, ::DecoupledLookback; + op, v::AbstractArray, backend::Backend, alg::DecoupledLookback; init, neutral, inclusive::Bool, - - # CPU settings - not used - max_tasks::Int, - min_elems::Int, - - # GPU settings - block_size::Int, - items_per_thread::Int, temp::Union{Nothing, AbstractArray}, temp_flags::Union{Nothing, AbstractArray}, ) - # Correctness checks - @argcheck block_size > 0 - @argcheck ispow2(block_size) - @argcheck items_per_thread > 0 + block_size, items_per_thread = alg.block_size, alg.items_per_thread # Nothing to accumulate if length(v) == 0 @@ -344,25 +333,14 @@ end # ScanPrefixes algorithm function accumulate_1d_gpu!( - op, v::AbstractArray, backend, ::ScanPrefixes; + op, v::AbstractArray, backend, alg::ScanPrefixes; init, neutral, inclusive::Bool, - - # CPU settings - not used - max_tasks::Int, - min_elems::Int, - - # GPU settings - block_size::Int, - items_per_thread::Int, temp::Union{Nothing, AbstractArray}, temp_flags::Union{Nothing, AbstractArray}, ) - # Correctness checks - @argcheck block_size > 0 - @argcheck ispow2(block_size) - @argcheck items_per_thread > 0 + block_size, items_per_thread = alg.block_size, alg.items_per_thread # Nothing to accumulate if length(v) == 0 diff --git a/src/accumulate/accumulate_nd.jl b/src/accumulate/accumulate_nd.jl index 3dea501a..782a0dfa 100644 --- a/src/accumulate/accumulate_nd.jl +++ b/src/accumulate/accumulate_nd.jl @@ -1,29 +1,12 @@ function accumulate_nd!( - op, v::AbstractArray, backend::Backend; + op, v::AbstractArray, backend::Backend, alg::Union{SliceScan, CPUThreads.Partitioned}; init, neutral, dims::Int, inclusive::Bool, - - # CPU settings - max_tasks::Int, - min_elems::Int, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int, ) - # Correctness checks - @argcheck block_size > 0 - @argcheck ispow2(block_size) - # Degenerate cases begin; order of priority matters - # Invalid dims - if dims < 1 - throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) - end - # Nothing to accumulate vsizes = size(v) if length(v) == 0 || dims > length(vsizes) @@ -35,9 +18,11 @@ function accumulate_nd!( # Degenerate cases end - if !use_gpu_algorithm(backend, prefer_threads) - _accumulate_nd_cpu_sections!(op, v; init, dims, inclusive, max_tasks, min_elems) + if alg isa CPUThreads.Partitioned + _accumulate_nd_cpu_sections!(op, v; init, dims, inclusive, + max_tasks=alg.max_tasks, min_elems=alg.min_elems) else + block_size = alg.block_size # On GPUs we have two parallelisation approaches, based on which dimension has more elements: # - If the dimension we are accumulating along has more elements than the "outer" dimensions, # (e.g. accumulate(+, rand(3, 1000), dims=2)), we use a block of threads per outer diff --git a/src/accumulate/tuning.jl b/src/accumulate/tuning.jl new file mode 100644 index 00000000..266b40d2 --- /dev/null +++ b/src/accumulate/tuning.jl @@ -0,0 +1,163 @@ +# Selection and tuning of scan algorithms; the steps are those of sorting (src/sort/tuning.jl). + + +""" + ScanPrefixes(; block_size=nothing, items_per_thread=nothing) + +GPU scan of a whole array: each block scans a tile of `block_size * items_per_thread` elements +(`block_size` a power of two up to 1024), then the tiles' totals are scanned and added to later +tiles. No block waits for another, so it runs on every backend. The default `items_per_thread` is +at most 8, fewer for wide element types, to keep the tile in local memory. +""" +Base.@kwdef struct ScanPrefixes <: ScanAlgorithm + block_size::Union{Nothing, Int} = nothing + items_per_thread::Union{Nothing, Int} = nothing +end + +""" + DecoupledLookback(; block_size=nothing, items_per_thread=nothing) + +GPU scan of a whole array in which each block looks back at earlier blocks' published prefixes +instead of waiting for a separate pass over the tiles' totals; the tiles are those of +[`ScanPrefixes`](@ref). It needs device-scope memory ordering and forward progress between +blocks, so it is rejected on backends that do not guarantee them (currently all but CUDA and +AMDGPU). Elements keep their order, but how they are grouped depends on which blocks have +finished, so floating-point results can differ between runs. +""" +Base.@kwdef struct DecoupledLookback <: ScanAlgorithm + block_size::Union{Nothing, Int} = nothing + items_per_thread::Union{Nothing, Int} = nothing +end + +""" + SliceScan(; block_size=nothing) + +GPU scan of each slice along `dims`: one thread per slice when there are more slices than +elements per slice, otherwise one block (`block_size` a power of two up to 1024) per slice. +The algorithm for scans along `dims`. +""" +Base.@kwdef struct SliceScan <: ScanAlgorithm + block_size::Union{Nothing, Int} = nothing +end + + +""" + ScanTuning(; kwargs...) + +Values that drive `Auto` selection and fill unset algorithm fields for scans on one device, as +returned by [`scan_tuning`](@ref). + +- `prefer_lookback`: `Auto` picks `DecoupledLookback` for whole-array scans where the backend + supports it (`_supports_lookback`), else `ScanPrefixes`. +- `block_size`: the block size of every scan kernel algorithm. +- `local_mem_bytes`, `max_items`: the default `items_per_thread` of `ScanPrefixes` and + `DecoupledLookback` is the largest that keeps a tile of the element type within + `local_mem_bytes`, and at most `max_items`. It is derived from the effective `block_size`, so an + explicit `block_size` gets a matching default. +- `threads_min_elems`: the default `min_elems` of `CPUThreads.Partitioned`; at least 2. + +The defaults reproduce AK's historical settings. Internal: the fields may change in any release. +""" +Base.@kwdef struct ScanTuning + prefer_lookback::Bool = false + block_size::Int = 256 + local_mem_bytes::Int = LOCAL_MEMORY_BUDGET + max_items::Int = 8 + threads_min_elems::Int = 2 +end + +""" + scan_tuning(backend, T) -> ScanTuning + +The scan tuning for element type `T` (the destination's) on `backend`'s current device; see +[`sort_tuning`](@ref) for the conventions. +""" +scan_tuning(::Backend, ::Type) = ScanTuning() + + +""" + _resolve_scan(alg, backend, T, dims) -> Algorithm + +Resolve `alg` for a scan of element type `T` along `dims` (`nothing` for the whole array in +linear order) on `backend`. Returns `ScanPrefixes`, `DecoupledLookback`, `SliceScan` or +`CPUThreads.Partitioned` with every field set, or throws an `ArgumentError`. +""" +function _resolve_scan(alg::Algorithm, backend::Backend, ::Type{T}, dims) where {T} + _checkdomain(alg) + t = scan_tuning(backend, T) + a = alg isa Auto ? _select_scan(backend, t, dims) : alg + a = _fill(a, t, T) + _check_scan(a, backend, T, dims) + return a +end + +function _select_scan(backend, t::ScanTuning, dims) + _runs_threads(backend) && return CPUThreads.Partitioned() + dims === nothing || return SliceScan() + return _supports_lookback(backend) && t.prefer_lookback ? DecoupledLookback() : ScanPrefixes() +end + +function _checkdomain(a::Union{ScanPrefixes, DecoupledLookback, SliceScan}) + _check_pow2(a, :block_size) + a.block_size === nothing || a.block_size <= 1024 || throw(ArgumentError( + "$(_algname(a)): `block_size` must be at most 1024, got $(a.block_size)")) + a isa SliceScan || _check_positive(a, :items_per_thread) + nothing +end + +# The largest number of items per thread whose tile fits the tuning's local-memory budget. Types +# without a definite size and block sizes that `_check_scan` rejects get a placeholder. +function _scan_items(t::ScanTuning, block_size::Int, ::Type{T}) where {T} + 1 <= block_size <= 1024 || return 1 + clamp(t.local_mem_bytes ÷ (block_size * (isbitstype(T) ? max(sizeof(T), 1) : 1)) - 1, + 1, max(t.max_items, 1)) +end + +function _fill(a::A, t::ScanTuning, ::Type{T}) where {A <: Union{ScanPrefixes, DecoupledLookback}, T} + block_size = something(a.block_size, t.block_size) + items_per_thread = isnothing(a.items_per_thread) ? _scan_items(t, block_size, T) : + a.items_per_thread + return A(block_size, items_per_thread) +end +_fill(a::SliceScan, t::ScanTuning, T) = SliceScan(something(a.block_size, t.block_size)) +_fill(a::CPUThreads.Partitioned, t::ScanTuning, T) = _fill_threads(a, t) +_fill(a::Algorithm, t::ScanTuning, T) = + throw(ArgumentError("$(_algname(a)) is not a scan algorithm")) + +function _check_scan(a::Union{ScanPrefixes, DecoupledLookback}, backend, ::Type{T}, dims) where {T} + _checkdomain(a) + _require_kernels(a, backend) + _check_scan_eltype(a, T) + dims === nothing || throw(ArgumentError( + "$(_algname(a)) scans whole arrays; scan along `dims` with `SliceScan` or `Auto()`")) + if a isa DecoupledLookback && !_supports_lookback(backend) + throw(ArgumentError( + "DecoupledLookback needs device-scope memory ordering and forward progress between " * + "blocks, which $(_backend_name(backend)) does not guarantee; use `ScanPrefixes`")) + end + # Keep the kernels' index arithmetic far from overflow + widemul(a.block_size, a.items_per_thread) <= typemax(Int32) || throw(ArgumentError( + "$(_algname(a)): `block_size * items_per_thread` must be at most $(typemax(Int32))")) + nothing +end + +function _check_scan(a::SliceScan, backend, ::Type{T}, dims) where {T} + _checkdomain(a) + _require_kernels(a, backend) + _check_scan_eltype(a, T) + dims === nothing && throw(ArgumentError( + "SliceScan scans along `dims`; scan whole arrays with `ScanPrefixes` or `Auto()`")) + nothing +end + +function _check_scan(a::CPUThreads.Partitioned, backend, T, dims) + _checkdomain(a) + _require_threads(a, backend) + # Each task's part must hold two elements for the exclusive scan's carries + a.min_elems >= 2 || throw(ArgumentError( + "CPUThreads.Partitioned: scans need `min_elems` of at least 2, got $(a.min_elems)")) + nothing +end + +_check_scan_eltype(a, ::Type{T}) where {T} = isbitstype(T) || throw(ArgumentError( + "$(_algname(a)): the element type $T is not a bits type, which the kernels need")) diff --git a/src/algorithms.jl b/src/algorithms.jl index 03235598..7cb66763 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -52,6 +52,15 @@ Supertype of the reduction algorithms: [`BlockReduce`](@ref). Reductions also ac """ abstract type ReduceAlgorithm <: Algorithm end +""" + ScanAlgorithm <: Algorithm + +Supertype of the scan (`accumulate`) algorithms: [`ScanPrefixes`](@ref), +[`DecoupledLookback`](@ref) and [`SliceScan`](@ref). Scans also accept +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned). +""" +abstract type ScanAlgorithm <: Algorithm end + """ AcceleratedKernels.CPUThreads @@ -119,6 +128,16 @@ backend of KernelAbstractions 0.10, which compiles kernels for PoCL. KernelAbstr _runs_kernels(::Backend) = true _runs_kernels(::HostBackend) = nameof(HostBackend) !== :CPU +""" + _supports_lookback(backend) + +Whether [`DecoupledLookback`](@ref) is correct on `backend`: it needs a device-scope memory fence +(`_decoupled_fence`), atomic loads and stores of the block flags, and forward progress between +workgroups, since a block spins until an earlier one publishes its prefix. AK's extensions +declare it for the backends where all three are known to hold. +""" +_supports_lookback(::Backend) = false + # Checks shared by the algorithm families diff --git a/src/arithmetics.jl b/src/arithmetics.jl index b2f6c485..ccd6c26d 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -129,23 +129,10 @@ end """ - cumsum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - neutral=zero(eltype(src)), - dims::Union{Nothing, Int}=nothing, + cumsum(src::AbstractArray; init=zero(eltype(src)), neutral=zero(eltype(src)), kwargs...) - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, - ) - -Cumulative sum of elements of an array, with optional `init` and `dims`. Arguments are the same as -for [`accumulate`](@ref). +Cumulative sum of elements of an array, with optional `init` and `dims`. The keywords are those +of [`accumulate`](@ref). # Examples Simple cumulative sum of elements in a vector: @@ -163,39 +150,16 @@ m = ROCArray(rand(Int32(1):Int32(100), 10, 100_000)) s = AK.cumsum(m, dims=1) ``` """ -function cumsum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - neutral=zero(eltype(src)), - kwargs... -) - accumulate( - +, src, backend; - init, neutral, - inclusive=true, - kwargs... - ) +function cumsum(src::AbstractArray; init=zero(eltype(src)), neutral=zero(eltype(src)), kwargs...) + accumulate(+, src; init, neutral, inclusive=true, kwargs...) end """ - cumprod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - neutral=one(eltype(src)), - dims::Union{Nothing, Int}=nothing, - - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, - ) + cumprod(src::AbstractArray; init=one(eltype(src)), neutral=one(eltype(src)), kwargs...) -Cumulative product of elements of an array, with optional `init` and `dims`. Arguments are the same -as for [`accumulate`](@ref). +Cumulative product of elements of an array, with optional `init` and `dims`. The keywords are +those of [`accumulate`](@ref). # Examples Simple cumulative product of elements in a vector: @@ -213,16 +177,6 @@ m = oneArray(rand(Int32(1):Int32(100), 10, 100_000)) p = AK.cumprod(m, dims=1) ``` """ -function cumprod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - neutral=one(eltype(src)), - kwargs... -) - accumulate( - *, src, backend; - init, neutral, - inclusive=true, - kwargs... - ) +function cumprod(src::AbstractArray; init=one(eltype(src)), neutral=one(eltype(src)), kwargs...) + accumulate(*, src; init, neutral, inclusive=true, kwargs...) end diff --git a/src/findall.jl b/src/findall.jl index 486733d1..8f9805f4 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -115,7 +115,7 @@ function findall_gpu( kernel! = findall_block!(backend, block_size) kernel!(nothing, bools, block_counts, input_indices, output_indices, items; ndrange=num_blocks * block_size) - accumulate!(+, block_counts, backend; init=0) + accumulate!(+, block_counts; backend, init=0) n = @allowscalar block_counts[end] out = similar(bools, I, n) diff --git a/src/sort/cpu_sample_sort.jl b/src/sort/cpu_sample_sort.jl index 9a145364..8c3cb38d 100644 --- a/src/sort/cpu_sample_sort.jl +++ b/src/sort/cpu_sample_sort.jl @@ -25,7 +25,8 @@ function _sample_sort_compute_offsets!(histograms, max_tasks) offsets[j] += histograms[j, itask] end end - accumulate!(+, offsets, init=0, inclusive=false, max_tasks=1) + accumulate!(+, offsets; init=0, inclusive=false, + alg=CPUThreads.Partitioned(max_tasks=1)) # Compute each task's local offset into each bucket for itask in 1:max_tasks @@ -33,7 +34,7 @@ function _sample_sort_compute_offsets!(histograms, max_tasks) +, @view(histograms[itask, 1:max_tasks]), init=0, inclusive=false, - max_tasks=1, + alg=CPUThreads.Partitioned(max_tasks=1), ) end end diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index 46bb4844..a77e6ed1 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -455,7 +455,7 @@ function _radix_sort!( shift32 = UInt32(shift) hist_kern!(hist, p1, shift32, descending, vitems; ndrange) - accumulate!(+, hist, backend; init=UInt32(0), inclusive=false, temp=acc_temp) + accumulate!(+, hist; backend, init=UInt32(0), inclusive=false, temp=acc_temp) scat_kern!(p2, p1, hist, shift32, descending, vitems; ndrange) p1, p2 = p2, p1 diff --git a/src/utils.jl b/src/utils.jl index bd4c446c..bd43bace 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -8,12 +8,6 @@ const CPU_BACKEND = get_backend([]) return backend != CPU_BACKEND || !prefer_threads end -# Backends may request more than the historical default of two items per thread. -@inline default_items_per_thread(backend) = 1 - -# Elements each thread scans in the GPU prefix-scan block kernel. -@inline default_scan_items_per_thread(backend) = 8 - # Local memory that kernels with a tunable footprint may use: the minimum an OpenCL device must # provide (`CL_DEVICE_LOCAL_MEM_SIZE`, full profile), and Metal's threadgroup memory limit. const LOCAL_MEMORY_BUDGET = 32 * 1024 @@ -22,13 +16,6 @@ const LOCAL_MEMORY_BUDGET = 32 * 1024 # `init=nothing`, which is an explicit initial value. struct _NoInit end -# Keep the default shared-memory use within the local-memory budget. -@inline function default_scan_items_per_thread(backend, ::Type{T}, block_size) where T - block_size > 0 || return default_scan_items_per_thread(backend) - max_items = max(1, LOCAL_MEMORY_BUDGET ÷ (block_size * max(sizeof(T), 1)) - 1) - min(default_scan_items_per_thread(backend), max_items) -end - """ struct TypeWrap{T} end TypeWrap(T) = TypeWrap{T}() diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index 94adb396..cc372270 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -1,8 +1,33 @@ -ALGS = AK.AccumulateAlgorithm[AK.ScanPrefixes()] +# The whole-array scan algorithms under test: each kernel algorithm the backend supports, or the +# threaded host algorithm. `scan_alg` and `slice_alg` build a whole-array or `dims` algorithm for +# the configuration from both kinds of settings; without settings they give `Auto()`, except in +# the `--cpu-ka` configuration, whose point is to run AK's kernels on the host backend. +SCAN_ALGS = if !TEST_KERNELS + [AK.CPUThreads.Partitioned] +elseif TEST_DL + [AK.ScanPrefixes, AK.DecoupledLookback] +else + [AK.ScanPrefixes] +end +scan_name(A) = A === AK.CPUThreads.Partitioned ? "threads" : string(nameof(A)) + +function scan_alg(A=nothing; block_size=nothing, items_per_thread=nothing, + max_tasks=nothing, min_elems=nothing) + if A === nothing + Base.all(isnothing, (block_size, items_per_thread, max_tasks, min_elems)) && + return HOST_KERNELS ? AK.ScanPrefixes() : AK.Auto() + A = TEST_KERNELS ? AK.ScanPrefixes : AK.CPUThreads.Partitioned + end + A === AK.CPUThreads.Partitioned ? A(; max_tasks, min_elems) : A(; block_size, items_per_thread) +end -TEST_DL && push!(ALGS, AK.DecoupledLookback()) +function slice_alg(; block_size=nothing, max_tasks=nothing, min_elems=nothing) + Base.all(isnothing, (block_size, max_tasks, min_elems)) && + return HOST_KERNELS ? AK.SliceScan() : AK.Auto() + TEST_KERNELS ? AK.SliceScan(; block_size) : AK.CPUThreads.Partitioned(; max_tasks, min_elems) +end -@testset "accumulate_1d $(alg isa AK.DecoupledLookback ? "DL" : "SP")" for alg in ALGS +@testset "accumulate_1d $(scan_name(A))" for A in SCAN_ALGS Random.seed!(0) @@ -10,7 +35,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) for num_elems in 1:256 x = array_from_host(ones(Int32, num_elems)) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, inclusive=false, block_size=128, alg) + AK.accumulate!(+, y; init=0, inclusive=false, alg=scan_alg(A; block_size=128)) yh = Array(y) @test all(yh .== 0:length(yh) - 1) end @@ -19,7 +44,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) for num_elems in 1:256 x = array_from_host(rand(1:1000, num_elems), Int32) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, block_size=128, alg) + AK.accumulate!(+, y; init=0, alg=scan_alg(A; block_size=128)) @test all(Array(y) .== accumulate(+, Array(x))) end @@ -28,7 +53,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) num_elems = rand(1:100_000) x = array_from_host(ones(Int32, num_elems)) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, inclusive=false, alg) + AK.accumulate!(+, y; init=0, inclusive=false, alg=scan_alg(A)) yh = Array(y) @test all(yh .== 0:length(yh) - 1) end @@ -41,8 +66,8 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) init = rand(Int32(-100):Int32(100)) xh = rand(Int32(-9):Int32(9), num_elems) y = array_from_host(xh) - AK.accumulate!(+, y; prefer_threads, init, inclusive=false, block_size, - items_per_thread, alg) + AK.accumulate!(+, y; init, inclusive=false, + alg=scan_alg(A; block_size, items_per_thread)) @test Array(y) == (cumsum(xh) .- xh) .+ init end end @@ -51,7 +76,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) if KernelAbstractions.supports_float64(BACKEND) xh = ComplexF64.(1:4097) y = array_from_host(xh) - AK.accumulate!(+, y; prefer_threads, init=0.0 + 0.0im, alg) + AK.accumulate!(+, y; init=0.0 + 0.0im, alg=scan_alg(A)) @test Array(y) == cumsum(xh) end @@ -60,7 +85,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) num_elems = rand(1:100_000) x = array_from_host(rand(1:1000, num_elems), Int32) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, alg) + AK.accumulate!(+, y; init=0, alg=scan_alg(A)) @test all(Array(y) .== accumulate(+, Array(x))) end @@ -69,7 +94,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) num_elems = rand(1:100_000) x = array_from_host(rand(1:1000, num_elems), Int32) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, block_size=16, alg) + AK.accumulate!(+, y; init=0, alg=scan_alg(A; block_size=16)) @test all(Array(y) .== accumulate(+, Array(x))) end @@ -80,7 +105,7 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) n3 = rand(1:100) vh = rand(Float32, n1, n2, n3) v = array_from_host(vh) - AK.accumulate!(+, v; prefer_threads, init=0, alg) + AK.accumulate!(+, v; init=0, alg=scan_alg(A)) @test all(Array(v) .≈ accumulate(+, vh)) end @@ -90,37 +115,37 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) x = array_from_host(rand(1:1000, num_elems), Int32) y = similar(x) init = rand(-1000:1000) - AK.accumulate!(+, y, x; prefer_threads, init=Int32(init), alg) + AK.accumulate!(+, y, x; init=Int32(init), alg=scan_alg(A)) @test all(Array(y) .== accumulate(+, Array(x); init)) end # Exclusive scan x = array_from_host(ones(Int32, 10)) y = copy(x) - AK.accumulate!(+, y; prefer_threads, init=0, inclusive=false, alg) + AK.accumulate!(+, y; init=0, inclusive=false, alg=scan_alg(A)) @test all(Array(y) .== 0:9) # Test init value is respected with exclusive scan too x = array_from_host(ones(Int32, 10)) y = copy(x) init = 10 - AK.accumulate!(+, y; prefer_threads, init=Int32(init), inclusive=false, alg) + AK.accumulate!(+, y; init=Int32(init), inclusive=false, alg=scan_alg(A)) @test all(Array(y) .== 10:19) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.accumulate(+, y; prefer_threads, init=10, dims=2, inclusive=false, bad=:kwarg) + @test_throws MethodError AK.accumulate(+, y; init=10, dims=2, inclusive=false, bad=:kwarg, + alg=slice_alg()) # Oversized temporaries are allowed. y = array_from_host(ones(Int32, 1000)) - AK.accumulate!(+, y; init=0, inclusive=false, prefer_threads, block_size=128, alg, - temp=array_from_host(zeros(Int32, 1000)), - temp_flags=array_from_host(zeros(Int8, 1000))) + AK.accumulate!(+, y; init=0, inclusive=false, temp=array_from_host(zeros(Int32, 1000)), + temp_flags=array_from_host(zeros(Int8, 1000)), alg=scan_alg(A; block_size=128)) @test Array(y) == 0:999 y = AK.accumulate(+, array_from_host(ones(Int32, 1000)); init=0, inclusive=false, - prefer_threads, block_size=128, alg, temp=array_from_host(zeros(Int64, 1000)), - temp_flags=array_from_host(zeros(Int8, 1000))) + temp_flags=array_from_host(zeros(Int8, 1000)), + alg=scan_alg(A; block_size=128)) @test Array(y) == 0:999 # Cross-block coherence: small tiles (block_size 16-64, 1 item/thread) maximise the number of @@ -133,14 +158,14 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) xh = rand(Int32(-9):Int32(9), num_elems) yi = array_from_host(xh) - AK.accumulate!(+, yi; prefer_threads, init=Int32(0), inclusive=true, - block_size, items_per_thread=1, alg) + AK.accumulate!(+, yi; init=Int32(0), inclusive=true, + alg=scan_alg(A; block_size, items_per_thread=1)) @test Array(yi) == cumsum(xh) init = rand(Int32(-50):Int32(50)) ye = array_from_host(xh) - AK.accumulate!(+, ye; prefer_threads, init, inclusive=false, - block_size, items_per_thread=1, alg) + AK.accumulate!(+, ye; init, inclusive=false, + alg=scan_alg(A; block_size, items_per_thread=1)) @test Array(ye) == (cumsum(xh) .- xh) .+ init end end @@ -185,7 +210,7 @@ function scan_reference(xh, dims; init, inclusive) end -@testset "accumulate_1d non-commutative $(alg isa AK.DecoupledLookback ? "DL" : "SP")" for alg in ALGS +@testset "accumulate_1d non-commutative $(scan_name(A))" for A in SCAN_ALGS Random.seed!(0) # Single and multiple blocks, and more blocks than one block can scan (block_size=16), so that @@ -197,8 +222,8 @@ end xh = [scan_randmat() for _ in 1:n] init = scan_randmat() y = array_from_host(xh) - AK.accumulate!(scan_matmul, y; prefer_threads, max_tasks=4, init, neutral=SCAN_I2, - inclusive, block_size, items_per_thread, alg) + AK.accumulate!(scan_matmul, y; init, neutral=SCAN_I2, inclusive, + alg=scan_alg(A; max_tasks=4, block_size, items_per_thread)) @test Array(y) == scan_reference(xh, 1; init, inclusive) end end @@ -219,8 +244,8 @@ end xh = [scan_randmat() for _ in CartesianIndices(sz)] init = scan_randmat() y = array_from_host(xh) - AK.accumulate!(scan_matmul, y; prefer_threads, max_tasks=4, init, neutral=SCAN_I2, - inclusive, dims, block_size) + AK.accumulate!(scan_matmul, y; init, neutral=SCAN_I2, inclusive, dims, + alg=slice_alg(; max_tasks=4, block_size)) @test Array(y) == scan_reference(xh, dims; init, inclusive) end end @@ -237,7 +262,7 @@ end for ksize in 0:3 sh = rand(Int32(1):Int32(100), isize, jsize, ksize) s = array_from_host(sh) - d = AK.accumulate(+, s; prefer_threads, init=Int32(0), dims) + d = AK.accumulate(+, s; init=Int32(0), dims, alg=slice_alg()) dh = Array(d) dhres = accumulate(+, sh; init=Int32(0), dims) @@ -257,7 +282,7 @@ end vh = rand(Int32(1):Int32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.accumulate(+, v; prefer_threads, init=Int32(0), dims) + s = AK.accumulate(+, v; init=Int32(0), dims, alg=slice_alg()) sh = Array(s) @test sh == accumulate(+, vh; init=Int32(0), dims) end @@ -271,7 +296,7 @@ end vh = rand(UInt32(1):UInt32(100), n1, n2, n3) v = array_from_host(vh) - s = AK.accumulate(+, v; prefer_threads, init=UInt32(0), dims) + s = AK.accumulate(+, v; init=UInt32(0), dims, alg=slice_alg()) sh = Array(s) @test sh == accumulate(+, vh; init=UInt32(0), dims) end @@ -285,7 +310,7 @@ end vh = rand(Float32, n1, n2, n3) v = array_from_host(vh) - s = AK.accumulate(+, v; prefer_threads, init=Float32(0), dims) + s = AK.accumulate(+, v; init=Float32(0), dims, alg=slice_alg()) sh = Array(s) @test all(sh .≈ accumulate(+, vh; init=Float32(0), dims)) end @@ -300,7 +325,7 @@ end vh = rand(Float32, n1, n2, n3) v = array_from_host(vh) init = rand(-1000:1000) - s = AK.accumulate(+, v; prefer_threads, init=Float32(init), dims) + s = AK.accumulate(+, v; init=Float32(init), dims, alg=slice_alg()) sh = Array(s) @test all(sh .≈ accumulate(+, vh; init=Float32(init), dims)) end @@ -309,19 +334,20 @@ end # Exclusive scan vh = ones(Int32, 10, 10) v = array_from_host(vh) - s = AK.accumulate(+, v; prefer_threads, init=0, dims=2, inclusive=false) + s = AK.accumulate(+, v; init=0, dims=2, inclusive=false, alg=slice_alg()) sh = Array(s) @test all([sh[i, :] == 0:9 for i in 1:10]) # Test init value is respected with exclusive scan too vh = ones(Int32, 10, 10) v = array_from_host(vh) - s = AK.accumulate(+, v; prefer_threads, init=10, dims=2, inclusive=false) + s = AK.accumulate(+, v; init=10, dims=2, inclusive=false, alg=slice_alg()) sh = Array(s) @test all([sh[i, :] == 10:19 for i in 1:10]) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.accumulate(+, v; prefer_threads, init=10, dims=2, inclusive=false, bad=:kwarg) + @test_throws MethodError AK.accumulate(+, v; init=10, dims=2, inclusive=false, bad=:kwarg, + alg=slice_alg()) # Test all options with bigger matrices for D in [(1_000_000,3), (3,1_000_000)], dims in [1,2] @@ -335,26 +361,13 @@ end end # Testing different settings - AK.accumulate( - (x, y) -> x + 1, - array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, - init=Int32(0), - neutral=Int32(0), - dims=2, - block_size=64, - temp=array_from_host(zeros(Int32, 3, 1, 5)), - ) - AK.accumulate( - (x, y) -> x + 1, - array_from_host(rand(Int32, 3, 4, 5)); - prefer_threads, - init=Int32(0), - neutral=Int32(0), - dims=3, - block_size=64, - temp=array_from_host(zeros(Int32, 3, 4, 1)), - ) + AK.accumulate((x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); init=Int32(0), + neutral=Int32(0), dims=2, alg=slice_alg(; block_size=64)) + AK.accumulate((x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); init=Int32(0), + neutral=Int32(0), dims=3, alg=slice_alg(; block_size=64)) + # The temporaries only apply to whole-array scans + @test_throws ArgumentError AK.accumulate(+, array_from_host(rand(Int32, 3, 4)); init=Int32(0), + dims=2, temp=array_from_host(zeros(Int32, 3))) end @testset "cumsum" begin @@ -363,14 +376,14 @@ end # Simple correctness tests v = array_from_host(1:100) vh = Array(v) - @test Array(AK.cumsum(v; prefer_threads)) == cumsum(vh) + @test Array(AK.cumsum(v; alg=scan_alg())) == cumsum(vh) # Fuzzy testing for _ in 1:100 num_elems = rand(1:100_000) vh = rand(Float32, num_elems) v = array_from_host(vh) - @test all(Array(AK.cumsum(v; prefer_threads)) .≈ cumsum(vh)) + @test all(Array(AK.cumsum(v; alg=scan_alg())) .≈ cumsum(vh)) end for _ in 1:100 @@ -382,10 +395,10 @@ end v = array_from_host(vh) # Indexing into array as if linear; not supported in Base - # @test all(Array(AK.cumsum(v; prefer_threads)) .== cumsum(vh)) + # @test all(Array(AK.cumsum(v; alg=scan_alg())) .== cumsum(vh)) # Along dimensions - r = Array(AK.cumsum(v; prefer_threads, dims)) + r = Array(AK.cumsum(v; dims, alg=slice_alg())) rh = cumsum(vh; dims) @test r == rh @@ -395,14 +408,14 @@ end # Test promotion to op-dictated type xh = rand(Bool, 16) x = array_from_host(xh) - @test Array(AK.cumsum(x; prefer_threads)) == cumsum(xh) + @test Array(AK.cumsum(x; alg=scan_alg())) == cumsum(xh) # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.cumsum(v; prefer_threads, block_size=64) + AK.cumsum(v; alg=scan_alg(; block_size=64)) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.cumsum(v; prefer_threads, init=10, bad=:kwarg) + @test_throws MethodError AK.cumsum(v; init=10, bad=:kwarg, alg=scan_alg()) # The other settings are stress-tested in reduce end @@ -415,11 +428,11 @@ end # Simple correctness tests v = array_from_host(1:100) vh = Array(v) - @test Array(AK.cumprod(v; prefer_threads)) == cumprod(vh) + @test Array(AK.cumprod(v; alg=scan_alg())) == cumprod(vh) vh = ones(Float32, 100_000) v = array_from_host(vh) - @test Array(AK.cumprod(v; prefer_threads)) == vh + @test Array(AK.cumprod(v; alg=scan_alg())) == vh # Fuzzy testing for _ in 1:100 @@ -431,10 +444,10 @@ end v = array_from_host(vh) # Indexing into array as if linear; not supported in Base - # @test all(Array(AK.cumprod(v; prefer_threads)) .== cumprod(vh)) + # @test all(Array(AK.cumprod(v; alg=scan_alg())) .== cumprod(vh)) # Along dimensions - r = Array(AK.cumprod(v; prefer_threads, dims)) + r = Array(AK.cumprod(v; dims, alg=slice_alg())) rh = cumprod(vh; dims) @test r == rh @@ -443,10 +456,10 @@ end # Testing different settings v = array_from_host(rand(-5:5, 100_000)) - AK.cumprod(v; prefer_threads, block_size=64) + AK.cumprod(v; alg=scan_alg(; block_size=64)) # Test that undefined kwargs are not accepted - @test_throws MethodError AK.cumprod(v; prefer_threads, init=10, bad=:kwarg) + @test_throws MethodError AK.cumprod(v; init=10, bad=:kwarg, alg=scan_alg()) # The other settings are stress-tested in reduce end diff --git a/test/generic/scan_resolve.jl b/test/generic/scan_resolve.jl new file mode 100644 index 00000000..45dc5363 --- /dev/null +++ b/test/generic/scan_resolve.jl @@ -0,0 +1,145 @@ +# Resolution of scan algorithms, without launching kernels. + +# A GPU backend with a tuning and a lookback capability the tests can change +struct ScanResolveTestBackend <: KernelAbstractions.Backend end +const SCAN_RESOLVE_TUNING = Ref(AK.ScanTuning()) +const SCAN_RESOLVE_LOOKBACK = Ref(false) +AK.scan_tuning(::ScanResolveTestBackend, ::Type) = SCAN_RESOLVE_TUNING[] +AK._supports_lookback(::ScanResolveTestBackend) = SCAN_RESOLVE_LOOKBACK[] + +# A backend that cannot run AK's kernels, like KernelAbstractions 0.9's `CPU` +struct ScanNoKernelsTestBackend <: KernelAbstractions.Backend end +AK._runs_kernels(::ScanNoKernelsTestBackend) = false + +const SRB = ScanResolveTestBackend() + +resolve_scan(alg; T=Float32, dims=nothing, backend=SRB) = AK._resolve_scan(alg, backend, T, dims) + +function with_scan_tuning(f; lookback=false, kwargs...) + old = SCAN_RESOLVE_TUNING[], SCAN_RESOLVE_LOOKBACK[] + SCAN_RESOLVE_TUNING[] = AK.ScanTuning(; kwargs...) + SCAN_RESOLVE_LOOKBACK[] = lookback + try + f() + finally + SCAN_RESOLVE_TUNING[], SCAN_RESOLVE_LOOKBACK[] = old + end +end + + +@testset "scan resolution: Auto" begin + # The default tuning reproduces AK's historical settings: 256 threads, at most 8 items per + # thread, fewer for wide element types + @test resolve_scan(AK.Auto()) === AK.ScanPrefixes(256, 8) + @test resolve_scan(AK.Auto(); T=ComplexF64) === AK.ScanPrefixes(256, 7) + @test resolve_scan(AK.Auto(); T=NTuple{16, Float64}) === AK.ScanPrefixes(256, 1) + @test resolve_scan(AK.Auto(); dims=2) === AK.SliceScan(256) + + # DecoupledLookback only where the tuning prefers it and the backend supports it + with_scan_tuning(; prefer_lookback=true) do + @test resolve_scan(AK.Auto()) === AK.ScanPrefixes(256, 8) + end + with_scan_tuning(; lookback=true) do + @test resolve_scan(AK.Auto()) === AK.ScanPrefixes(256, 8) + end + with_scan_tuning(; prefer_lookback=true, lookback=true) do + @test resolve_scan(AK.Auto()) === AK.DecoupledLookback(256, 8) + @test resolve_scan(AK.Auto(); dims=1) === AK.SliceScan(256) + end + + # On the host backend: the threaded algorithm, filled + host = AK.HOST_BACKEND + @test resolve_scan(AK.Auto(); backend=host) === AK.CPUThreads.Partitioned(Threads.nthreads(), 2) + @test resolve_scan(AK.Auto(); backend=host, dims=1) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 2) +end + + +@testset "scan resolution: explicit algorithms" begin + # Explicit fields win over the tuning; `items_per_thread` is derived from the effective + # `block_size`, so an explicit block size gets a matching default + with_scan_tuning(; block_size=128, max_items=16) do + @test resolve_scan(AK.ScanPrefixes()) === AK.ScanPrefixes(128, 16) + @test resolve_scan(AK.ScanPrefixes(block_size=1024)) === AK.ScanPrefixes(1024, 7) + @test resolve_scan(AK.ScanPrefixes(block_size=1024, items_per_thread=12)) === + AK.ScanPrefixes(1024, 12) + @test resolve_scan(AK.SliceScan(); dims=1) === AK.SliceScan(128) + end + with_scan_tuning(; local_mem_bytes=1024) do + @test resolve_scan(AK.ScanPrefixes(block_size=64)) === AK.ScanPrefixes(64, 3) + @test resolve_scan(AK.ScanPrefixes(block_size=1024)) === AK.ScanPrefixes(1024, 1) + end + host = AK.HOST_BACKEND + @test resolve_scan(AK.CPUThreads.Partitioned(max_tasks=3); backend=host) === + AK.CPUThreads.Partitioned(3, 2) + + # Domains of explicit fields + for bad in (AK.ScanPrefixes(block_size=0), AK.ScanPrefixes(block_size=96), + AK.ScanPrefixes(block_size=2048), AK.ScanPrefixes(items_per_thread=0), + AK.SliceScan(block_size=3)) + @test_throws ArgumentError resolve_scan(bad; dims=bad isa AK.SliceScan ? 1 : nothing) + end + @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(block_size=1024, items_per_thread=1 << 22)) + # A tuning's settings are checked like explicit ones + for block_size in (0, 1 << 62) + with_scan_tuning(; block_size) do + @test_throws ArgumentError resolve_scan(AK.Auto()) + end + end + # Exclusive scans on threads need two elements per task + @test_throws ArgumentError resolve_scan(AK.CPUThreads.Partitioned(min_elems=1); backend=host) + + # Whole-array algorithms do not scan along `dims`, and SliceScan only does; they used to be + # ignored silently along `dims` + @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(); dims=1) + @test_throws ArgumentError resolve_scan(AK.SliceScan()) + + # DecoupledLookback needs the capability, whatever the tuning says + @test_throws ArgumentError resolve_scan(AK.DecoupledLookback()) + with_scan_tuning(; lookback=true) do + @test resolve_scan(AK.DecoupledLookback()) === AK.DecoupledLookback(256, 8) + end + + # Kernels need a bits-type element + @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(); T=String) + @test_throws ArgumentError resolve_scan(AK.Auto(); T=Union{Missing, Int}, dims=1) + + # Capabilities: the threaded algorithm only on the host, kernels only where they run + @test_throws ArgumentError resolve_scan(AK.CPUThreads.Partitioned()) + @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(); backend=ScanNoKernelsTestBackend()) + + # Algorithms of other families + @test_throws ArgumentError resolve_scan(AK.BlockReduce()) + @test_throws ArgumentError resolve_scan(AK.CPUThreads.SampleSort(); backend=host) + @test_throws ArgumentError AK._resolve_reduce(AK.ScanPrefixes(), SRB, Float32, :) +end + + +@testset "scan resolution: inference" begin + host = typeof(AK.HOST_BACKEND) + for (B, dims) in ((ScanResolveTestBackend, Nothing), (ScanResolveTestBackend, Int), + (host, Nothing), (host, Int)) + rt = only(Base.return_types(AK._resolve_scan, (AK.Auto, B, Type{Float32}, dims))) + @test rt <: Union{AK.ScanPrefixes, AK.DecoupledLookback, AK.SliceScan, + AK.CPUThreads.Partitioned} + end +end + + +@testset "scan resolution: entry points" begin + # The keyword backend: derived from the arrays, destination first + v = rand(Int32(1):Int32(9), 1000) + @test AK.accumulate!(+, copy(v); init=Int32(0)) == cumsum(v) + @test AK.accumulate!(+, similar(v), v; init=Int32(0)) == cumsum(v) + @test AK.accumulate(+, v; init=Int32(0), backend=AK.HOST_BACKEND) == cumsum(v) + # An explicit backend is used as given, and checked against the algorithm + @test_throws ArgumentError AK.accumulate!(+, similar(v), v; init=Int32(0), backend=SRB, + alg=AK.CPUThreads.Partitioned()) + # A rejected algorithm leaves the destination untouched + dst = zeros(Int32, 1000) + @test_throws ArgumentError AK.accumulate!(+, dst, v; init=Int32(0), + alg=AK.CPUThreads.Partitioned(min_elems=1)) + @test Base.all(iszero, dst) + # Invalid `dims` + @test_throws ArgumentError AK.accumulate(+, reshape(v, 10, 100); init=Int32(0), dims=0) +end From 473d42c160e8a419b800fa86995d190463a025e6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 07:44:01 +0200 Subject: [PATCH 6/8] Define the scan contract 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. --- docs/src/api/differences.md | 4 + src/accumulate/accumulate.jl | 196 +++++++++++++++++++++------- src/accumulate/accumulate_1d_cpu.jl | 48 ++++--- src/accumulate/accumulate_1d_gpu.jl | 94 +++++++------ src/accumulate/accumulate_nd.jl | 90 +++++++------ src/accumulate/tuning.jl | 20 +-- src/arithmetics.jl | 22 ++-- src/reduce/utilities.jl | 18 +++ test/Project.toml | 1 + test/generic/accumulate.jl | 177 ++++++++++++++++++++++++- test/generic/scan_resolve.jl | 7 + 11 files changed, 504 insertions(+), 173 deletions(-) diff --git a/docs/src/api/differences.md b/docs/src/api/differences.md index a6938666..61fbc25e 100644 --- a/docs/src/api/differences.md +++ b/docs/src/api/differences.md @@ -17,3 +17,7 @@ GPU arrays on top of AcceleratedKernels and reproduces Base's results. The diffe | an `init` that is not a neutral element of `op` | outside the contract: `init` must be neutral, and it is unspecified whether it is used for non-empty collections | any value, applied exactly once | | `reduce(op, A; dims)` for an `op` without `Base.reducedim_init`, such as a closure | `MethodError` | works | | a non-commutative `op` in a reduction | works (elements keep their order) | unsupported | +| the running-value type of `accumulate!(op, B, A)` | depends on the code path: the fold type from `init` or the elements for vectors and `dims=1`, `eltype(B)` along `dims ≥ 2` | the fold type from `eltype(B)` joined with `init`'s type and the elements (unless `acctype` is given), so for the usual operators not narrower than `eltype(B)` | +| the element type of `accumulate(op, A)` on Julia 1.10 | `Base.promote_op(op, T, T)` | the fold type, as Base's on Julia 1.13 | +| `accumulate(op, A; dims, init)` with `dims > ndims(A)` | copies `A`, ignoring `init` | applies `init` to every element | +| a non-associative `op` in a scan, such as `-` | works (a sequential recurrence) | unsupported | diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index 0e7ba211..63c29fe9 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -1,23 +1,6 @@ include("tuning.jl") -# Helpers -# Check the given dst is compatible with src and init -function _accumulate_check_types(dst, src, init) - eltype(dst) === eltype(src) && return - eltype(dst) === typeof(init) && return - eltype(dst) === promote_type(eltype(src), typeof(init)) && return - - throw(ArgumentError( - """ - destination array type `$(eltype(dst))` (temp) is incompatible with source array type - `$(eltype(src))` and initial value type `$(typeof(init))`; eltype(dst) must be either - like eltype(src) or typeof(init) or promote_type(eltype(src), typeof(init)). - """ - )) -end - - # Implementations, then interfaces include("accumulate_1d_cpu.jl") include("accumulate_1d_gpu.jl") @@ -31,10 +14,11 @@ include("accumulate_nd.jl") accumulate!( op, dst, src; backend=nothing, - init, - neutral=neutral_element(op, eltype(dst)), + init=, + neutral=nothing, dims::Union{Nothing, Integer}=nothing, inclusive::Bool=true, + acctype=nothing, alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_flags::Union{Nothing, AbstractArray}=nothing, @@ -42,17 +26,39 @@ include("accumulate_nd.jl") Compute accumulated running totals along a sequence by applying a binary operator to all elements up to the current one; often used in GPU programming as a first step in finding / extracting -subsets of data. The first form scans `v` in place; the second writes the scan of `src` to `dst` -(which may be `src` itself). +subsets of data. The first form scans `v` in place; the second writes the scan of `src` to `dst`, +which must be `src` itself or not overlap it. **Other names**: prefix sum, `thrust::scan`, cumulative sum; inclusive (or exclusive) if the first element is included in the accumulation (or not). -The operator `op` must be associative, as elements are combined in parallel; it does not need to -be commutative, as every combination keeps the elements in order (e.g. matrix products are fine). - With `dims=nothing` the array is scanned in linear order, whatever its number of dimensions; with -an integer `dims`, each slice along that dimension is scanned independently. +an integer `dims`, each slice along that dimension is scanned independently. A `dims` beyond the +array's dimensions makes every slice one element long. + +The contract, which is AcceleratedKernels' own (see [Differences from Base](@ref)): + +- **Algebra.** `op` must be associative, as elements are combined in parallel; it does not need + to be commutative, as every combination keeps the elements in order (matrix products are + fine). Non-associative operators such as `-` are not supported. +- **Inclusive scans** (the default): with running values `y`, `y[1] = Base.reduce_first(op, + src[1])` without `init`, and `op(init, src[1])` with it; then `y[k] = op(y[k-1], src[k])`, and + `dst[k] = y[k]`. `init` is applied once per slice. +- **Exclusive scans** (`inclusive=false`): `y[1]` is `init`, or the neutral element of `op` + without it (an `ArgumentError` when none is known); then `y[k] = op(y[k-1], src[k-1])`. +- **Neutral element.** `neutral` (a two-sided identity of `op`) seeds partial results; it + defaults to `GPUArraysCore.neutral_element(op, T)` where that is defined. For other operators, + partial results start from their first element instead, so inclusive scans need no neutral + element. +- **Running-value type.** The running values have one type, and are converted to `eltype(dst)` + only when stored: `acctype` when it is given, else the type the fold of `op` settles on, + starting from `eltype(dst)` (joined with `init`'s type) and the elements, as for + [`mapreducedim!`](@ref) (e.g. `Int` for `add_sum` over `Int8`). So with the usual operators and + no `acctype`, a scan does not run in a type narrower than its destination: `Float32`s scanned into a `Float64` + array are summed in `Float64`. Where the running type differs from `eltype(dst)`, the scan runs in a scratch array + of that type. An `acctype` that cannot hold the running values at all is an `ArgumentError`; + whether their values fit is the caller's obligation. [`accumulate`](@ref) allocates the fold + type from `init`'s type and the elements. `alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host, and on GPUs [`ScanPrefixes`](@ref) for @@ -64,12 +70,12 @@ On the host, accumulation is typically a memory-bound operation, so multithreade only becomes faster for more compute-heavy operations that hide memory latency, e.g. accumulating tuples or structs, or expensive operators. -The temporaries only apply to whole-array GPU scans: `temp` stores per-block aggregates, with -`eltype(temp) === eltype(dst)`; `temp_flags` stores `DecoupledLookback`'s block flags (any -integer type). Both need at least `cld(length(dst), block_size * items_per_thread)` elements -of the resolved algorithm. Multi-block exclusive scans with `DecoupledLookback` use two epilogue -kernels to shift the result in place; they reuse `temp` for tile-boundary values and do not -allocate a full-array copy. +The temporaries only apply to whole-array GPU scans with a known neutral element: `temp` stores +per-block aggregates, with the accumulator element type (see **Types** above); `temp_flags` stores +`DecoupledLookback`'s block flags (any integer type). Both need at least +`cld(length(dst), block_size * items_per_thread)` elements of the resolved algorithm. +Multi-block exclusive scans with `DecoupledLookback` use two epilogue kernels to shift the result +in place; they reuse `temp` for tile-boundary values and do not allocate a full-array copy. # Examples Example computing an inclusive prefix sum (the typical GPU "scan"): @@ -78,10 +84,11 @@ import AcceleratedKernels as AK using oneAPI v = oneAPI.ones(Int32, 100_000) -AK.accumulate!(+, v; init=0) +AK.accumulate!(+, v) # 1, 2, 3, ... +AK.accumulate!(+, v; inclusive=false) # 0, 1, 3, ...: starts from the neutral element # Choose the algorithm and its settings -AK.accumulate!(+, v; init=0, alg=AK.ScanPrefixes(block_size=512)) +AK.accumulate!(+, v; alg=AK.ScanPrefixes(block_size=512)) ``` """ function accumulate!(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) @@ -93,6 +100,13 @@ function accumulate!( backend::Union{Nothing, Backend}=nothing, kwargs... ) + if dst !== src + Base.mightalias(dst, src) && throw(ArgumentError( + "the destination of a scan must be its source or not overlap it")) + axes(dst) == axes(src) || throw(DimensionMismatch( + "the destination of a scan must have the source's axes $(axes(src)), " * + "got $(axes(dst))")) + end _accumulate_impl!(op, dst, src, _resolve_backend(backend, dst, src); kwargs...) end @@ -100,44 +114,130 @@ end # Scan `src` into `v` (which may be `src`), after checking the algorithm function _accumulate_impl!( op, v::AbstractArray, src::AbstractArray, backend::Backend; - init, - neutral=neutral_element(op, eltype(v)), + init=_NoInit(), + neutral=nothing, dims::Union{Nothing, Integer}=nothing, inclusive::Bool=true, + acctype=nothing, alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_flags::Union{Nothing, AbstractArray}=nothing, ) dims isa Integer && dims < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) - a = _resolve_scan(alg, backend, eltype(v), dims) + A = _scan_acctype(op, eltype(v), eltype(src), init, inclusive, + _scan_combines(size(src), dims, init, inclusive), acctype) + neutral = _reduce_seed(op, A, neutral) + # The kernels' tiles hold partial results: lanes, when `op` has no known neutral element + a = _resolve_scan(alg, backend, A, dims, typeof(neutral)) if dims !== nothing && (temp !== nothing || temp_flags !== nothing) throw(ArgumentError( "`temp` and `temp_flags` only apply to whole-array scans (`dims=nothing`)")) end - v === src || copyto!(v, src) + if !inclusive && init isa _NoInit && neutral isa _Lane && !_valid(neutral) + throw(ArgumentError( + "an exclusive scan without `init` starts from the neutral element of `op`, " * + "which is not known for $op; pass `init` or `neutral`")) + end + + # As Base, scan in the accumulator type and convert to `v`'s element type only when storing: + # in `v` when it has that type, else in a scratch array. Elements enter as one-element + # reductions (`Base.reduce_first`), and an inclusive scan applies `init` to the first element + # of each slice, as `op(init, x)`. + w = eltype(v) === A ? v : similar(v, A) + launch = _scan_launch(a) + if w !== src + _foreachindex(eachindex(w, src), backend; launch...) do i + @inbounds w[i] = Base.reduce_first(op, src[i]) + end + end + # (`init` of the accumulator type seeds the kernels directly instead, except where every + # slice has one element and no kernel runs) + if inclusive && !(init isa _NoInit) && !isempty(w) && + (!(init isa A) || dims !== nothing && dims > ndims(w)) + _scan_apply_init!(w, op, init, dims, backend, launch) + init = _NoInit() + end + if dims === nothing if a isa CPUThreads.Partitioned - accumulate_1d_cpu!(op, v, backend, a; init, neutral, inclusive) + accumulate_1d_cpu!(op, w, backend, a; init, neutral, inclusive) else - accumulate_1d_gpu!(op, v, backend, a; init, neutral, inclusive, temp, temp_flags) + accumulate_1d_gpu!(op, w, backend, a; init, neutral, inclusive, temp, temp_flags) end + elseif dims > ndims(w) + # Every slice has one element: inclusive scans applied `init` above, exclusive ones start + # over + inclusive || fill!(w, _unlane(_scan_first_seed(w, init, neutral))) else - accumulate_nd!(op, v, backend, a; init, neutral, dims=Int(dims), inclusive) + accumulate_nd!(op, w, backend, a; init, neutral, dims=Int(dims), inclusive) end + w === v || copyto!(v, w) return v end +# The running-value type of a scan of elements of type `T` into an array of element type `D` +# (`Union{}` for none): `acctype`, else the accumulator type of a reduction from `D` joined with +# `init`'s type; an exclusive scan's `init` is a running value too. `combines` says whether `op` +# is called at all. +function _scan_acctype(op, ::Type{D}, ::Type{T}, init, inclusive::Bool, combines::Bool, + acctype=nothing) where {D, T} + acctype === nothing || return _acctype(op, Union{}, T, acctype) + S = init isa _NoInit ? D : promote_type(D, typeof(init)) + A = _reduce_acctype(op, S, T) + if !inclusive && !(init isa _NoInit) && A !== Union{} + A = promote_type(A, typeof(init)) + end + if A === Union{} + # `op` always throws for these types, which only a scan that never calls it gets past + combines && _check_acctype(op, identity, A) + F = _first_type(op, T) + return F === Union{} ? T : F + end + return A +end + +# Whether a scan of an array of size `sz` along `dims` calls `op`: some slice has two elements, +# or an inclusive scan applies `init` to a first element (along a `dims` beyond the array's +# dimensions, every element is one) +function _scan_combines(sz, dims, init, inclusive) + n = Base.prod(sz; init=1) + (n == 0 || dims isa Integer && dims < 1) && return false + len = dims === nothing ? n : dims <= length(sz) ? sz[dims] : 1 + return len >= 2 || (inclusive && !(init isa _NoInit)) +end + +_scan_launch(a::Union{ScanPrefixes, DecoupledLookback, SliceScan}) = (; block_size=a.block_size) +_scan_launch(a::CPUThreads.Partitioned) = (; max_tasks=a.max_tasks, min_elems=a.min_elems) + +# `w[i] = op(init, w[i])` for the first element `i` of the array (`dims === nothing`) or of each +# slice along `dims` (every element, for a `dims` beyond the array's) +function _scan_apply_init!(w, op, init, dims, backend, launch) + ax = axes(w) + firsts = dims === nothing ? CartesianIndices(Base.map(a -> first(a):first(a), ax)) : + dims > length(ax) ? CartesianIndices(ax) : + CartesianIndices(Base.setindex(ax, first(ax[dims]):first(ax[dims]), dims)) + _foreachindex(firsts, backend; launch...) do I + @inbounds w[I] = op(init, w[I]) + end + return w +end + """ - accumulate(op, v::AbstractArray; init, kwargs...) + accumulate(op, v::AbstractArray; init=, kwargs...) -Out-of-place version of [`accumulate!`](@ref), with the same keywords; the result's element type -is `Base.promote_op(op, eltype(v), typeof(init))`. +Out-of-place version of [`accumulate!`](@ref), with the same keywords. The result's element type +is `acctype` when it is given, else the type the fold of `op` settles on from `init`'s type (when +given) and the elements, e.g. an `Int` array for `accumulate(+, Int8[1, 2]; init=0)` and for +`accumulate(Base.add_sum, Int8[1, 2])`. """ -function accumulate(op, v::AbstractArray; init, kwargs...) - dst_type = Base.promote_op(op, eltype(v), typeof(init)) - vcopy = similar(v, dst_type) - copyto!(vcopy, v) - accumulate!(op, vcopy; init, kwargs...) +function accumulate(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + init=_NoInit(), dims=nothing, inclusive::Bool=true, acctype=nothing, + kwargs...) + backend = _resolve_backend(backend, v) + combines = _scan_combines(size(v), dims, init, inclusive) + D = _scan_acctype(op, Union{}, eltype(v), init, inclusive, combines, acctype) + accumulate!(op, _similar(backend, v, D), v; backend, init, dims, inclusive, acctype, + kwargs...) end diff --git a/src/accumulate/accumulate_1d_cpu.jl b/src/accumulate/accumulate_1d_cpu.jl index 0165313d..fb73f2b5 100644 --- a/src/accumulate/accumulate_1d_cpu.jl +++ b/src/accumulate/accumulate_1d_cpu.jl @@ -14,34 +14,45 @@ function accumulate_1d_cpu!( # resolution checks `min_elems` max_tasks, min_elems = alg.max_tasks, alg.min_elems + seed = _scan_first_seed(v, init, neutral) + _, op = _lanefuncs(_Partials(), op, neutral) + # First accumulate chunks independently tp = TaskPartitioner(length(v), max_tasks, min_elems) if tp.num_tasks == 1 - _accumulate_1d_cpu_section!(op, v; init, inclusive) + _accumulate_1d_cpu_section!(op, v; seed, neutral, inclusive) return v end # Scan each task's section with the requested inclusivity, seeding only the first one with - # `init`, and save each section's total - shared = Vector{eltype(v)}(undef, tp.num_tasks) + # `seed`, and save each section's total + shared = Vector{typeof(neutral)}(undef, tp.num_tasks) itask_partition(tp) do itask, irange shared[itask] = _accumulate_1d_cpu_section!( op, @view(v[irange]); - init=itask == 1 ? init : neutral, + seed=itask == 1 ? seed : neutral, + neutral, inclusive, ) end # Now accumulate the totals of each task; the number of tasks is small enough (even for # 144-thread HPC nodes) that there is no need to do decoupled lookbacks - _accumulate_1d_cpu_section!(op, shared; init=neutral, inclusive=true) + _accumulate_1d_cpu_section!(op, shared; seed=neutral, neutral, inclusive=true) - # Now prepend the running total of all previous tasks to each element, except in the first task + # Now prepend the running total of all previous tasks to each element, except in the first + # task. An exclusive section starts with its seed, `neutral`, which may be an empty lane that + # cannot be stored: its first element is the running total itself. itask_partition(tp) do itask, irange @inbounds begin if itask != 1 + carry = shared[itask - 1] for i in irange - v[i] = op(shared[itask - 1], v[i]) + v[i] = if !inclusive && i == first(irange) + _lower(eltype(v), carry) + else + _lower(eltype(v), op(carry, _lift(neutral, v[i]))) + end end end end @@ -51,18 +62,19 @@ function accumulate_1d_cpu!( end -# Scan a section sequentially, returning its total (including `init`). -function _accumulate_1d_cpu_section!(op, v; init, inclusive) +# Scan a section sequentially from `seed`, returning its total. `op` combines partial results +# (lanes when `neutral` is an empty lane); elements are lifted to them and stored back lowered. +function _accumulate_1d_cpu_section!(op, v; seed, neutral, inclusive) @inbounds begin - running = init - if inclusive - for i in eachindex(v) - running = op(running, v[i]) - v[i] = running - end - else - for i in eachindex(v) - v[i], running = running, op(running, v[i]) + running = seed + for i in eachindex(v) + x = _lift(neutral, v[i]) + if inclusive + running = op(running, x) + _store!(v, i, running) + else + _store!(v, i, running) + running = op(running, x) end end end diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 469ed176..280d9ce4 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -59,17 +59,19 @@ function _decoupled_fence end end -# Register-raking block scan with striped loads and stores. +# Register-raking block scan with striped loads and stores. `neutral` seeds partial results (an +# empty `_Lane` when `op` has no known neutral element), and `seed` is the first block's. @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_block!( - op, v, init, neutral, + op, v, seed, neutral, inclusive, flags, prefixes, ::Val{ITEMS}, ) where ITEMS # `block_size` is a power of two. @uniform block_size = @groupsize()[1] - tile = @localmem eltype(v) (block_size * ITEMS,) - thread_totals = @localmem eltype(v) (block_size,) + tile = @localmem typeof(neutral) (block_size * ITEMS,) + thread_totals = @localmem typeof(neutral) (block_size,) + _, op = _lanefuncs(_Partials(), op, neutral) # Internal indices are zero-based; add one only when indexing arrays. len = length(v) @@ -82,7 +84,7 @@ end while j < ITEMS p = j * block_size + ithread gi = block_offset + p - tile[p + 0x1] = gi < len ? v[gi + 0x1] : neutral + tile[p + 0x1] = gi < len ? _lift(neutral, v[gi + 0x1]) : neutral j += 1 end @synchronize() @@ -100,9 +102,9 @@ end # Scan the per-thread totals. Later blocks receive their carry from the # second kernel. - seed = iblock == 0x0 ? init : neutral + block_seed = iblock == 0x0 ? seed : neutral thread_prefix, block_total = block_exclusive_scan!( - @context, op, thread_totals, seed, block_size, ithread, + @context, op, thread_totals, block_seed, block_size, ithread, ) # DecoupledLookback keeps later blocks inclusive until the carry pass. @@ -133,7 +135,7 @@ end p = j * block_size + ithread gi = block_offset + p if gi < len - v[gi + 0x1] = tile[p + 0x1] + _store!(v, gi + 0x1, tile[p + 0x1]) end j += 1 end @@ -142,8 +144,9 @@ end # Add each block's running prefix, stopping at a completed predecessor. @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_previous!( - op, v, flags, @Const(prefixes), ::Val{ITEMS}, + op, v, flags, @Const(prefixes), neutral, ::Val{ITEMS}, ) where ITEMS + _, op = _lanefuncs(_Partials(), op, neutral) len = length(v) @uniform block_size = @groupsize()[1] @@ -161,7 +164,8 @@ end ) if flag == ACC_FLAG_A _decoupled_fence() # acquire: order the `v` read after the flag load - running_prefix = op(v[(inspected_block + 0x1) * block_size * ITEMS], running_prefix) + running_prefix = op(_lift(neutral, v[(inspected_block + 0x1) * block_size * ITEMS]), + running_prefix) break else running_prefix = op(prefixes[inspected_block + 0x1], running_prefix) @@ -175,7 +179,7 @@ end while j < ITEMS gi = block_offset + j * block_size + ithread if gi < len - v[gi + 0x1] = op(running_prefix, v[gi + 0x1]) + v[gi + 0x1] = _lower(eltype(v), op(running_prefix, _lift(neutral, v[gi + 0x1]))) end j += 1 end @@ -195,8 +199,9 @@ end # Add pre-scanned block prefixes to each tile. @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_previous_coupled_preblocks!( - op, v, prefixes, ::Val{ITEMS}, + op, v, prefixes, neutral, inclusive, ::Val{ITEMS}, ) where ITEMS + _, op = _lanefuncs(_Partials(), op, neutral) len = length(v) @uniform block_size = @groupsize()[1] @@ -217,11 +222,17 @@ end running_prefix = op(carry, running_prefix) end + # An exclusive tile starts with its seed, `neutral`, which may be an empty lane that cannot + # be stored: its first element is the running prefix itself j = 0 while j < ITEMS gi = block_offset + j * block_size + ithread if gi < len - v[gi + 0x1] = op(running_prefix, v[gi + 0x1]) + v[gi + 0x1] = if !inclusive && gi == block_offset + _lower(eltype(v), running_prefix) + else + _lower(eltype(v), op(running_prefix, _lift(neutral, v[gi + 0x1]))) + end end j += 1 end @@ -230,10 +241,10 @@ end # Save the value preceding each tile before shifting the array in place. @kernel cpu=false inbounds=true function exclusive_prefixes_kernel!( - prefixes, @Const(v), init, elems_per_block, + prefixes, @Const(v), seed, neutral, elems_per_block, ) iblock = @index(Global, Linear) - 0x1 - prefixes[iblock + 0x1] = iblock == 0x0 ? init : v[iblock * elems_per_block] + prefixes[iblock + 0x1] = iblock == 0x0 ? seed : _lift(neutral, v[iblock * elems_per_block]) end @@ -264,13 +275,30 @@ end p = j * block_size + ithread gi = block_offset + p if gi < len - v[gi + 0x1] = p == 0x0 ? prefixes[iblock + 0x1] : tile[p] + v[gi + 0x1] = p == 0x0 ? _lower(eltype(v), prefixes[iblock + 0x1]) : tile[p] end j += 1 end end +# The partial-result buffers of a whole-array scan: the tiles' totals, and DecoupledLookback's +# block flags +function _scan_prefixes(v, neutral, num_blocks, temp) + isnothing(temp) && return similar(v, typeof(neutral), num_blocks) + neutral isa _Lane && throw(ArgumentError( + "`temp` needs a neutral element of the operator; pass `neutral`")) + eltype(temp) === eltype(v) || throw(ArgumentError( + "`temp` must have the scanned element type $(eltype(v)), got $(eltype(temp))")) + @argcheck length(temp) >= num_blocks + return view(temp, 1:num_blocks) +end + +# The first block's seed: `init` in the element type, else the neutral seed +_scan_first_seed(v, init, neutral) = _lift(neutral, convert(eltype(v), init)) +_scan_first_seed(v, ::_NoInit, neutral) = neutral + + # DecoupledLookback algorithm function accumulate_1d_gpu!( op, v::AbstractArray, backend::Backend, alg::DecoupledLookback; @@ -290,14 +318,8 @@ function accumulate_1d_gpu!( elems_per_block = block_size * items_per_thread num_blocks = (length(v) + elems_per_block - 1) ÷ elems_per_block items = Val(items_per_thread) - - if isnothing(temp) - prefixes = similar(v, eltype(v), num_blocks) - else - @argcheck eltype(temp) === eltype(v) - @argcheck length(temp) >= num_blocks - prefixes = view(temp, 1:num_blocks) - end + prefixes = _scan_prefixes(v, neutral, num_blocks, temp) + seed = _scan_first_seed(v, init, neutral) if isnothing(temp_flags) flags = similar(v, UInt8, num_blocks) @@ -311,18 +333,18 @@ function accumulate_1d_gpu!( block_inclusive = inclusive || shift_to_exclusive kernel1! = _accumulate_block!(backend, block_size) - kernel1!(op, v, init, neutral, block_inclusive, flags, prefixes, items, + kernel1!(op, v, seed, neutral, block_inclusive, flags, prefixes, items, ndrange=num_blocks * block_size) if num_blocks > 1 kernel2! = _accumulate_previous!(backend, block_size) - kernel2!(op, v, flags, prefixes, items, + kernel2!(op, v, flags, prefixes, neutral, items, ndrange=(num_blocks - 1) * block_size) end if shift_to_exclusive - exclusive_prefixes_kernel!(backend, block_size)(prefixes, v, init, elems_per_block, - ndrange=num_blocks) + exclusive_prefixes_kernel!(backend, block_size)(prefixes, v, seed, neutral, + elems_per_block, ndrange=num_blocks) exclusive_shift_kernel!(backend, block_size)(v, prefixes, items, ndrange=num_blocks * block_size) end @@ -350,22 +372,16 @@ function accumulate_1d_gpu!( elems_per_block = block_size * items_per_thread num_blocks = (length(v) + elems_per_block - 1) ÷ elems_per_block items = Val(items_per_thread) - - if isnothing(temp) - prefixes = similar(v, eltype(v), num_blocks) - else - @argcheck eltype(temp) === eltype(v) - @argcheck length(temp) >= num_blocks - prefixes = view(temp, 1:num_blocks) - end + prefixes = _scan_prefixes(v, neutral, num_blocks, temp) + seed = _scan_first_seed(v, init, neutral) kernel1! = _accumulate_block!(backend, block_size) - kernel1!(op, v, init, neutral, inclusive, nothing, prefixes, items, + kernel1!(op, v, seed, neutral, inclusive, nothing, prefixes, items, ndrange=num_blocks * block_size) if num_blocks > 1 - # Accumulate prefixes of all blocks; use neutral as init here to not reinclude init + # Accumulate prefixes of all blocks; seed with neutral here to not reinclude init num_blocks_prefixes = (length(prefixes) + elems_per_block - 1) ÷ elems_per_block kernel1!(op, prefixes, neutral, neutral, true, nothing, nothing, items, ndrange=num_blocks_prefixes * block_size) @@ -373,7 +389,7 @@ function accumulate_1d_gpu!( # Prefixes are pre-accumulated (completely accumulated if num_blocks_prefixes == 1, or # partially, which we will account for in the coupled lookback) kernel2! = _accumulate_previous_coupled_preblocks!(backend, block_size) - kernel2!(op, v, prefixes, items, + kernel2!(op, v, prefixes, neutral, inclusive, items, ndrange=(num_blocks - 1) * block_size) end diff --git a/src/accumulate/accumulate_nd.jl b/src/accumulate/accumulate_nd.jl index 782a0dfa..7cb2ea1d 100644 --- a/src/accumulate/accumulate_nd.jl +++ b/src/accumulate/accumulate_nd.jl @@ -18,8 +18,9 @@ function accumulate_nd!( # Degenerate cases end + seed = _scan_first_seed(v, init, neutral) if alg isa CPUThreads.Partitioned - _accumulate_nd_cpu_sections!(op, v; init, dims, inclusive, + _accumulate_nd_cpu_sections!(op, v; seed, neutral, dims, inclusive, max_tasks=alg.max_tasks, min_elems=alg.min_elems) else block_size = alg.block_size @@ -38,7 +39,7 @@ function accumulate_nd!( blocks = (length_outer + block_size - 1) ÷ block_size kernel1! = _accumulate_nd_by_thread!(backend, block_size) kernel1!( - v, op, init, dims, inclusive, + v, op, seed, neutral, dims, inclusive, ndrange=(block_size * blocks,), ) else @@ -46,7 +47,7 @@ function accumulate_nd!( blocks = length_outer kernel2! = _accumulate_nd_by_block!(backend, block_size) kernel2!( - v, op, init, neutral, dims, inclusive, + v, op, seed, neutral, dims, inclusive, ndrange=(block_size, blocks), ) end @@ -56,13 +57,18 @@ function accumulate_nd!( end +# The kernels index `v` linearly, so they step through it with the strides of a column-major array +# of its size (`Base.size_to_strides`), not with `strides(v)`: a wrapper's storage layout, such as a +# `PermutedDimsArray`'s, differs from its linear indices. + function _accumulate_nd_cpu_sections!( op, v::AbstractArray; - init, dims, inclusive, + seed, neutral, dims, inclusive, max_tasks, min_elems, ) + _, op = _lanefuncs(_Partials(), op, neutral) vsizes = size(v) - vstrides = strides(v) + vstrides = Base.size_to_strides(1, vsizes...) ndims = length(vsizes) @@ -85,18 +91,16 @@ function _accumulate_nd_cpu_sections!( end # Go over each element in the accumulated dimension - if inclusive - running = init - for i in 0:length_dims - 1 - v_idx = input_base_idx + i * vstrides[dims] - running = op(running, v[v_idx + 1]) - v[v_idx + 1] = running - end - else - running = init - for i in 0:length_dims - 1 - v_idx = input_base_idx + i * vstrides[dims] - v[v_idx + 1], running = running, op(running, v[v_idx + 1]) + running = seed + for i in 0:length_dims - 1 + v_idx = input_base_idx + i * vstrides[dims] + x = _lift(neutral, v[v_idx + 1]) + if inclusive + running = op(running, x) + v[v_idx + 1] = _unlane(running) + else + v[v_idx + 1] = _unlane(running) + running = op(running, x) end end end @@ -107,13 +111,14 @@ end @kernel inbounds=true cpu=false unsafe_indices=true function _accumulate_nd_by_thread!( - v, op, init, dims, inclusive, + v, op, seed, neutral, dims, inclusive, ) + _, op = _lanefuncs(_Partials(), op, neutral) # One thread per outer dimension element, when there are more outer elements than in the # reduced dim e.g. accumulate(+, rand(3, 1000), dims=1) => only 3 elements in the accumulated # dim vsizes = size(v) - vstrides = strides(v) + vstrides = Base.size_to_strides(1, vsizes...) ndims = length(vsizes) @@ -148,18 +153,16 @@ end # Go over each element in the accumulated dimension; this implementation assumes that there # are so many outer elements (each processed by an independent thread) that we afford to # loop sequentially over the accumulated dimension (e.g. reduce(+, rand(3, 1000), dims=1)) - if inclusive - running = init - for i in 0x0:length_dims - 0x1 - v_idx = input_base_idx + i * vstrides[dims] - running = op(running, v[v_idx + 0x1]) - v[v_idx + 0x1] = running - end - else - running = init - for i in 0x0:length_dims - 0x1 - v_idx = input_base_idx + i * vstrides[dims] - v[v_idx + 0x1], running = running, op(running, v[v_idx + 0x1]) + running = seed + for i in 0x0:length_dims - 0x1 + v_idx = input_base_idx + i * vstrides[dims] + x = _lift(neutral, v[v_idx + 0x1]) + if inclusive + running = op(running, x) + v[v_idx + 0x1] = _unlane(running) + else + v[v_idx + 0x1] = _unlane(running) + running = op(running, x) end end end @@ -167,7 +170,7 @@ end @kernel inbounds=true cpu=false unsafe_indices=true function _accumulate_nd_by_block!( - v, op, init, neutral, dims, inclusive, + v, op, seed, neutral, dims, inclusive, ) # NOTE: shmem_size MUST be greater than 2 * block_size # NOTE: block_size MUST be a power of 2 @@ -176,7 +179,7 @@ end # than in outer dimensions, e.g. accumulate(+, rand(3, 1000), dims=2) => only 3 elements in # outer dimensions vsizes = size(v) - vstrides = strides(v) + vstrides = Base.size_to_strides(1, vsizes...) ndims = length(vsizes) @@ -185,8 +188,9 @@ end @uniform block_size = @groupsize()[1] - temp = @localmem eltype(v) (0x2 * block_size + conflict_free_offset(0x2 * block_size),) - running_prefix = @localmem eltype(v) (1,) + temp = @localmem typeof(neutral) (0x2 * block_size + conflict_free_offset(0x2 * block_size),) + running_prefix = @localmem typeof(neutral) (1,) + _, op = _lanefuncs(_Partials(), op, neutral) # NOTE: for many index calculations in this library, computation using zero-indexing leads to # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero @@ -211,13 +215,13 @@ end end # We have a block of threads to accumulate along the dims axis; do it in chunks of - # 2 * block_size and carry the total of all previous chunks (seeded with `init`) into each one. + # 2 * block_size and carry the total of all previous chunks (seeded with `seed`) into each one. # Operands are combined in element order, so `op` need not be commutative. ichunk = typeof(iblock)(0) num_chunks = (length_dims + (0x2 * block_size) - 0x1) ÷ (0x2 * block_size) if ithread == 0x0 - running_prefix[0x1] = init + running_prefix[0x1] = seed end while ichunk < num_chunks @@ -231,20 +235,20 @@ end bank_offset_b = conflict_free_offset(bi) xa = if block_offset + ai < length_dims - v[ + _lift(neutral, v[ input_base_idx + # Outer element axis starting index (block_offset + ai) * vstrides[dims] + # Move along dims axis in strides 0x1 # - to 1-indexing - ] + ]) else neutral end xb = if block_offset + bi < length_dims - v[ + _lift(neutral, v[ input_base_idx + (block_offset + bi) * vstrides[dims] + 0x1 - ] + ]) else neutral end @@ -312,14 +316,14 @@ end input_base_idx + (block_offset + ai) * vstrides[dims] + 0x1 - ] = op(carry, ra) + ] = _unlane(op(carry, ra)) end if block_offset + bi < length_dims v[ input_base_idx + (block_offset + bi) * vstrides[dims] + 0x1 - ] = op(carry, rb) + ] = _unlane(op(carry, rb)) end # Every thread has read the carry; the last thread extends it by this chunk's total (the diff --git a/src/accumulate/tuning.jl b/src/accumulate/tuning.jl index 266b40d2..c8756ff8 100644 --- a/src/accumulate/tuning.jl +++ b/src/accumulate/tuning.jl @@ -69,25 +69,29 @@ end """ scan_tuning(backend, T) -> ScanTuning -The scan tuning for element type `T` (the destination's) on `backend`'s current device; see -[`sort_tuning`](@ref) for the conventions. +The scan tuning for element type `T` (the running-value type the scan computes in; see +[`accumulate!`](@ref)) on `backend`'s current device; see [`sort_tuning`](@ref) for the +conventions. """ scan_tuning(::Backend, ::Type) = ScanTuning() """ - _resolve_scan(alg, backend, T, dims) -> Algorithm + _resolve_scan(alg, backend, T, dims, S=T) -> Algorithm Resolve `alg` for a scan of element type `T` along `dims` (`nothing` for the whole array in -linear order) on `backend`. Returns `ScanPrefixes`, `DecoupledLookback`, `SliceScan` or -`CPUThreads.Partitioned` with every field set, or throws an `ArgumentError`. +linear order) on `backend`, whose kernels hold partial results of type `S` in local memory (a +lane type when the operator has no known neutral element). Returns `ScanPrefixes`, +`DecoupledLookback`, `SliceScan` or `CPUThreads.Partitioned` with every field set, or throws an +`ArgumentError`. """ -function _resolve_scan(alg::Algorithm, backend::Backend, ::Type{T}, dims) where {T} +function _resolve_scan(alg::Algorithm, backend::Backend, ::Type{T}, dims, + ::Type{S}=T) where {T, S} _checkdomain(alg) t = scan_tuning(backend, T) a = alg isa Auto ? _select_scan(backend, t, dims) : alg - a = _fill(a, t, T) - _check_scan(a, backend, T, dims) + a = _fill(a, t, S) + _check_scan(a, backend, S, dims) return a end diff --git a/src/arithmetics.jl b/src/arithmetics.jl index ccd6c26d..f0d6a8f9 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -129,10 +129,11 @@ end """ - cumsum(src::AbstractArray; init=zero(eltype(src)), neutral=zero(eltype(src)), kwargs...) + cumsum(src::AbstractArray; kwargs...) -Cumulative sum of elements of an array, with optional `init` and `dims`. The keywords are those -of [`accumulate`](@ref). +Cumulative sum of elements of an array, as `Base.cumsum` (Base's `add_sum`, so small integers are +summed as `Int`), except that without `dims` a multidimensional array is summed in linear order. +The keywords are those of [`accumulate`](@ref). # Examples Simple cumulative sum of elements in a vector: @@ -150,16 +151,15 @@ m = ROCArray(rand(Int32(1):Int32(100), 10, 100_000)) s = AK.cumsum(m, dims=1) ``` """ -function cumsum(src::AbstractArray; init=zero(eltype(src)), neutral=zero(eltype(src)), kwargs...) - accumulate(+, src; init, neutral, inclusive=true, kwargs...) -end +cumsum(src::AbstractArray; kwargs...) = accumulate(Base.add_sum, src; kwargs...) """ - cumprod(src::AbstractArray; init=one(eltype(src)), neutral=one(eltype(src)), kwargs...) + cumprod(src::AbstractArray; kwargs...) -Cumulative product of elements of an array, with optional `init` and `dims`. The keywords are -those of [`accumulate`](@ref). +Cumulative product of elements of an array, as `Base.cumprod` (Base's `mul_prod`), except that +without `dims` a multidimensional array is multiplied in linear order. The keywords are those of +[`accumulate`](@ref). # Examples Simple cumulative product of elements in a vector: @@ -177,6 +177,4 @@ m = oneArray(rand(Int32(1):Int32(100), 10, 100_000)) p = AK.cumprod(m, dims=1) ``` """ -function cumprod(src::AbstractArray; init=one(eltype(src)), neutral=one(eltype(src)), kwargs...) - accumulate(*, src; init, neutral, inclusive=true, kwargs...) -end +cumprod(src::AbstractArray; kwargs...) = accumulate(Base.mul_prod, src; kwargs...) diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index f89b2229..47e62112 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -133,6 +133,24 @@ _AccOp{A}(op::OP) where {A, OP} = _AccOp{A, OP}(op) @inline _unlane(x) = x @inline _unlane(x::_Lane) = x.value +# Scans keep elements and partial results in one array or tile: an element enters the seed's +# representation with `_lift` (partials pass through), and a partial is stored in an array of +# element type `E` with `_lower` (as a lane only in arrays of lanes). +@inline _lift(neutral, x) = x +@inline _lift(::_Lane{T}, x) where {T} = _Lane{T}(x) +@inline _lift(::_Lane{T}, x::_Lane{T}) where {T} = x +@inline _lower(::Type{E}, p) where {E} = _unlane(p) +@inline _lower(::Type{E}, p) where {E <: _Lane} = p + +# Store the partial `p` in `v[i]`, except an empty lane where elements are stored (a scan's +# carry pass fills that element later; its value field may not even be defined) +@inline function _store!(v, i, p) + if !(p isa _Lane) || eltype(v) <: _Lane || _valid(p) + @inbounds v[i] = _lower(eltype(v), p) + end + return nothing +end + # The seed of every partial result with accumulator type `A`: the caller's `neutral`, else # GPUArraysCore's neutral element for `op` when it has one, else an empty lane. A seed that is not # a lane stands for the accumulator type, so an abstract one (only on the host) is a lane: holding diff --git a/test/Project.toml b/test/Project.toml index bc0d462b..9ce1e943 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -18,3 +18,4 @@ pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" [compat] ParallelTestRunner = "2" +oneAPI = "2.9.2" diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index cc372270..f404479a 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -138,7 +138,7 @@ end # Oversized temporaries are allowed. y = array_from_host(ones(Int32, 1000)) - AK.accumulate!(+, y; init=0, inclusive=false, temp=array_from_host(zeros(Int32, 1000)), + AK.accumulate!(+, y; init=Int32(0), inclusive=false, temp=array_from_host(zeros(Int32, 1000)), temp_flags=array_from_host(zeros(Int8, 1000)), alg=scan_alg(A; block_size=128)) @test Array(y) == 0:999 @@ -216,13 +216,16 @@ end # Single and multiple blocks, and more blocks than one block can scan (block_size=16), so that # the block scan, the lookback / block-prefix carry and the chunked prefix carry are all # exercised. On the CPU, `max_tasks=4` exercises the carry between tasks. - for inclusive in (true, false), (block_size, items_per_thread) in ((256, nothing), (16, 1), (32, 3)) + # With a neutral element, and without one (partial results then start from their first + # element) + for neutral in (SCAN_I2, nothing), inclusive in (true, false), + (block_size, items_per_thread) in ((256, nothing), (16, 1), (32, 3)) # Includes exact tile boundaries (16 and 96 elements) and one past them for n in (1, 2, 5, 16, 17, 96, 97, 100, 1000, 5000, 70_000) xh = [scan_randmat() for _ in 1:n] init = scan_randmat() y = array_from_host(xh) - AK.accumulate!(scan_matmul, y; init, neutral=SCAN_I2, inclusive, + AK.accumulate!(scan_matmul, y; init, neutral, inclusive, alg=scan_alg(A; max_tasks=4, block_size, items_per_thread)) @test Array(y) == scan_reference(xh, 1; init, inclusive) end @@ -235,7 +238,7 @@ end # Both GPU strategies: one thread per slice when there are more slices than elements per # slice, else one block per slice, processing the slice in several chunks of 2 * block_size - for inclusive in (true, false), block_size in (64, 256) + for neutral in (SCAN_I2, nothing), inclusive in (true, false), block_size in (64, 256) # Slices of exactly one and two chunks, and one element more, for both block sizes chunk = 2 * block_size for (sz, dims) in (((3, 2000), 2), ((2000, 3), 1), ((2000, 3), 2), ((7, 600), 2), @@ -244,7 +247,7 @@ end xh = [scan_randmat() for _ in CartesianIndices(sz)] init = scan_randmat() y = array_from_host(xh) - AK.accumulate!(scan_matmul, y; init, neutral=SCAN_I2, inclusive, dims, + AK.accumulate!(scan_matmul, y; init, neutral, inclusive, dims, alg=slice_alg(; max_tasks=4, block_size)) @test Array(y) == scan_reference(xh, dims; init, inclusive) end @@ -463,3 +466,167 @@ end # The other settings are stress-tested in reduce end + + +@testset "scan contract" begin + Random.seed!(0) + xh = rand(Int32(-9):Int32(9), 20_000) + x = array_from_host(xh) + + # Inclusive scans without `init` start from the first element; with it, `init` is applied + # once, also when the scan spans many blocks + for A in SCAN_ALGS + alg = scan_alg(A; block_size=32, items_per_thread=1, max_tasks=4) + @test Array(AK.accumulate(+, x; alg)) == cumsum(xh) + @test Array(AK.accumulate(+, x; init=Int32(5), alg)) == accumulate(+, xh; init=Int32(5)) + # Exclusive scans start from `init`, else from the neutral element + @test Array(AK.accumulate(+, x; inclusive=false, alg)) == cumsum(xh) .- xh + @test Array(AK.accumulate(+, x; init=Int32(5), inclusive=false, alg)) == + cumsum(xh) .- xh .+ 5 + # Operators without a known neutral element + @test Array(AK.accumulate((a, b) -> a + b, x; alg)) == cumsum(xh) + @test Array(AK.accumulate((a, b) -> max(a, b), x; init=Int32(0), alg)) == + accumulate(max, xh; init=Int32(0)) + @test Array(AK.accumulate((a, b) -> a + b, x; init=Int32(5), inclusive=false, alg)) == + cumsum(xh) .- xh .+ 5 + @test_throws ArgumentError AK.accumulate((a, b) -> a + b, x; inclusive=false, alg) + end + + # The same along `dims` + mh = rand(Int32(-9):Int32(9), 70, 300) + m = array_from_host(mh) + for dims in (1, 2), alg in (slice_alg(), slice_alg(; block_size=32, max_tasks=4)) + @test Array(AK.accumulate(+, m; dims, alg)) == cumsum(mh; dims) + @test Array(AK.accumulate((a, b) -> a + b, m; dims, alg)) == cumsum(mh; dims) + @test Array(AK.accumulate(+, m; dims, init=Int32(2), inclusive=false, alg)) == + cumsum(mh; dims) .- mh .+ 2 + @test_throws ArgumentError AK.accumulate((a, b) -> a + b, m; dims, inclusive=false, alg) + end + # A `dims` beyond the array's makes every slice one element long: inclusive scans apply `init` + # to every element, and exclusive ones start every slice over + @test Array(AK.accumulate(+, m; dims=3, alg=slice_alg())) == mh + @test Array(AK.accumulate(+, m; dims=3, init=Int32(7), alg=slice_alg())) == mh .+ Int32(7) + @test Array(AK.accumulate(+, m; dims=3, init=7, alg=slice_alg())) == mh .+ 7 + @test Array(AK.accumulate(+, m; dims=3, init=Int32(7), inclusive=false, alg=slice_alg())) == + fill(Int32(7), size(mh)) + + # `accumulate` allocates the fold type of `op` from `init`'s type and the elements + h8 = Int8[1, 2, 100] + for (r, ref) in ((AK.accumulate(+, array_from_host(h8); alg=scan_alg()), accumulate(+, h8)), + (AK.accumulate(+, array_from_host(h8); init=0, alg=scan_alg()), + accumulate(+, h8; init=0)), + (AK.cumsum(array_from_host(h8); alg=scan_alg()), cumsum(h8)), + (AK.cumsum(array_from_host(Bool[1, 1, 0]); alg=scan_alg()), cumsum(Bool[1, 1, 0])), + (AK.cumprod(array_from_host(Int8[2, 3]); alg=scan_alg()), cumprod(Int8[2, 3]))) + @test eltype(r) === eltype(ref) && Array(r) == ref + end + @test Array(AK.accumulate(+, array_from_host(Int32[]); alg=scan_alg())) == Int32[] + + # `init=nothing` is an initial value, not an omitted one: `op` applies it + @test_throws Exception AK.accumulate(+, x; init=nothing, alg=scan_alg()) + something_add(a, b) = something(a, Int32(0)) + something(b, Int32(0)) + @test Array(AK.accumulate(something_add, x; init=nothing, alg=scan_alg())) == + accumulate(something_add, xh; init=nothing) + + # The scan runs in the running-value type and converts only when storing (on oneAPI, a kernel + # whose conversion may throw an `InexactError` needs oneAPI.jl 2.9.2's device heap) + fh = Float32[0.5, 1.0, 2.0] + @test Array(AK.accumulate!(+, array_from_host(zeros(Int32, 3)), array_from_host(fh); + init=0.5f0, alg=scan_alg())) == + accumulate!(+, zeros(Int32, 3), fh; init=0.5f0) + @test Array(AK.accumulate!(+, array_from_host(zeros(Int32, 1, 3)), + array_from_host(reshape(fh, 1, 3)); + init=0.5f0, dims=2, alg=slice_alg())) == + accumulate!(+, zeros(Int32, 1, 3), reshape(fh, 1, 3); init=0.5f0, dims=2) + # ... in which the elements are one-element partial results: `addmod`'s operands need not + # fit its result type + addmod(a, b) = Int32(mod(Int64(a) + Int64(b), 97)) + bigs = Int64[2147483648, 1, 2] + @test Array(AK.accumulate(addmod, array_from_host(bigs); init=Int32(0), alg=scan_alg())) == + accumulate(addmod, bigs; init=Int32(0)) + # ... which starts from the destination's element type as well, so a wider destination + # holds the running values without a scratch array + @test Array(AK.accumulate!(*, array_from_host(zeros(Int, 4)), array_from_host(fill(0x10, 4)); + alg=scan_alg())) == [16, 256, 4096, 65536] + @test Array(AK.accumulate!(*, array_from_host(zeros(Int, 2, 3)), + array_from_host(fill(0x10, 2, 3)); dims=2, alg=slice_alg())) == + repeat([16 256 4096], 2) + if Float64 in valid_backend_eltypes(BACKEND, (Float64,)) + @test Array(AK.accumulate!(+, array_from_host(zeros(Float64, 3)), + array_from_host(Float32[1.0f8, 1, -1.0f8]); alg=scan_alg())) == + [1e8, 100000001, 1] + end + # `acctype` sets the running-value type, and `accumulate`'s element type + r = AK.accumulate(+, array_from_host(Int8[100, 100]); acctype=Int8, alg=scan_alg()) + @test eltype(r) === Int8 && Array(r) == Int8[100, -56] + r = AK.cumsum(array_from_host(Int8[100, 100]); acctype=Int16, alg=scan_alg()) + @test eltype(r) === Int16 && Array(r) == Int16[100, 200] + r = AK.accumulate!(+, array_from_host(zeros(Int32, 3)), array_from_host(Int8[100, 100, 100]); + acctype=Int8, alg=scan_alg()) + @test Array(r) == Int32[100, -56, 44] + @test_throws ArgumentError AK.accumulate(+, x; acctype=String, alg=scan_alg()) + # A scan that never calls `op` only copies: along a `dims` beyond the array's without `init`, + # or of one element + @test Array(AK.accumulate!(+, array_from_host(zeros(Int32, 2)), array_from_host(Int32[1, 2]); + dims=2, alg=slice_alg())) == [1, 2] + @test Array(AK.accumulate!(+, array_from_host(zeros(Int32, 2)), array_from_host(Int32[1, 2]); + dims=2, init=Int32(1), alg=slice_alg())) == [2, 3] + if !TEST_KERNELS + @test AK.accumulate(+, ['a']; alg=scan_alg()) == ['a'] # (Base's fails on Julia 1.10) + @test AK.accumulate(coalesce, [missing, 1, 2]; init=0, alg=scan_alg()) == + accumulate(coalesce, [missing, 1, 2]; init=0) + end + # ... from the exact neutral element, so the first element stays as it is + z = Array(AK.accumulate(+, array_from_host([-0.0f0, -0.0f0]); alg=scan_alg())) + @test Base.all(signbit, z) + if !TEST_KERNELS + # Values that are not bits types, on the host: elements enter as `Base.reduce_first` + # (a `Char` becomes a `String`), and the lanes of an operator without a neutral element + # are never stored empty + @test AK.accumulate(*, collect("abcdef"); alg=scan_alg(max_tasks=3, min_elems=2)) == + accumulate(*, collect("abcdef")) + @test AK.accumulate((a, b) -> a * b, fill("x", 8); init="s", inclusive=false, + alg=scan_alg(max_tasks=4, min_elems=2)) == + ["s" * "x"^k for k in 0:7] + end + @test length(AK.accumulate(+, array_from_host(Int32[]); init=nothing, alg=scan_alg())) == 0 + if !TEST_KERNELS + # An abstract running type, with the caller's neutral element + for alg in (scan_alg(), scan_alg(max_tasks=2, min_elems=2)) + @test AK.accumulate(+, Real[1, 2.5, 3]; neutral=0, inclusive=false, alg) == [0, 1, 3.5] + @test AK.accumulate(+, [1, 2]; acctype=Real, neutral=0, inclusive=false, alg) == [0, 1] + end + end + + # The destination is the source or does not overlap it + y = copy(x) + @test AK.accumulate!(+, y, y; alg=scan_alg()) === y + @test Array(y) == cumsum(xh) + # (where the array type lets `Base.mightalias` see the overlap: Metal's contiguous views are + # separate arrays) + if Base.mightalias(view(x, 2:100), view(x, 1:99)) + @test_throws ArgumentError AK.accumulate!(+, view(x, 2:100), view(x, 1:99); alg=scan_alg()) + end + @test_throws DimensionMismatch AK.accumulate!(+, array_from_host(zeros(Int32, 5)), x; + alg=scan_alg()) +end + + +@testset "accumulate: backend-free inputs" begin + # A range has no backend: the result is allocated on the one given + r = AK.cumsum(1:5; backend=BACKEND, alg=scan_alg()) + @test get_backend(r) == BACKEND + @test Array(r) == cumsum(1:5) + r = AK.accumulate(+, 1:5; backend=BACKEND, init=10, alg=scan_alg()) + @test Array(r) == accumulate(+, 1:5; init=10) +end + + +@testset "accumulate: dims of wrapped arrays" begin + # Linear indices of a permuted or strided view are not its storage offsets + h = reshape(Int32.(1:24), 4, 6) + for wrap in (A -> PermutedDimsArray(A, (2, 1)), A -> view(A, 1:2:4, :)), dims in (1, 2) + v = wrap(array_from_host(h)) + @test Array(AK.accumulate!(+, v; dims, alg=slice_alg())) == accumulate(+, wrap(h); dims) + end +end diff --git a/test/generic/scan_resolve.jl b/test/generic/scan_resolve.jl index 45dc5363..1723cbd3 100644 --- a/test/generic/scan_resolve.jl +++ b/test/generic/scan_resolve.jl @@ -100,6 +100,13 @@ end @test resolve_scan(AK.DecoupledLookback()) === AK.DecoupledLookback(256, 8) end + # Tiles of partial results: lanes (a Bool and a value) are wider than elements, so fewer fit. + # NTuple{3, Int32} is 4-byte aligned everywhere, so its lanes take 16 bytes on every platform + # (Int64 lanes take 12 on 32-bit x86) + @test resolve_scan(AK.Auto(); T=NTuple{3, Int32}) === AK.ScanPrefixes(256, 8) + @test AK._resolve_scan(AK.Auto(), SRB, NTuple{3, Int32}, nothing, + AK._Lane{NTuple{3, Int32}}) === AK.ScanPrefixes(256, 7) + # Kernels need a bits-type element @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(); T=String) @test_throws ArgumentError resolve_scan(AK.Auto(); T=Union{Missing, Int}, dims=1) From 135188c95933f4ae56ebf6bb7e4fa148ec140d86 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 07:52:43 +0200 Subject: [PATCH 7/8] Move findall and the predicates onto algorithm selection `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. --- docs/src/api/algorithms.md | 6 +- docs/src/api/custom_structs.md | 2 +- docs/src/api/differences.md | 2 + docs/src/api/findall.md | 17 ++ docs/src/api/predicates.md | 16 +- docs/src/tuning.md | 9 +- ext/oneAPIExt.jl | 34 +--- src/algorithms.jl | 17 ++ src/arithmetics.jl | 8 +- src/findall.jl | 226 ++++++++++++-------- src/predicates.jl | 362 ++++++++++++++------------------- src/reduce/utilities.jl | 6 + src/utils.jl | 6 - test/generic/findall.jl | 175 ++++++++++++---- test/generic/predicates.jl | 104 ++++++++-- test/generic/reduce.jl | 12 -- test/runtests.jl | 9 +- 17 files changed, 587 insertions(+), 424 deletions(-) diff --git a/docs/src/api/algorithms.md b/docs/src/api/algorithms.md index 79a58d76..5470e66e 100644 --- a/docs/src/api/algorithms.md +++ b/docs/src/api/algorithms.md @@ -12,7 +12,8 @@ AcceleratedKernels.CPUThreads ``` The algorithm types of each operation are documented with it, for [sorting](sort.md), -[reductions](reduce.md) and [scans](accumulate.md). +[reductions](reduce.md), [scans](accumulate.md), [`findall`](findall.md) and +[`any`/`all`](predicates.md). #### How `Auto` chooses @@ -22,7 +23,8 @@ host backend it chooses the family's `CPUThreads` algorithm, which runs on Julia it follows per-device tuning values: for sorting, `BitonicSort` for short arrays and slices, `RadixSort` for long whole arrays of the element types it supports, and `MergeSort` otherwise, subject to `Auto`'s requirements (`stable=true` by default); for reductions, `BlockReduce`; for -scans, `ScanPrefixes` for whole arrays and `SliceScan` along `dims`. +scans, `ScanPrefixes` for whole arrays and `SliceScan` along `dims`; for `findall`, +`ScanScatter`; for `any`/`all`, `ConcurrentWrite` (`ViaReduce` on oneAPI). Algorithms carry their settings as fields, and a field left at `nothing` takes the device's tuned value: diff --git a/docs/src/api/custom_structs.md b/docs/src/api/custom_structs.md index a00b8758..814794a0 100644 --- a/docs/src/api/custom_structs.md +++ b/docs/src/api/custom_structs.md @@ -14,7 +14,7 @@ using CUDA function complex_any(x, y) # Calling `any` on a normal Julia range, but running on x's backend - AK.any(1:length(x), AK.get_backend(x)) do i + AK.any(1:length(x); backend=AK.get_backend(x)) do i x[i] < 0 && y[i] > 0 end end diff --git a/docs/src/api/differences.md b/docs/src/api/differences.md index 61fbc25e..53d5bcd9 100644 --- a/docs/src/api/differences.md +++ b/docs/src/api/differences.md @@ -21,3 +21,5 @@ GPU arrays on top of AcceleratedKernels and reproduces Base's results. The diffe | the element type of `accumulate(op, A)` on Julia 1.10 | `Base.promote_op(op, T, T)` | the fold type, as Base's on Julia 1.13 | | `accumulate(op, A; dims, init)` with `dims > ndims(A)` | copies `A`, ignoring `init` | applies `init` to every element | | a non-associative `op` in a scan, such as `-` | works (a sequential recurrence) | unsupported | +| `findall(pred, A)` of a 0-dimensional `A` | `Int` indices | `CartesianIndex{0}`, `keys(A)`'s, unless `items=LinearIndices(A)` | +| `any`, `all` with a predicate that returns `missing` | three-valued logic | an error: the predicate must return a `Bool` (an `ArgumentError` before launching where inference shows it cannot) | diff --git a/docs/src/api/findall.md b/docs/src/api/findall.md index 5f2aa606..d580becf 100644 --- a/docs/src/api/findall.md +++ b/docs/src/api/findall.md @@ -2,5 +2,22 @@ ```@docs AcceleratedKernels.findall +``` + +`findall` selects items: indices by default, or anything else of the input's length, such as +the values themselves. + +```julia +v = CuArray(rand(Float32, 1000)) +AK.findall(x -> x > 0.5f0, v) # indices +AK.findall(x -> x > 0.5f0, v; items=v) # the selected values, as v[v .> 0.5f0] +``` + +[`Auto()`](@ref AcceleratedKernels.Auto) uses +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host and +[`ScanScatter`](@ref AcceleratedKernels.ScanScatter) on GPUs, with the device's settings. + +```@docs +AcceleratedKernels.FindallAlgorithm AcceleratedKernels.ScanScatter ``` diff --git a/docs/src/api/predicates.md b/docs/src/api/predicates.md index a407c3ae..d1bb4972 100644 --- a/docs/src/api/predicates.md +++ b/docs/src/api/predicates.md @@ -9,4 +9,18 @@ AcceleratedKernels.any AcceleratedKernels.all ``` -**Note on the `cooperative` keyword**: some older platforms crash when multiple threads write to the same memory location in a global array (e.g. old Intel Graphics); if all threads were to write the same value, it is well-defined on others (e.g. CUDA F4.2 says "If a non-atomic instruction executed by a warp writes to the same location in global memory for more than one of the threads of the warp, only one thread performs a write and which thread does it is undefined."). This "cooperative" thread behaviour allows for a faster implementation; if you have a platform - the only one I know is Intel UHD Graphics - that crashes, set `cooperative=false` to use a safer `mapreduce`-based implementation. +[`Auto()`](@ref AcceleratedKernels.Auto) uses +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host and +[`ConcurrentWrite`](@ref AcceleratedKernels.ConcurrentWrite) on GPUs, in which many threads write +the same value to one memory location. That is well-defined (CUDA F4.2: "If a non-atomic +instruction executed by a warp writes to the same location in global memory for more than one of +the threads of the warp, only one thread performs a write and which thread does it is +undefined."), but some older platforms (Intel UHD Graphics) have been reported to hang on it, so +on oneAPI `Auto` uses the `mapreduce`-based [`ViaReduce`](@ref AcceleratedKernels.ViaReduce) +instead. An explicit `alg=ConcurrentWrite()` runs on every GPU. + +```@docs +AcceleratedKernels.PredicateAlgorithm +AcceleratedKernels.ConcurrentWrite +AcceleratedKernels.ViaReduce +``` diff --git a/docs/src/tuning.md b/docs/src/tuning.md index 48ee2acb..ebf1cd2c 100644 --- a/docs/src/tuning.md +++ b/docs/src/tuning.md @@ -21,11 +21,14 @@ AcceleratedKernels._resolve_sort ``` The other families follow the same steps, with a one-line selection: reductions resolve with -`_resolve_reduce`, scans with `_resolve_scan`. +`_resolve_reduce`, scans with `_resolve_scan`, `findall` with `_resolve_findall` and +`any`/`all` with `_resolve_predicate`. ```@docs AcceleratedKernels._resolve_reduce AcceleratedKernels._resolve_scan +AcceleratedKernels._resolve_findall +AcceleratedKernels._resolve_predicate ``` `Auto` and explicit algorithms share steps 3 and 4, so a tuning cannot make an invalid algorithm @@ -43,6 +46,10 @@ AcceleratedKernels.ReduceTuning AcceleratedKernels.reduce_tuning AcceleratedKernels.ScanTuning AcceleratedKernels.scan_tuning +AcceleratedKernels.FindallTuning +AcceleratedKernels.findall_tuning +AcceleratedKernels.PredicateTuning +AcceleratedKernels.predicate_tuning ``` AcceleratedKernels defines the hook's generic method, whose values reproduce the library's diff --git a/ext/oneAPIExt.jl b/ext/oneAPIExt.jl index 66c2cea0..fa5bdfd0 100644 --- a/ext/oneAPIExt.jl +++ b/ext/oneAPIExt.jl @@ -12,36 +12,10 @@ oneAPI.@device_override AK._decoupled_fence() = SPIRV.atomic_work_item_fence(SPIRV.GLOBAL_MEM_FENCE, SPIRV.memory_order_seq_cst, SPIRV.memory_scope_device) -# On oneAPI, use the MapReduce algorithm by default as on some Intel GPUs ConcurrentWrite hangs -# the device. -function AK.any( - pred, v::AbstractArray, backend::oneAPIBackend; - - # Algorithm choice - alg::AK.PredicatesAlgorithm=AK.MapReduce(), - kwargs... -) - AK._any_impl( - pred, v, backend; - alg, - kwargs... - ) -end - - -function AK.all( - pred, v::AbstractArray, backend::oneAPIBackend; - - # Algorithm choice - alg::AK.PredicatesAlgorithm=AK.MapReduce(), - kwargs... -) - AK._all_impl( - pred, v, backend; - alg, - kwargs... - ) -end +# Some Intel GPUs (reportedly Intel UHD Graphics) hang when many threads write one global location, +# as `ConcurrentWrite` does. An Iris Xe does not, but the affected devices are not known, so `Auto` +# keeps using `ViaReduce` for `any`/`all` on oneAPI; an explicit `ConcurrentWrite` is allowed. +AK.predicate_tuning(::oneAPIBackend, ::Type) = AK.PredicateTuning(prefer_concurrent_write=false) end # module oneAPIExt diff --git a/src/algorithms.jl b/src/algorithms.jl index 7cb66763..a6df18af 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -61,6 +61,23 @@ Supertype of the scan (`accumulate`) algorithms: [`ScanPrefixes`](@ref), """ abstract type ScanAlgorithm <: Algorithm end +""" + FindallAlgorithm <: Algorithm + +Supertype of the stream-compaction (`findall`) algorithms: [`ScanScatter`](@ref). `findall` +also accepts [`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned). +""" +abstract type FindallAlgorithm <: Algorithm end + +""" + PredicateAlgorithm <: Algorithm + +Supertype of the algorithms of `any` and `all`: [`ConcurrentWrite`](@ref) and +[`ViaReduce`](@ref). They also accept +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned). +""" +abstract type PredicateAlgorithm <: Algorithm end + """ AcceleratedKernels.CPUThreads diff --git a/src/arithmetics.jl b/src/arithmetics.jl index f0d6a8f9..68b48f24 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -119,13 +119,7 @@ c = AK.count(CuArray(rand(Bool, 10, 100_000)); init=Int32(0), dims=2) """ count(src::AbstractArray; kwargs...) = count(identity, src; kwargs...) count(f, src::AbstractArray; init=0, kwargs...) = - mapreduce(_CountBool(f), Base.add_sum, src; init, kwargs...) - -# `f`, required to return a `Bool` (Base's `_bool`) -struct _CountBool{F} <: Function - f::F -end -@inline (c::_CountBool)(x) = c.f(x)::Bool + mapreduce(_BoolValued(f), Base.add_sum, src; init, kwargs...) """ diff --git a/src/findall.jl b/src/findall.jl index 8f9805f4..6fdbfc44 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -1,32 +1,92 @@ -abstract type FindallAlgorithm end - """ - ScanScatter(; block_size=256, items_per_thread=16) + ScanScatter(; block_size=nothing, items_per_thread=nothing) -Stable GPU stream compaction using per-block counts, a prefix scan, and a scatter pass. -`block_size` must be a power of two between 1 and 1024; `items_per_thread` must be positive. +Stable GPU stream compaction: each block counts the selected elements of its tile of +`block_size * items_per_thread` elements (`block_size` a power of two up to 1024), the counts are +scanned, and a second pass scatters the selected indices. """ Base.@kwdef struct ScanScatter <: FindallAlgorithm + block_size::Union{Nothing, Int} = nothing + items_per_thread::Union{Nothing, Int} = nothing +end + + +""" + FindallTuning(; kwargs...) + +Values that fill unset algorithm fields for `findall` on one device, as returned by +[`findall_tuning`](@ref): `block_size` and `items_per_thread` for `ScanScatter`, and +`threads_min_elems` for `CPUThreads.Partitioned`. The defaults reproduce AK's historical +settings. Internal: the fields may change in any release. +""" +Base.@kwdef struct FindallTuning block_size::Int = 256 items_per_thread::Int = 16 + threads_min_elems::Int = 1 +end + +""" + findall_tuning(backend, T) -> FindallTuning + +The `findall` tuning for element type `T` (the input's) on `backend`'s current device; see +[`sort_tuning`](@ref) for the conventions. +""" +findall_tuning(::Backend, ::Type) = FindallTuning() + +""" + _resolve_findall(alg, backend, T) -> Algorithm + +Resolve `alg` for `findall` over elements of type `T` on `backend`: `ScanScatter` or +`CPUThreads.Partitioned` with every field set, or an `ArgumentError`. +""" +function _resolve_findall(alg::Algorithm, backend::Backend, ::Type{T}) where {T} + _checkdomain(alg) + t = findall_tuning(backend, T) + a = alg isa Auto ? (_runs_threads(backend) ? CPUThreads.Partitioned() : ScanScatter()) : alg + a = _fill(a, t, T) + _check_findall(a, backend) + return a end +function _checkdomain(a::ScanScatter) + _check_pow2(a, :block_size) + a.block_size === nothing || a.block_size <= 1024 || throw(ArgumentError( + "ScanScatter: `block_size` must be at most 1024, got $(a.block_size)")) + _check_positive(a, :items_per_thread) + nothing +end -findall_algorithm(alg::ScanScatter) = alg -function findall_algorithm(alg::FindallAlgorithm) - throw(ArgumentError("$(typeof(alg)) is not supported by findall")) +_fill(a::ScanScatter, t::FindallTuning, T) = + ScanScatter(something(a.block_size, t.block_size), + something(a.items_per_thread, t.items_per_thread)) +_fill(a::CPUThreads.Partitioned, t::FindallTuning, T) = _fill_threads(a, t) +_fill(a::Algorithm, t::FindallTuning, T) = + throw(ArgumentError("$(_algname(a)) is not a findall algorithm")) + +function _check_findall(a::ScanScatter, backend) + _checkdomain(a) + _require_kernels(a, backend) + widemul(a.block_size, a.items_per_thread) <= typemax(Int32) || throw(ArgumentError( + "ScanScatter: `block_size * items_per_thread` must be at most $(typemax(Int32))")) + nothing end +function _check_findall(a::CPUThreads.Partitioned, backend) + _checkdomain(a) + _require_threads(a, backend) +end + +# The element of `indices` at the ordinal `position` (1-based, whatever the axes) @inline findall_index(indices::AbstractUnitRange, position) = first(indices) + position - 1 @inline findall_index(indices::LinearIndices{1}, position) = first(indices) + position - 1 -@inline findall_index(indices, position) = @inbounds indices[position] +@inline findall_index(indices, position) = @inbounds indices[firstindex(indices) + position - 1] # With `out === nothing`, compute block counts. Otherwise, `block_counts` contains their -# inclusive prefix scan and the kernel scatters the selected indices. +# inclusive prefix scan and the kernel scatters the selected items. @kernel cpu=false inbounds=true unsafe_indices=true function findall_block!( out, @Const(bools), block_counts, input_indices, output_indices, ::Val{ITEMS}, ) where ITEMS @@ -100,11 +160,9 @@ function findall_gpu( ) where I block_size = alg.block_size items_per_thread = alg.items_per_thread - @argcheck 1 <= block_size <= 1024 - @argcheck ispow2(block_size) - @argcheck items_per_thread > 0 - isempty(bools) && return similar(bools, I, 0) + # The output on the resolved backend, which a range does not determine + isempty(bools) && return KernelAbstractions.allocate(backend, I, 0) elems_per_block = block_size * items_per_thread num_blocks = cld(length(bools), elems_per_block) @@ -118,7 +176,7 @@ function findall_gpu( accumulate!(+, block_counts; backend, init=0) n = @allowscalar block_counts[end] - out = similar(bools, I, n) + out = KernelAbstractions.allocate(backend, I, n) if n > 0 kernel!(out, bools, block_counts, input_indices, output_indices, items; ndrange=num_blocks * block_size) @@ -128,13 +186,12 @@ end function findall_cpu( - bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend; - max_tasks::Int, - min_elems::Int, + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, + alg::CPUThreads.Partitioned; temp, ) where I input_indices = eachindex(bools) - tp = TaskPartitioner(length(bools), max_tasks, min_elems) + tp = TaskPartitioner(length(bools), alg.max_tasks, alg.min_elems) if tp.num_tasks == 1 out = similar(bools, I, Base.count(bools)) findall_section!(out, bools, input_indices, output_indices, Base.OneTo(length(bools)), 0) @@ -170,33 +227,25 @@ end function findall_impl( - bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend; - alg::FindallAlgorithm=ScanScatter(), - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - prefer_threads::Bool=true, - temp::Union{Nothing, AbstractArray}=nothing, + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg; + temp::Union{Nothing, AbstractArray}, ) where I - alg = findall_algorithm(alg) - - if use_gpu_algorithm(backend, prefer_threads) + if alg isa ScanScatter findall_gpu(bools, I, output_indices, backend, alg; temp) else - findall_cpu(bools, I, output_indices, backend; max_tasks, min_elems, temp) + findall_cpu(bools, I, output_indices, backend, alg; temp) end end -findall_output_indices(v, ::Type{Int}) = LinearIndices(v) -findall_output_indices(v, ::Type) = CartesianIndices(axes(v)) +# The launch settings of the mask pass for a resolved findall algorithm +_findall_launch(a::ScanScatter) = (; block_size=a.block_size) +_findall_launch(a::CPUThreads.Partitioned) = (; max_tasks=a.max_tasks, min_elems=a.min_elems) - -function findall_bools( - pred, v::AbstractArray, backend::Backend, temp_bools; - max_tasks, min_elems, prefer_threads, block_size, -) +function findall_bools(pred, v::AbstractArray, backend::Backend, alg, temp_bools) bools = if isnothing(temp_bools) - similar(v, Bool) + # On the resolved backend, which a range does not determine + KernelAbstractions.allocate(backend, Bool, size(v)) else @argcheck get_backend(temp_bools) === backend @argcheck eltype(temp_bools) === Bool @@ -206,7 +255,7 @@ function findall_bools( end input_indices = eachindex(v) bool_indices = eachindex(bools) - _foreachindex(Base.OneTo(length(v)), backend; max_tasks, min_elems, block_size) do position + _foreachindex(Base.OneTo(length(v)), backend; _findall_launch(alg)...) do position input_index = findall_index(input_indices, position) bool_index = findall_index(bool_indices, position) @inbounds bools[bool_index] = pred(v[input_index]) ? true : false @@ -216,37 +265,37 @@ end """ - findall(A::AbstractArray, backend::Backend=get_backend(A); - alg::FindallAlgorithm=ScanScatter(), - max_tasks::Int=Threads.nthreads(), min_elems::Int=1, - prefer_threads::Bool=true, - temp::Union{Nothing, AbstractArray}=nothing, - temp_bools::Union{Nothing, AbstractArray}=nothing) - findall(pred, A::AbstractArray, backend::Backend=get_backend(A); - alg::FindallAlgorithm=ScanScatter(), - max_tasks::Int=Threads.nthreads(), min_elems::Int=1, - prefer_threads::Bool=true, - temp::Union{Nothing, AbstractArray}=nothing, - temp_bools::Union{Nothing, AbstractArray}=nothing) - -Return the indices of the `true` elements of `A`, or of the elements for which `pred` returns -`true`, in the same order and with the same index types as `Base.findall`. Values used as -conditions must be `Bool`. + findall(A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), temp=nothing, + temp_bools=nothing) + findall(pred, A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), temp=nothing, + temp_bools=nothing) + +Stream compaction: select, in order, the elements of `items` at the positions where `A` is `true`, +or where `pred` returns `true` for `A`'s elements. Values used as conditions, and `pred`'s results, +must be `Bool`. The result is a new vector of `eltype(items)` on the backend. + +Positions are ordinal: the `k`-th element of `A`, in the order of `eachindex(A)`, selects the +`k`-th element of `items`, which may be any array with `A`'s length (so offset axes are no +problem). What is selected is the caller's choice: + +- `keys(A)` (the default) gives `A`'s indices: `Int`s for a vector, `CartesianIndex`es for + other arrays, including 0-dimensional ones. +- `LinearIndices(A)` gives linear indices of any array. +- An array of values selects those values: `AK.findall(mask; items=A)` is `A[mask]` for a mask of + `A`'s shape, in one pass instead of `findall` and a gather. The supported inputs are arrays. Dictionaries, other iterables, and scalar inputs accepted by `Base.findall` are outside the scope of this package. -## Settings - -- `alg=ScanScatter()` selects the GPU algorithm and its tuning parameters. -- `max_tasks=Threads.nthreads()` and `min_elems=1` control CPU task partitioning. -- `temp=nothing` may provide the `Int` buffer used for block or task counts. -- `temp_bools=nothing` may provide the Bool mask for the predicate form or for a mask whose - element type is not `Bool`. It must have the same axes as `A` and must not alias it. +`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref +AcceleratedKernels.CPUThreads.Partitioned) on the host and [`ScanScatter`](@ref) on GPUs, with +the device's settings. `backend` is derived from `A` and `items`. -On a GPU, `temp` needs at least -`cld(length(A), alg.block_size * alg.items_per_thread)` elements. On a CPU, it needs one element -per task used. Omitted buffers are allocated automatically. +`temp` may provide the `Int` buffer used for block or task counts: at least +`cld(length(A), block_size * items_per_thread)` elements for `ScanScatter`, one per task on the +host. `temp_bools` may provide the Bool mask for the predicate form or for a mask whose element +type is not `Bool`; it must have the same axes as `A` and must not alias it. Omitted buffers are +allocated automatically. # Examples ```julia @@ -255,50 +304,51 @@ import AcceleratedKernels as AK v = CUDA.CuArray(Int32[5, -2, 8, -1, 3]) AK.findall(x -> x > 0, v) # [1, 3, 5] +AK.findall(x -> x > 0, v; items=v) # Int32[5, 8, 3] m = CUDA.CuArray(Bool[1 0; 0 1]) AK.findall(m) # [CartesianIndex(1, 1), CartesianIndex(2, 2)] +AK.findall(m; items=LinearIndices(m)) # [1, 4] ``` """ function findall( - values::AbstractArray, backend::Backend=get_backend(values); - alg::FindallAlgorithm=ScanScatter(), - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - prefer_threads::Bool=true, + values::AbstractArray; + items::AbstractArray=keys(values), + backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_bools::Union{Nothing, AbstractArray}=nothing, ) - alg = findall_algorithm(alg) + values isa AbstractArray{Bool} || _check_bool_result(identity, values) + _check_findall_items(values, items) + backend = _resolve_backend(backend, values, items) + a = _resolve_findall(alg, backend, eltype(values)) bools = if values isa AbstractArray{Bool} isnothing(temp_bools) || throw(ArgumentError("temp_bools is not used for a Bool mask")) values else - findall_bools(identity, values, backend, temp_bools; - max_tasks, min_elems, prefer_threads, block_size=alg.block_size) + findall_bools(identity, values, backend, a, temp_bools) end - I = keytype(values) - output_indices = findall_output_indices(values, I) - findall_impl(bools, I, output_indices, backend; - alg, max_tasks, min_elems, prefer_threads, temp) + findall_impl(bools, eltype(items), items, backend, a; temp) end function findall( - pred, v::AbstractArray, backend::Backend=get_backend(v); - alg::FindallAlgorithm=ScanScatter(), - max_tasks::Int=Threads.nthreads(), - min_elems::Int=1, - prefer_threads::Bool=true, + pred, v::AbstractArray; + items::AbstractArray=keys(v), + backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto(), temp::Union{Nothing, AbstractArray}=nothing, temp_bools::Union{Nothing, AbstractArray}=nothing, ) - alg = findall_algorithm(alg) - bools = findall_bools(pred, v, backend, temp_bools; - max_tasks, min_elems, prefer_threads, block_size=alg.block_size) - I = ndims(v) == 0 ? Int : keytype(v) - output_indices = findall_output_indices(v, I) - findall_impl(bools, I, output_indices, backend; - alg, max_tasks, min_elems, prefer_threads, temp) + _check_bool_result(pred, v) + _check_findall_items(v, items) + backend = _resolve_backend(backend, v, items) + a = _resolve_findall(alg, backend, eltype(v)) + bools = findall_bools(pred, v, backend, a, temp_bools) + findall_impl(bools, eltype(items), items, backend, a; temp) end + +_check_findall_items(v, items) = length(items) == length(v) || throw(DimensionMismatch( + "`items` must have the length of the array, $(length(v)), got $(length(items))")) diff --git a/src/predicates.jl b/src/predicates.jl index f02e480d..598f1db8 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -1,8 +1,96 @@ -abstract type PredicatesAlgorithm end -struct ConcurrentWrite <: PredicatesAlgorithm end -Base.@kwdef struct MapReduce{T <: Union{Nothing, AbstractArray}} <: PredicatesAlgorithm - temp::T = nothing - switch_below::Int = 0 +""" + ConcurrentWrite(; block_size=nothing) + +GPU `any`/`all` in which every thread whose element decides the result stores the same value to +one global flag. It is the fastest predicate algorithm, but some older Intel GPUs have been reported +to hang when many threads write one location, so `Auto` picks [`ViaReduce`](@ref) on oneAPI. +""" +Base.@kwdef struct ConcurrentWrite <: PredicateAlgorithm + block_size::Union{Nothing, Int} = nothing +end + +""" + ViaReduce(; reduce::ReduceAlgorithm=BlockReduce()) + +`any`/`all` as a reduction with `|` or `&`, using the reduction algorithm `reduce`. It runs on +every backend that runs AcceleratedKernels' kernels, but reads every element. +""" +Base.@kwdef struct ViaReduce{R <: ReduceAlgorithm} <: PredicateAlgorithm + reduce::R = BlockReduce() +end + + +""" + PredicateTuning(; kwargs...) + +Values that drive `Auto` selection and fill unset algorithm fields for `any` and `all` on one +device, as returned by [`predicate_tuning`](@ref): + +- `prefer_concurrent_write`: whether `Auto` picks `ConcurrentWrite`, rather than `ViaReduce()`. +- `block_size`: for `ConcurrentWrite`. +- `threads_min_elems`: the default `min_elems` of `CPUThreads.Partitioned`. + +The defaults reproduce AK's historical settings. Internal: the fields may change in any release. +""" +Base.@kwdef struct PredicateTuning + prefer_concurrent_write::Bool = true + block_size::Int = 256 + threads_min_elems::Int = 1 +end + +""" + predicate_tuning(backend, T) -> PredicateTuning + +The `any`/`all` tuning for element type `T` (the input's) on `backend`'s current device; see +[`sort_tuning`](@ref) for the conventions. +""" +predicate_tuning(::Backend, ::Type) = PredicateTuning() + +""" + _resolve_predicate(alg, backend, T) -> Algorithm + +Resolve `alg` for `any`/`all` over elements of type `T` on `backend`: `ConcurrentWrite`, +`ViaReduce` (with its reduction resolved) or `CPUThreads.Partitioned`, with every field set, or an +`ArgumentError`. +""" +function _resolve_predicate(alg::Algorithm, backend::Backend, ::Type{T}) where {T} + _checkdomain(alg) + t = predicate_tuning(backend, T) + a = alg isa Auto ? _select_predicate(backend, t) : alg + a = _fill(a, t, backend) + _check_predicate(a, backend) + return a +end + +function _select_predicate(backend, t::PredicateTuning) + _runs_threads(backend) && return CPUThreads.Partitioned() + return t.prefer_concurrent_write ? ConcurrentWrite() : ViaReduce() +end + +function _checkdomain(a::ConcurrentWrite) + _check_pow2(a, :block_size) + a.block_size === nothing || a.block_size <= 1024 || throw(ArgumentError( + "ConcurrentWrite: `block_size` must be at most 1024, got $(a.block_size)")) + nothing +end + +_fill(a::ConcurrentWrite, t::PredicateTuning, backend) = + ConcurrentWrite(something(a.block_size, t.block_size)) +# The nested reduction reduces `Bool`s along the whole array +_fill(a::ViaReduce, t::PredicateTuning, backend) = + ViaReduce(_resolve_reduce(a.reduce, backend, Bool, :)) +_fill(a::CPUThreads.Partitioned, t::PredicateTuning, backend) = _fill_threads(a, t) +_fill(a::Algorithm, t::PredicateTuning, backend) = + throw(ArgumentError("$(_algname(a)) is not an algorithm for `any` and `all`")) + +function _check_predicate(a::ConcurrentWrite, backend) + _checkdomain(a) + _require_kernels(a, backend) +end +_check_predicate(a::ViaReduce, backend) = nothing # its reduction was resolved by `_fill` +function _check_predicate(a::CPUThreads.Partitioned, backend) + _checkdomain(a) + _require_threads(a, backend) end @@ -30,44 +118,23 @@ end """ - any( - pred, v::AbstractArray, backend::Backend=get_backend(v); - - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - ) + any(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) -Check if any element of `v` satisfies the predicate `pred` (i.e. some `pred(v[i]) == true`). -Optimised differently to `mapreduce` due to shortcircuiting behaviour of booleans. +Check if any element of `v` satisfies the predicate `pred` (i.e. some `pred(v[i]) == true`); `pred` +must return a `Bool`. Optimised differently to `mapreduce` due to shortcircuiting behaviour of +booleans. **Other names**: not often implemented standalone on GPUs, typically included as part of a reduction. -## CPU -Multithreaded parallelisation is only worth it for large arrays, relatively expensive predicates, -and/or rare occurrence of true; use `max_tasks` and `min_elems` to only use parallelism when worth -it in your application. When only one thread is needed, there is no overhead. - -## GPU -There are two possible `alg` choices: -- `ConcurrentWrite()`: the default algorithm, using concurrent writing to a global flag; there is - only one platform we are aware of (Intel UHD 620 integrated graphics cards) where multiple - threads writing to the same memory location - even if writing the same value - hang the device. -- `MapReduce(; temp=nothing, switch_below=0)`: a conservative [`mapreduce`](@ref)-based - implementation which can be used on all platforms, but does not use shortcircuiting - optimisations. You can set the `temp` and `switch_below` keyword arguments to be forwarded to - [`mapreduce`](@ref). - -# Platform-Specific Notes -On oneAPI, `alg=MapReduce()` is the default as on some Intel GPUs concurrent global writes hang -the device. +`alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref +AcceleratedKernels.CPUThreads.Partitioned) on the host, and [`ConcurrentWrite`](@ref) on GPUs +([`ViaReduce`](@ref) on oneAPI). `backend` is derived from `v`; pass it for inputs that do not +determine it, such as index ranges. + +On the host, multithreaded parallelisation is only worth it for large arrays, relatively expensive +predicates, and/or rare occurrence of true; use `CPUThreads.Partitioned(; max_tasks, min_elems)` +to only use parallelism when worth it in your application. # Examples ```julia @@ -76,17 +143,13 @@ using CUDA v = CuArray(rand(Float32, 100_000)) AK.any(x -> x < 1, v) -``` - -Using a different algorithm: -```julia -AK.any(x -> x < 1, v, alg=AK.MapReduce(switch_below=100)) +AK.any(x -> x < 1, v; alg=AK.ViaReduce(AK.BlockReduce(switch_below=100))) ``` Checking a more complex condition with unmaterialised index ranges: ```julia function complex_any(x, y) - AK.any(eachindex(x), AK.get_backend(x)) do i + AK.any(eachindex(x); backend=AK.get_backend(x)) do i x[i] < 0 && y[i] > 0 end end @@ -94,112 +157,56 @@ end complex_any(CuArray(rand(Float32, 100)), CuArray(rand(Float32, 100))) ``` """ -function any( - pred, v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - _any_impl( - pred, v, backend; - kwargs... - ) +function any(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto()) + _check_bool_result(pred, v) + backend = _resolve_backend(backend, v) + _any(pred, v, backend, _resolve_predicate(alg, backend, eltype(v))) end +# The predicate must return a `Bool`, as in Base. Where inference shows it cannot, the call fails +# on the host, before launching, with an `ArgumentError` instead of an exception from the kernel +# (which GPU backends report without Base's `TypeError`). +function _check_bool_result(pred, v) + isempty(v) && return nothing # (Base never calls the predicate then) + R = Base.promote_op(pred, eltype(v)) + R === Union{} || Bool <: R || throw(ArgumentError( + "the predicate must return a `Bool`, but returns `$R` for elements of type `$(eltype(v))`")) + nothing +end -function _any_impl( - pred, v::AbstractArray, backend::Backend; - - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int=256, -) - if use_gpu_algorithm(backend, prefer_threads) - @argcheck block_size > 0 - - # Some platforms crash when multiple threads write to the same memory location in a global - # array (e.g. old Intel Graphics); if it is the same value, it is well-defined on others (e.g. - # CUDA). If not cooperative, we need to do a mapreduce - if alg === ConcurrentWrite() - out = KernelAbstractions.zeros(backend, Int8, 1) - _any_global!(backend, block_size)(out, pred, v, ndrange=length(v)) - outh = @allowscalar(out[1]) - return outh == 0 ? false : true - else - return mapreduce( - pred, - (x, y) -> x || y, - v; - backend, - init=false, - neutral=false, - alg=BlockReduce(; block_size, switch_below=alg.switch_below), - temp=alg.temp, - ) - end - else - overall = Ref(false) - task_partition(length(v), max_tasks, min_elems) do irange - for i in irange - if pred(v[i]) - # Again, this is technically a thread race, but it doesn't matter as all threads - # would write the same value; no data corruption can occur - overall[] = true - break - end +function _any(pred, v, backend, alg::ConcurrentWrite) + # `ndrange` must not be zero + isempty(v) && return false + out = KernelAbstractions.zeros(backend, Int8, 1) + _any_global!(backend, alg.block_size)(out, pred, v, ndrange=length(v)) + return @allowscalar(out[1]) != 0 +end + +_any(pred, v, backend, alg::ViaReduce) = + mapreduce(_BoolValued(pred), |, v; backend, init=false, alg=alg.reduce) + +function _any(pred, v, backend, alg::CPUThreads.Partitioned) + overall = Ref(false) + task_partition(length(v), alg.max_tasks, alg.min_elems) do irange + for i in irange + if pred(v[i]) + # Again, this is technically a thread race, but it doesn't matter as all threads + # would write the same value; no data corruption can occur + overall[] = true + break end end - return overall[] end + return overall[] end - - """ - all( - pred, v::AbstractArray, backend::Backend=get_backend(v); - - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), + all(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - ) - -Check if all elements of `v` satisfy the predicate `pred` (i.e. all `pred(v[i]) == true`). -Optimised differently to `mapreduce` due to shortcircuiting behaviour of booleans. - -**Other names**: not often implemented standalone on GPUs, typically included as part of a -reduction. - -## CPU -Multithreaded parallelisation is only worth it for large arrays, relatively expensive predicates, -and/or rare occurrence of true; use `max_tasks` and `min_elems` to only use parallelism when worth -it in your application. When only one thread is needed, there is no overhead. - -## GPU -There are two possible `alg` choices: -- `ConcurrentWrite()`: the default algorithm, using concurrent writing to a global flag; there is - only one platform we are aware of (Intel UHD 620 integrated graphics cards) where multiple - threads writing to the same memory location - even if writing the same value - hang the device. -- `MapReduce(; temp=nothing, switch_below=0)`: a conservative [`mapreduce`](@ref)-based - implementation which can be used on all platforms, but does not use shortcircuiting - optimisations. You can set the `temp` and `switch_below` keyword arguments to be forwarded to - [`mapreduce`](@ref). - -# Platform-Specific Notes -On oneAPI, `alg=MapReduce()` is the default as on some Intel GPUs concurrent global writes hang -the device. +Check if all elements of `v` satisfy the predicate `pred` (i.e. all `pred(v[i]) == true`); `pred` +must return a `Bool`. The keywords are those of [`any`](@ref). # Examples ```julia @@ -210,80 +217,25 @@ v = MtlArray(rand(Float32, 100_000)) AK.all(x -> x > 0, v) ``` -Using a different algorithm: -```julia -AK.all(x -> x > 0, v, alg=AK.MapReduce(switch_below=100)) -``` - Checking a more complex condition with unmaterialised index ranges: ```julia function complex_all(x, y) - AK.all(eachindex(x), AK.get_backend(x)) do i + AK.all(eachindex(x); backend=AK.get_backend(x)) do i x[i] > 0 && y[i] < 0 end end -complex_all(CuArray(rand(Float32, 100)), CuArray(rand(Float32, 100))) +complex_all(MtlArray(rand(Float32, 100)), MtlArray(rand(Float32, 100))) ``` """ -function all( - pred, v::AbstractArray, backend::Backend=get_backend(v); - kwargs... -) - _all_impl( - pred, v, backend; - kwargs..., - ) +function all(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + alg::Algorithm=Auto()) + _check_bool_result(pred, v) + backend = _resolve_backend(backend, v) + _all(pred, v, backend, _resolve_predicate(alg, backend, eltype(v))) end - -function _all_impl( - pred, v::AbstractArray, backend::Backend; - - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, - - # GPU settings - block_size::Int=256, -) - if use_gpu_algorithm(backend, prefer_threads) - @argcheck block_size > 0 - - # Some platforms crash when multiple threads write to the same memory location in a global - # array (e.g. old Intel Graphics); if it is the same value, it is well-defined on others (e.g. - # CUDA). If not cooperative, we need to do a mapreduce - if alg === ConcurrentWrite() - out = KernelAbstractions.zeros(backend, Int8, 1) - _any_global!(backend, block_size)(out, (!pred), v, ndrange=length(v)) - outh = @allowscalar(out[1]) - return outh == 0 ? true : false - else - return mapreduce( - pred, - (x, y) -> x && y, - v; - backend, - init=true, - neutral=true, - alg=BlockReduce(; block_size, switch_below=alg.switch_below), - temp=alg.temp, - ) - end - else - overall = Ref(true) - task_partition(length(v), max_tasks, min_elems) do irange - for i in irange - if !pred(v[i]) - overall[] = false - break - end - end - end - return overall[] - end -end +_all(pred, v, backend, alg::ConcurrentWrite) = !_any(!pred, v, backend, alg) +_all(pred, v, backend, alg::ViaReduce) = + mapreduce(_BoolValued(pred), &, v; backend, init=true, alg=alg.reduce) +_all(pred, v, backend, alg::CPUThreads.Partitioned) = !_any(!pred, v, backend, alg) diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index 47e62112..d7cf0b6a 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -1,3 +1,9 @@ +# `f`, required to return a `Bool` (Base's `_bool`), for `count` and the predicates +struct _BoolValued{F} <: Function + f::F +end +@inline (c::_BoolValued)(x) = c.f(x)::Bool + # The reduction contract's machinery (see `mapreducedim!`): the accumulator type, the value that # seeds partial results, and how a finished partial becomes an output value. diff --git a/src/utils.jl b/src/utils.jl index bd43bace..82a2a3e4 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -2,12 +2,6 @@ function ispow2(x) count_ones(x) == 1 end -# Helper function to check whether the package cpu implementation of an algorithm should be used -const CPU_BACKEND = get_backend([]) -@inline function use_gpu_algorithm(backend, prefer_threads) - return backend != CPU_BACKEND || !prefer_threads -end - # Local memory that kernels with a tunable footprint may use: the minimum an OpenCL device must # provide (`CL_DEVICE_LOCAL_MEM_SIZE`, full profile), and Metal's threadgroup memory limit. const LOCAL_MEMORY_BUDGET = 32 * 1024 diff --git a/test/generic/findall.jl b/test/generic/findall.jl index 3adbdb25..1f92526e 100644 --- a/test/generic/findall.jl +++ b/test/generic/findall.jl @@ -19,11 +19,20 @@ Base.similar(v::FindallOffsetVector, ::Type{T}) where T = Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v.data, T, dims) +# Tests that do not choose an algorithm use `FINDALL_ALG`: `Auto()`, except in the `--cpu-ka` +# configuration, whose point is to run AK's kernels on the host backend. `findall_alg` builds an +# explicitly tuned algorithm for the configuration from both kinds of settings. +FINDALL_ALG = HOST_KERNELS ? AK.ScanScatter() : AK.Auto() +findall_alg(; block_size=nothing, items_per_thread=nothing, max_tasks=nothing, min_elems=nothing) = + TEST_KERNELS ? AK.ScanScatter(; block_size, items_per_thread) : + AK.CPUThreads.Partitioned(; max_tasks, min_elems) + + @testset "findall" begin Random.seed!(0) - default_alg = AK.ScanScatter() - tile_size = default_alg.block_size * default_alg.items_per_thread + default_tuning = AK.FindallTuning() + tile_size = default_tuning.block_size * default_tuning.items_per_thread edge_sizes = [0, 1, 2, 3, tile_size - 1, tile_size, tile_size + 1, 2tile_size - 1, 2tile_size, 2tile_size + 1, 10_000] test_types = valid_backend_eltypes(BACKEND, (Int32, Float32, Float64)) @@ -33,20 +42,20 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. pred = T <: AbstractFloat ? (x -> x > oftype(x, 0.5)) : (x -> x > zero(x)) h = T <: AbstractFloat ? rand(T, n) : rand(T(-5):T(5), n) v = array_from_host(h) - @test Array(AK.findall(pred, v; prefer_threads)) == findall(pred, h) + @test Array(AK.findall(pred, v; alg=FINDALL_ALG)) == findall(pred, h) end h = collect(Int32, -10:10) v = array_from_host(h) - @test Array(AK.findall(FindallCallable(), v; prefer_threads)) == findall(x -> x > 0, h) + @test Array(AK.findall(FindallCallable(), v; alg=FINDALL_ALG)) == findall(x -> x > 0, h) - if prefer_threads + if !TEST_KERNELS calls = Ref(0) pred = x -> (calls[] += 1; isodd(x)) h = collect(1:100) - @test AK.findall(pred, h; max_tasks=1) == findall(isodd, h) + @test AK.findall(pred, h; alg=AK.CPUThreads.Partitioned(max_tasks=1)) == findall(isodd, h) @test calls[] == length(h) - @test_throws TypeError AK.findall(Returns(1), [1]; max_tasks=1) + @test_throws ArgumentError AK.findall(Returns(1), [1]; alg=AK.CPUThreads.Partitioned(max_tasks=1)) end end @@ -54,17 +63,18 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. for n in edge_sizes h = rand(Bool, n) v = array_from_host(h) - out = AK.findall(v; prefer_threads) + out = AK.findall(v; alg=FINDALL_ALG) @test Array(out) == findall(h) @test eltype(out) == Int end - if prefer_threads + if !TEST_KERNELS @test AK.findall(Any[true, false, true]) == findall(Any[true, false, true]) @test AK.findall(Any[true, false, true]; temp_bools=Vector{Bool}(undef, 3)) == findall(Any[true, false, true]) - @test_throws TypeError AK.findall([1]) - @test_throws TypeError AK.findall(Any[true, missing]; max_tasks=1) + @test_throws ArgumentError AK.findall([1]) + @test_throws TypeError AK.findall(Any[true, missing]; + alg=AK.CPUThreads.Partitioned(max_tasks=1)) scalar = Array{Any}(undef) scalar[] = true @@ -76,31 +86,38 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. for shape in ([4, 2], [1, 6], [64, 64], [8, 8, 8]) h = rand(Float32, shape...) v = array_from_host(h) - out = AK.findall(x -> x > 0.5f0, v; prefer_threads) + out = AK.findall(x -> x > 0.5f0, v; alg=FINDALL_ALG) @test Array(out) == findall(x -> x > 0.5f0, h) @test eltype(out) == CartesianIndex{length(shape)} hb = rand(Bool, shape...) - @test Array(AK.findall(array_from_host(hb); prefer_threads)) == findall(hb) + @test Array(AK.findall(array_from_host(hb); alg=FINDALL_ALG)) == findall(hb) end for value in (false, true) h = fill(value) - @test Array(AK.findall(array_from_host(h); prefer_threads)) == findall(h) + @test Array(AK.findall(array_from_host(h); alg=FINDALL_ALG)) == findall(h) end + # The default items are `keys(A)`, also for a 0-dimensional array and the predicate form for value in (0.25f0, 0.75f0) h = fill(value) pred = x -> x > 0.5f0 - @test Array(AK.findall(pred, array_from_host(h); prefer_threads)) == findall(pred, h) + out = AK.findall(pred, array_from_host(h); alg=FINDALL_ALG) + @test eltype(out) === CartesianIndex{0} && Array(out) == findall(fill(pred(h[]))) + out = AK.findall(pred, array_from_host(h); items=LinearIndices(h), alg=FINDALL_ALG) + @test eltype(out) === Int && Array(out) == (value > 0.5f0 ? [1] : Int[]) end - if prefer_threads + if !TEST_KERNELS + # Positions are ordinal, so offset axes pair with any items h = FindallOffsetVector([-1, 1, -2, 2, 0], -3) mask = FindallOffsetVector(Bool[false, true, true, false, true], -3) for max_tasks in (1, 4) - @test AK.findall(x -> x > 0, h, BACKEND; max_tasks, min_elems=1) == - [-1, 1] - @test AK.findall(mask, BACKEND; max_tasks, min_elems=1) == findall(mask) + alg = AK.CPUThreads.Partitioned(; max_tasks, min_elems=1) + @test AK.findall(x -> x > 0, h; backend=BACKEND, alg) == [-1, 1] + @test AK.findall(mask; backend=BACKEND, alg) == findall(mask) + @test AK.findall(x -> x > 0, h; items=10:14, backend=BACKEND, alg) == [11, 13] + @test AK.findall(mask; items=h, backend=BACKEND, alg) == [1, -2, 0] end h = collect(1:20) @@ -111,15 +128,56 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. end end + @testset "items" begin + h = rand(Float32, 37, 29) + v = array_from_host(h) + pred = x -> x > 0.5f0 + sel = vec(pred.(h)) + # Linear indices of a matrix, and values + @test Array(AK.findall(pred, v; items=LinearIndices(v), alg=FINDALL_ALG)) == + findall(sel) + out = AK.findall(pred, v; items=v, alg=FINDALL_ALG) + @test eltype(out) === Float32 && Array(out) == h[sel] + hb = rand(Bool, size(h)) + @test Array(AK.findall(array_from_host(hb); items=v, alg=FINDALL_ALG)) == h[hb] + @test Array(AK.findall(array_from_host(hb); items=CartesianIndices(v), alg=FINDALL_ALG)) == + findall(hb) + # ... of another shape or on another array, paired by position + w = array_from_host(collect(Int32, 1:length(h))) + @test Array(AK.findall(pred, v; items=w, alg=FINDALL_ALG)) == findall(sel) + @test Array(AK.findall(isodd, 1:10; items=array_from_host(collect(Int32, 11:20)), + alg=FINDALL_ALG)) == 11:2:19 + # A wrapper whose `eachindex` is Cartesian, with its own linear indices + hv = view(h, 1:2:37, :) + vv = view(v, 1:2:37, :) + @test Array(AK.findall(pred, vv; items=LinearIndices(vv), alg=FINDALL_ALG)) == + findall(vec(pred.(hv))) + @test Array(AK.findall(pred, vv; alg=FINDALL_ALG)) == findall(pred, hv) + # A mask whose wrappers `@Const` rebuilds on the device (a reshaped view) + hbm = rand(Bool, 8, 8) + vbm = vec(view(array_from_host(hbm), 1:4, 1:4)) + @test Array(AK.findall(vbm; alg=FINDALL_ALG)) == findall(vec(view(hbm, 1:4, 1:4))) + # `items` must have the array's length + @test_throws DimensionMismatch AK.findall(pred, v; items=1:3, alg=FINDALL_ALG) + @test_throws DimensionMismatch AK.findall(array_from_host(hb); items=1:3, alg=FINDALL_ALG) + end + @testset "selection extremes" begin for n in (0, 1, 2, 1000), h in (trues(n), falses(n)) values = collect(h) - @test Array(AK.findall(array_from_host(values); prefer_threads)) == findall(values) + @test Array(AK.findall(array_from_host(values); alg=FINDALL_ALG)) == findall(values) end + # A range, with its backend given explicitly + @test Array(AK.findall(isodd, 1:10; backend=BACKEND, alg=FINDALL_ALG)) == findall(isodd, 1:10) + # (indexing a `Bool` range may throw an `InexactError`, which on oneAPI needs oneAPI.jl + # 2.9.2's device heap) + @test Array(AK.findall(false:true; backend=BACKEND, alg=FINDALL_ALG)) == findall(false:true) + @test Array(AK.findall(true:false; backend=BACKEND, alg=FINDALL_ALG)) == Int[] + v = array_from_host(collect(Int32, 1:1000)) - @test Array(AK.findall(x -> x > 0, v; prefer_threads)) == collect(1:1000) - @test Array(AK.findall(x -> x < 0, v; prefer_threads)) == Int[] + @test Array(AK.findall(x -> x > 0, v; alg=FINDALL_ALG)) == collect(1:1000) + @test Array(AK.findall(x -> x < 0, v; alg=FINDALL_ALG)) == Int[] end @testset "random sizes" begin @@ -127,7 +185,7 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. n = rand(1:100_000) h = rand(Float32, n) v = array_from_host(h) - @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads)) == + @test Array(AK.findall(x -> x > 0.5f0, v; alg=FINDALL_ALG)) == findall(x -> x > 0.5f0, h) end end @@ -136,39 +194,72 @@ Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v. h = rand(Float32, 10_000) v = array_from_host(h) for block_size in (32, 64, 128, 256), items_per_thread in (1, 3, 8) - alg = AK.ScanScatter(; block_size, items_per_thread) - @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads, alg)) == + alg = findall_alg(; block_size, items_per_thread) + @test Array(AK.findall(x -> x > 0.5f0, v; alg)) == findall(x -> x > 0.5f0, h) end for (max_tasks, min_elems) in ((1, 1), (2, 100), (4, 1000)) - @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads, max_tasks, min_elems)) == + @test Array(AK.findall(x -> x > 0.5f0, v; alg=findall_alg(; max_tasks, min_elems))) == findall(x -> x > 0.5f0, h) end - alg = AK.ScanScatter(block_size=64, items_per_thread=3) - temp = similar(v, Int, max(4, cld(length(v), alg.block_size * alg.items_per_thread))) + alg = findall_alg(block_size=64, items_per_thread=3, max_tasks=4) + temp = similar(v, Int, max(4, cld(length(v), 64 * 3))) temp_bools = similar(v, Bool) - @test Array(AK.findall(x -> x > 0.5f0, v; - prefer_threads, max_tasks=4, alg, temp, temp_bools)) == + @test Array(AK.findall(x -> x > 0.5f0, v; alg, temp, temp_bools)) == findall(x -> x > 0.5f0, h) - @test_throws ArgumentError AK.findall(v; prefer_threads, alg=OtherFindallAlgorithm()) + @test_throws ArgumentError AK.findall(v; alg=OtherFindallAlgorithm()) @test_throws ArgumentError AK.findall(identity, temp_bools; - prefer_threads, temp_bools) - @test_throws ArgumentError AK.findall(identity, v; prefer_threads, - temp_bools=reshape(similar(v, Bool), :, 1)) + temp_bools) + @test_throws ArgumentError AK.findall(identity, v; temp_bools=reshape(similar(v, Bool), :, 1)) - if !prefer_threads + if TEST_KERNELS bools = array_from_host(rand(Bool, length(v))) - @test_throws ArgumentError AK.findall(bools; prefer_threads, - alg=AK.ScanScatter(block_size=192)) - @test_throws ArgumentError AK.findall(bools; prefer_threads, - alg=AK.ScanScatter(items_per_thread=0)) - @test_throws ArgumentError AK.findall(bools; prefer_threads, + @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(block_size=192)) + @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(items_per_thread=0)) + @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(), temp=similar(v, Int32, 100)) - @test_throws ArgumentError AK.findall(bools; prefer_threads, - temp=similar(v, Int, 1)) + @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(), temp=similar(v, Int, 1)) end end end + + +# A GPU backend with its own tuning +struct FindallResolveTestBackend <: KernelAbstractions.Backend end +AK.findall_tuning(::FindallResolveTestBackend, ::Type) = + AK.FindallTuning(block_size=128, items_per_thread=4) + +@testset "findall resolution" begin + B = FindallResolveTestBackend() + @test AK._resolve_findall(AK.Auto(), B, Bool) === AK.ScanScatter(128, 4) + @test AK._resolve_findall(AK.ScanScatter(block_size=64), B, Bool) === AK.ScanScatter(64, 4) + @test AK._resolve_findall(AK.Auto(), AK.HOST_BACKEND, Bool) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 1) + for bad in (AK.ScanScatter(block_size=0), AK.ScanScatter(block_size=96), + AK.ScanScatter(block_size=2048), AK.ScanScatter(items_per_thread=0), + AK.ScanScatter(block_size=1024, items_per_thread=1 << 22)) + @test_throws ArgumentError AK._resolve_findall(bad, B, Bool) + end + @test_throws ArgumentError AK._resolve_findall(AK.CPUThreads.Partitioned(), B, Bool) + @test_throws ArgumentError AK._resolve_findall(AK.BlockReduce(), B, Bool) + + # Removed keywords + v = array_from_host(rand(Bool, 10)) + @test_throws MethodError AK.findall(v; prefer_threads=true) + @test_throws MethodError AK.findall(v; max_tasks=2) + @test_throws MethodError AK.any(identity, v; prefer_threads=true) + @test_throws MethodError AK.all(identity, v; block_size=64) +end + + +@testset "findall: Bool values" begin + # The predicate (or the values) must give a `Bool`, as in Base; where inference shows they + # cannot, before launching (not every backend reports an error thrown in a kernel) + @test_throws ArgumentError AK.findall(x -> 1, array_from_host(Int32[1, 2])) + @test_throws ArgumentError AK.findall(array_from_host(Int32[1, 0])) + # ... except on an empty array, where Base never calls it + @test isempty(AK.findall(x -> 1, array_from_host(Int32[]))) +end diff --git a/test/generic/predicates.jl b/test/generic/predicates.jl index 16bd7b58..a6359b05 100644 --- a/test/generic/predicates.jl +++ b/test/generic/predicates.jl @@ -1,38 +1,100 @@ -@testset "truth" begin +# The algorithms under test: the kernel algorithms the backend supports, or the threaded host +# algorithm, each with default and explicit settings; and `Auto()`. +PRED_ALGS = if TEST_KERNELS + AK.Algorithm[ + AK.ConcurrentWrite(), AK.ConcurrentWrite(block_size=64), + AK.ViaReduce(), AK.ViaReduce(AK.BlockReduce(block_size=64, switch_below=100))] +else + AK.Algorithm[AK.CPUThreads.Partitioned(), AK.CPUThreads.Partitioned(max_tasks=2, min_elems=100)] +end +HOST_KERNELS || pushfirst!(PRED_ALGS, AK.Auto()) + +@testset "truth $alg" for alg in PRED_ALGS Random.seed!(0) # Simple correctness tests v = array_from_host(1:100) - @test AK.any(x->x<0, v; prefer_threads) === false - @test AK.any(x->x>99, v; prefer_threads) === true + @test AK.any(x->x<0, v; alg) === false + @test AK.any(x->x>99, v; alg) === true - @test AK.all(x->x>0, v; prefer_threads) === true - @test AK.all(x->x<100, v; prefer_threads) === false + @test AK.all(x->x>0, v; alg) === true + @test AK.all(x->x<100, v; alg) === false for _ in 1:100 num_elems = rand(1:100_000) v = array_from_host(rand(Float32, num_elems)) - @test AK.any(x->x<0, v; prefer_threads) === false - @test AK.any(x->x<1, v; prefer_threads) === true - @test AK.all(x->x<1, v; prefer_threads) === true - @test AK.all(x->x<0, v; prefer_threads) === false + @test AK.any(x->x<0, v; alg) === false + @test AK.any(x->x<1, v; alg) === true + @test AK.all(x->x<1, v; alg) === true + @test AK.all(x->x<0, v; alg) === false end - for _ in 1:100 - num_elems = rand(1:100_000) - v = array_from_host(rand(Float32, num_elems)) - @test AK.any(x->x<0, v; prefer_threads) === false - @test AK.any(x->x<1, v; prefer_threads) === true - @test AK.all(x->x<1, v; prefer_threads) === true - @test AK.all(x->x<0, v; prefer_threads) === false + # Empty inputs, as Base + e = array_from_host(Float32[]) + @test AK.any(x->x<0, e; alg) === false + @test AK.all(x->x<0, e; alg) === true + + # Unmaterialised index ranges, which need `backend` + x = array_from_host(rand(Float32, 1000)) + @test AK.any(i -> x[i] > 2, 1:length(x); backend=BACKEND, alg) === false + @test AK.all(i -> x[i] < 2, 1:length(x); backend=BACKEND, alg) === true + + # The predicate must return a Bool, as in Base + @test_throws ArgumentError AK.any(identity, array_from_host(Int32[2, 3]); alg) + @test_throws ArgumentError AK.all(identity, array_from_host(Int32[2, 3]); alg) + # ... which it never calls on an empty array + @test AK.any(identity, array_from_host(Int32[]); alg) === false + @test AK.all(identity, array_from_host(Int32[]); alg) === true + + # A reshaped view, whose wrappers `@Const` rebuilds on the device + hr = rand(Int32(0):Int32(9), 50, 40) + vr = vec(view(array_from_host(hr), 1:40, 1:30)) + @test AK.any(x -> x > 8, vr; alg) === any(x -> x > 8, vec(view(hr, 1:40, 1:30))) + @test AK.all(x -> x < 10, vr; alg) === true + + # A source of a bits-union element type (where the backend's arrays can hold one) + if TEST_BITS_UNIONS + hm = rand([true, false, missing], 10_000) + dm = array_from_host(hm) + @test AK.any(ismissing, dm; alg) === true + @test AK.all(x -> x !== true, dm; alg) === false + @test AK.any(x -> x === true, array_from_host(Union{Missing, Bool}[missing, false]); alg) === + false end - # Testing different settings - v = array_from_host(rand(-5:5, 100_000)) - AK.any(x->x<5, v; prefer_threads, max_tasks=2, min_elems=100, block_size=64) - AK.all(x->x<5, v; prefer_threads, max_tasks=2, min_elems=100, block_size=64) + # Test that undefined kwargs are not accepted + @test_throws MethodError AK.any(x->x<0, v; alg, bad=:kwarg) +end + + +# GPU backends whose tuning prefers `ViaReduce` and `ConcurrentWrite` +struct PredicateResolveTestBackend <: KernelAbstractions.Backend end +AK.predicate_tuning(::PredicateResolveTestBackend, ::Type) = + AK.PredicateTuning(prefer_concurrent_write=false) +struct PredicateCWTestBackend <: KernelAbstractions.Backend end - # The MapReduce algorithm is tested with the reductions tests +@testset "predicate resolution" begin + B = PredicateResolveTestBackend() + host = AK.HOST_BACKEND + @test AK._resolve_predicate(AK.Auto(), host, Float32) === + AK.CPUThreads.Partitioned(Threads.nthreads(), 1) + @test AK._resolve_predicate(AK.Auto(), PredicateCWTestBackend(), Float32) === + AK.ConcurrentWrite(256) + @test AK._resolve_predicate(AK.Auto(), B, Float32) === AK.ViaReduce(AK.BlockReduce(256, 2, 0)) + # The nested reduction is resolved like any reduction + @test AK._resolve_predicate(AK.ViaReduce(AK.BlockReduce(block_size=64)), B, Float32) === + AK.ViaReduce(AK.BlockReduce(64, 2, 0)) + @test_throws ArgumentError AK._resolve_predicate(AK.ViaReduce(AK.BlockReduce(block_size=3)), B, Float32) + # A tuning only guides `Auto` + @test AK._resolve_predicate(AK.ConcurrentWrite(), B, Float32) === AK.ConcurrentWrite(256) + for bad in (AK.ConcurrentWrite(block_size=0), AK.ConcurrentWrite(block_size=96), + AK.ConcurrentWrite(block_size=2048)) + @test_throws ArgumentError AK._resolve_predicate(bad, PredicateCWTestBackend(), Float32) + end + # Capabilities and other families + @test_throws ArgumentError AK._resolve_predicate(AK.CPUThreads.Partitioned(), B, Float32) + @test_throws ArgumentError AK._resolve_predicate(AK.BlockReduce(), B, Float32) + @test_throws ArgumentError AK._resolve_predicate(AK.ScanScatter(), B, Float32) end diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index 7e01bd5d..cb7a081c 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1053,18 +1053,6 @@ end # The other settings are stress-tested in reduce end -@testset "MapReduce predicates" begin - for _ in 1:100 - num_elems = rand(1:100_000) - v = array_from_host(rand(Float32, num_elems)) - alg=AK.MapReduce(temp=similar(v, Bool), switch_below=100) - @test AK.any(x->x<0, v; prefer_threads, alg) === false - @test AK.any(x->x<1, v; prefer_threads, alg) === true - @test AK.all(x->x<1, v; prefer_threads, alg) === true - @test AK.all(x->x<0, v; prefer_threads, alg) === false - end -end - @testset "reduce block sizes" begin # The block-level tree reduction must cover every power-of-two group size the device allows; diff --git a/test/runtests.jl b/test/runtests.jl index 27a48107..e30eea89 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -27,7 +27,7 @@ const _testset_setup_code = quote global array_from_host array_from_host(h_arr::AbstractArray, dtype=nothing) = array_from_host(BACKEND, h_arr, dtype) function array_from_host(backend, h_arr::AbstractArray, dtype=nothing) - d_arr = d_arr = if prefer_threads # Don't use KA zeros if not using KA algorithms + d_arr = if !TEST_KERNELS # host arrays for the threaded algorithms zeros(isnothing(dtype) ? eltype(h_arr) : dtype, size(h_arr)) else KernelAbstractions.zeros(backend, isnothing(dtype) ? eltype(h_arr) : dtype, size(h_arr)) @@ -56,7 +56,6 @@ if in_test_env("CUDACore") using CUDACore global BACKEND = CUDABackend() global MAX_BLOCK_SIZE = Int(CUDACore.attribute(CUDACore.device(), CUDACore.DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)) - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = false # ... on the host backend global TEST_DL = true @@ -73,7 +72,6 @@ if in_test_env("AMDGPU") using AMDGPU global BACKEND = ROCBackend() global MAX_BLOCK_SIZE = 1024 - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = false # ... on the host backend global TEST_DL = true @@ -89,7 +87,6 @@ if in_test_env("Metal") using Metal global BACKEND = MetalBackend() global MAX_BLOCK_SIZE = Int(Metal.device().maxThreadsPerThreadgroup.width) - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = false # ... on the host backend global TEST_DL = false @@ -105,7 +102,6 @@ if in_test_env("oneAPI") using oneAPI global BACKEND = oneAPIBackend() global MAX_BLOCK_SIZE = Int(oneAPI.oneL0.compute_properties(oneAPI.device()).maxTotalGroupSize) - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = false # ... on the host backend global TEST_DL = false @@ -122,7 +118,6 @@ if in_test_env("OpenCL") using OpenCL global BACKEND = OpenCLBackend() global MAX_BLOCK_SIZE = Int(OpenCL.cl.device().max_work_group_size) - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = false # ... on the host backend global TEST_DL = false @@ -138,7 +133,6 @@ if args.custom["cpu-ka"] !== nothing push!(backends, "cpu-ka" => quote global BACKEND = get_backend([]) global MAX_BLOCK_SIZE = 1024 - global prefer_threads = false # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = true # AK's kernels run in this configuration global HOST_KERNELS = true # ... on the host backend global TEST_DL = false @@ -151,7 +145,6 @@ if args.custom["cpu"] !== nothing || isempty(backends) push!(backends, "cpu" => quote global BACKEND = get_backend([]) global MAX_BLOCK_SIZE = 1024 - global prefer_threads = true # Also used to determine whether to run the CPU or GPU tests global TEST_KERNELS = false global HOST_KERNELS = false global TEST_DL = false From 0bdbd3f846e71a9df8cad659df8d8e8ff9518205 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 17:46:37 +0200 Subject: [PATCH 8/8] Replace the scratch keywords with workspaces 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. --- docs/make.jl | 1 + docs/src/api/sort.md | 8 +- docs/src/api/workspace.md | 33 ++++ docs/src/performance.md | 9 +- src/AcceleratedKernels.jl | 1 + src/accumulate/accumulate.jl | 126 ++++++++++----- src/accumulate/accumulate_1d_gpu.jl | 30 +--- src/arithmetics.jl | 38 +++-- src/findall.jl | 135 +++++++--------- src/predicates.jl | 63 +++++--- src/reduce/mapreduce_1d_gpu.jl | 34 ++-- src/reduce/mapreduce_nd.jl | 241 +++++++++++++++------------- src/reduce/reduce.jl | 171 +++++++++++++------- src/sort/cpu_sample_sort.jl | 29 ++-- src/sort/merge_sort.jl | 20 +-- src/sort/merge_sort_by_key.jl | 26 +-- src/sort/merge_sortperm.jl | 34 +--- src/sort/radix_sort.jl | 88 ++++++---- src/sort/sort.jl | 191 ++++++++++++++++------ src/workspace.jl | 153 ++++++++++++++++++ test/generic/accumulate.jl | 17 +- test/generic/findall.jl | 22 ++- test/generic/reduce.jl | 17 +- test/generic/reduce_resolve.jl | 2 +- test/generic/sort.jl | 121 +++++--------- test/generic/sort_resolve.jl | 2 +- test/generic/workspace.jl | 240 +++++++++++++++++++++++++++ test/runtests.jl | 5 + 28 files changed, 1229 insertions(+), 628 deletions(-) create mode 100644 docs/src/api/workspace.md create mode 100644 src/workspace.jl create mode 100644 test/generic/workspace.jl diff --git a/docs/make.jl b/docs/make.jl index c332e76f..ebf812e0 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -33,6 +33,7 @@ makedocs(; "Reverse" => "api/reverse.md", "Predicates" => "api/predicates.md", "Arithmetics" => "api/arithmetics.md", + "Scratch Memory" => "api/workspace.md", "Custom Structs" => "api/custom_structs.md", "Task Partitioning" => "api/task_partition.md", "Utilities" => "api/utilities.md", diff --git a/docs/src/api/sort.md b/docs/src/api/sort.md index fb4fe391..b6fedcdd 100644 --- a/docs/src/api/sort.md +++ b/docs/src/api/sort.md @@ -62,10 +62,10 @@ AK.sort!(A; dims=1) # each column sorted ix = AK.sortperm(A; dims=2) # A[ix] has each row sorted ``` -As GPU memory is more expensive, the sorts accept their temporary arrays (`temp`, and `temp_keys` -and `temp_values` for `sort_by_key!`), so that you can supply your own buffers, e.g.: +The sorts' scratch memory (the merge and radix sorts' swap buffers, radix sort's histograms) can +be allocated once and reused with a [workspace](workspace.md): ```julia v = ROCArray(rand(Float32, 100_000)) -temp = similar(v) -AK.sort!(v; temp) +ws = AK.workspace(AK.sort!, v) +AK.sort!(v; workspace=ws) ``` diff --git a/docs/src/api/workspace.md b/docs/src/api/workspace.md new file mode 100644 index 00000000..48631658 --- /dev/null +++ b/docs/src/api/workspace.md @@ -0,0 +1,33 @@ +### Scratch Memory + +Many operations need scratch memory: merge and radix sorts swap between two buffers, reductions +keep partial results per block, scans keep the totals of their tiles, and `findall` counts the +selected elements of each block. By default each call allocates what it needs. To allocate it once +and reuse it, make a workspace for the call and pass it through the operation's `workspace` +keyword: + +```julia +import AcceleratedKernels as AK +using CUDA + +v = CuArray(rand(Float32, 1_000_000)) +ws = AK.workspace(AK.sort!, v) # the same arguments as the call +for _ in 1:10 + rand!(v) + AK.sort!(v; workspace=ws) # allocates no scratch memory +end + +AK.workspace_size(AK.sum, v) # what the call needs, without allocating it +``` + +A workspace belongs to one kind of call: the operation resolves its algorithm and computes its +buffers as it would without one, and a workspace made for another backend or device, another +algorithm (`Auto()` may choose differently for another length) or other buffer sizes is an +`ArgumentError`. It holds no state between calls, so calls may reuse it in turn, but not at the +same time on several tasks or streams. + +```@docs +AcceleratedKernels.workspace +AcceleratedKernels.workspace_size +AcceleratedKernels.Workspace +``` diff --git a/docs/src/performance.md b/docs/src/performance.md index 37c387eb..3cc4e999 100644 --- a/docs/src/performance.md +++ b/docs/src/performance.md @@ -16,13 +16,14 @@ Similarly, for performance on the CPU the overhead of spawning threads should be ``` -### Temporary Arrays +### Scratch Memory -As GPU memory is more expensive, all functions in AcceleratedKernels.jl expose any temporary arrays they will use (the `temp` argument); you can supply your own buffers to make the algorithms not allocate additional GPU storage, e.g.: +As GPU memory is more expensive, every operation that needs scratch memory can take it from a +workspace, allocated once for a kind of call and reused (see [Scratch Memory](api/workspace.md)): ```julia v = ROCArray(rand(Float32, 100_000)) -temp = similar(v) -AK.sort!(v, temp=temp) +ws = AK.workspace(AK.sort!, v) +AK.sort!(v; workspace=ws) ``` diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index e92b7918..3790ca33 100644 --- a/src/AcceleratedKernels.jl +++ b/src/AcceleratedKernels.jl @@ -27,6 +27,7 @@ const get_backend = KernelAbstractions.get_backend # Include code from other files include("utils.jl") include("algorithms.jl") +include("workspace.jl") include("task_partitioner.jl") include("foreachindex.jl") include("map.jl") diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index 63c29fe9..ce2cb5ea 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -20,8 +20,7 @@ include("accumulate_nd.jl") inclusive::Bool=true, acctype=nothing, alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) Compute accumulated running totals along a sequence by applying a binary operator to all elements @@ -70,12 +69,8 @@ On the host, accumulation is typically a memory-bound operation, so multithreade only becomes faster for more compute-heavy operations that hide memory latency, e.g. accumulating tuples or structs, or expensive operators. -The temporaries only apply to whole-array GPU scans with a known neutral element: `temp` stores -per-block aggregates, with the accumulator element type (see **Types** above); `temp_flags` stores -`DecoupledLookback`'s block flags (any integer type). Both need at least -`cld(length(dst), block_size * items_per_thread)` elements of the resolved algorithm. -Multi-block exclusive scans with `DecoupledLookback` use two epilogue kernels to shift the result -in place; they reuse `temp` for tile-boundary values and do not allocate a full-array copy. +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call, so that +the scan allocates none of its own. # Examples Example computing an inclusive prefix sum (the typical GPU "scan"): @@ -91,60 +86,96 @@ AK.accumulate!(+, v; inclusive=false) # 0, 1, 3, ...: starts from the neut AK.accumulate!(+, v; alg=AK.ScanPrefixes(block_size=512)) ``` """ -function accumulate!(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) - _accumulate_impl!(op, v, v, _resolve_backend(backend, v); kwargs...) +function accumulate!(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, + workspace=nothing, kwargs...) + s = _accumulate_setup(op, eltype(v), eltype(v), size(v), _resolve_backend(backend, v); + kwargs...) + _accumulate_run!(op, v, v, s, _buffers(s.plan, workspace, v)) end function accumulate!( op, dst::AbstractArray, src::AbstractArray; backend::Union{Nothing, Backend}=nothing, + workspace=nothing, kwargs... ) - if dst !== src - Base.mightalias(dst, src) && throw(ArgumentError( - "the destination of a scan must be its source or not overlap it")) - axes(dst) == axes(src) || throw(DimensionMismatch( - "the destination of a scan must have the source's axes $(axes(src)), " * - "got $(axes(dst))")) - end - _accumulate_impl!(op, dst, src, _resolve_backend(backend, dst, src); kwargs...) + _check_scan_destination(dst, src) + s = _accumulate_setup(op, eltype(dst), eltype(src), size(dst), + _resolve_backend(backend, dst, src); kwargs...) + _accumulate_run!(op, dst, src, s, _buffers(s.plan, workspace, dst, src)) +end + +_plan(::typeof(accumulate!), op, v::AbstractArray; backend=nothing, kwargs...) = + _accumulate_setup(op, eltype(v), eltype(v), size(v), _resolve_backend(backend, v); + kwargs...).plan +function _plan(::typeof(accumulate!), op, dst::AbstractArray, src::AbstractArray; + backend=nothing, kwargs...) + _check_scan_destination(dst, src) + _accumulate_setup(op, eltype(dst), eltype(src), size(dst), + _resolve_backend(backend, dst, src); kwargs...).plan +end + +function _check_scan_destination(dst, src) + dst === src && return nothing + Base.mightalias(dst, src) && throw(ArgumentError( + "the destination of a scan must be its source or not overlap it")) + axes(dst) == axes(src) || throw(DimensionMismatch( + "the destination of a scan must have the source's axes $(axes(src)), got $(axes(dst))")) + nothing end -# Scan `src` into `v` (which may be `src`), after checking the algorithm -function _accumulate_impl!( - op, v::AbstractArray, src::AbstractArray, backend::Backend; +# Everything a scan of elements of type `T` into an array of element type `D` and size `sz` +# resolves before touching data: its plan (backend, algorithm, scratch), running-value type and +# partial-result seed (see `_reduce_seed`) +function _accumulate_setup( + op, ::Type{D}, ::Type{T}, sz::Dims, backend::Backend; init=_NoInit(), neutral=nothing, dims::Union{Nothing, Integer}=nothing, inclusive::Bool=true, acctype=nothing, alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, -) +) where {D, T} dims isa Integer && dims < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) - A = _scan_acctype(op, eltype(v), eltype(src), init, inclusive, - _scan_combines(size(src), dims, init, inclusive), acctype) - neutral = _reduce_seed(op, A, neutral) + n = Base.prod(sz; init=1) + A = _scan_acctype(op, D, T, init, inclusive, _scan_combines(sz, dims, init, inclusive), + acctype) + seed = _reduce_seed(op, A, neutral) # The kernels' tiles hold partial results: lanes, when `op` has no known neutral element - a = _resolve_scan(alg, backend, A, dims, typeof(neutral)) - if dims !== nothing && (temp !== nothing || temp_flags !== nothing) - throw(ArgumentError( - "`temp` and `temp_flags` only apply to whole-array scans (`dims=nothing`)")) - end - if !inclusive && init isa _NoInit && neutral isa _Lane && !_valid(neutral) + a = _resolve_scan(alg, backend, A, dims, typeof(seed)) + if !inclusive && init isa _NoInit && seed isa _Lane && !_valid(seed) throw(ArgumentError( "an exclusive scan without `init` starts from the neutral element of `op`, " * "which is not known for $op; pass `init` or `neutral`")) end + # A destination of another element type scans in a scratch array of the accumulator type + work = D === A ? (;) : (; work=_buffer(A, sz)) + prefixes = if a isa Union{ScanPrefixes, DecoupledLookback} + num_blocks = cld(n, a.block_size * a.items_per_thread) + p = _buffer(typeof(seed), num_blocks) + a isa DecoupledLookback ? (; prefixes=p, flags=_buffer(UInt8, num_blocks)) : + (; prefixes=p) + else + (;) + end + return (; plan=_Plan(backend, a, merge(work, prefixes)), A=Val(A), seed, init, dims, + inclusive) +end + +# Scan `src` into `v` (which may be `src`), as set up by `_accumulate_setup`, with the plan's +# scratch buffers `bufs` +function _accumulate_run!(op, v, src, s, bufs) + backend, a = s.plan.backend, s.plan.alg + init, neutral, dims, inclusive = s.init, s.seed, s.dims, s.inclusive + A = _unval(s.A) # As Base, scan in the accumulator type and convert to `v`'s element type only when storing: - # in `v` when it has that type, else in a scratch array. Elements enter as one-element + # in `v` when it has that type, else in the scratch array. Elements enter as one-element # reductions (`Base.reduce_first`), and an inclusive scan applies `init` to the first element # of each slice, as `op(init, x)`. - w = eltype(v) === A ? v : similar(v, A) + w = haskey(bufs, :work) ? bufs.work : v launch = _scan_launch(a) if w !== src _foreachindex(eachindex(w, src), backend; launch...) do i @@ -163,7 +194,9 @@ function _accumulate_impl!( if a isa CPUThreads.Partitioned accumulate_1d_cpu!(op, w, backend, a; init, neutral, inclusive) else - accumulate_1d_gpu!(op, w, backend, a; init, neutral, inclusive, temp, temp_flags) + accumulate_1d_gpu!(op, w, backend, a; init, neutral, inclusive, + prefixes=get(bufs, :prefixes, nothing), + flags=get(bufs, :flags, nothing)) end elseif dims > ndims(w) # Every slice has one element: inclusive scans applied `init` above, exclusive ones start @@ -223,6 +256,12 @@ function _scan_apply_init!(w, op, init, dims, backend, launch) return w end +# A scan nested in another operation, with the scratch buffers of its plan in the outer one +function _accumulate_nested!(op, v, bufs; backend, kwargs...) + s = _accumulate_setup(op, eltype(v), eltype(v), size(v), backend; kwargs...) + _accumulate_run!(op, v, v, s, bufs) +end + """ accumulate(op, v::AbstractArray; init=, kwargs...) @@ -236,8 +275,19 @@ function accumulate(op, v::AbstractArray; backend::Union{Nothing, Backend}=nothi init=_NoInit(), dims=nothing, inclusive::Bool=true, acctype=nothing, kwargs...) backend = _resolve_backend(backend, v) - combines = _scan_combines(size(v), dims, init, inclusive) - D = _scan_acctype(op, Union{}, eltype(v), init, inclusive, combines, acctype) + D = _accumulate_eltype(op, v, init, dims, inclusive, acctype) accumulate!(op, _similar(backend, v, D), v; backend, init, dims, inclusive, acctype, kwargs...) end + +function _plan(::typeof(accumulate), op, v::AbstractArray; backend=nothing, init=_NoInit(), + dims=nothing, inclusive::Bool=true, acctype=nothing, kwargs...) + D = _accumulate_eltype(op, v, init, dims, inclusive, acctype) + _accumulate_setup(op, D, eltype(v), size(v), _resolve_backend(backend, v); init, dims, + inclusive, acctype, kwargs...).plan +end + +# The element type `accumulate` allocates: the running-value type from `init`'s type alone +_accumulate_eltype(op, v, init, dims, inclusive, acctype) = + _scan_acctype(op, Union{}, eltype(v), init, inclusive, + _scan_combines(size(v), dims, init, inclusive), acctype) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 280d9ce4..0c06a81e 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -282,18 +282,6 @@ end end -# The partial-result buffers of a whole-array scan: the tiles' totals, and DecoupledLookback's -# block flags -function _scan_prefixes(v, neutral, num_blocks, temp) - isnothing(temp) && return similar(v, typeof(neutral), num_blocks) - neutral isa _Lane && throw(ArgumentError( - "`temp` needs a neutral element of the operator; pass `neutral`")) - eltype(temp) === eltype(v) || throw(ArgumentError( - "`temp` must have the scanned element type $(eltype(v)), got $(eltype(temp))")) - @argcheck length(temp) >= num_blocks - return view(temp, 1:num_blocks) -end - # The first block's seed: `init` in the element type, else the neutral seed _scan_first_seed(v, init, neutral) = _lift(neutral, convert(eltype(v), init)) _scan_first_seed(v, ::_NoInit, neutral) = neutral @@ -305,8 +293,8 @@ function accumulate_1d_gpu!( init, neutral, inclusive::Bool, - temp::Union{Nothing, AbstractArray}, - temp_flags::Union{Nothing, AbstractArray}, + prefixes::Union{Nothing, AbstractArray}, + flags::Union{Nothing, AbstractArray}, ) block_size, items_per_thread = alg.block_size, alg.items_per_thread @@ -318,17 +306,8 @@ function accumulate_1d_gpu!( elems_per_block = block_size * items_per_thread num_blocks = (length(v) + elems_per_block - 1) ÷ elems_per_block items = Val(items_per_thread) - prefixes = _scan_prefixes(v, neutral, num_blocks, temp) seed = _scan_first_seed(v, init, neutral) - if isnothing(temp_flags) - flags = similar(v, UInt8, num_blocks) - else - @argcheck eltype(temp_flags) <: Integer - @argcheck length(temp_flags) >= num_blocks - flags = view(temp_flags, 1:num_blocks) - end - shift_to_exclusive = !inclusive && num_blocks > 1 block_inclusive = inclusive || shift_to_exclusive @@ -359,8 +338,8 @@ function accumulate_1d_gpu!( init, neutral, inclusive::Bool, - temp::Union{Nothing, AbstractArray}, - temp_flags::Union{Nothing, AbstractArray}, + prefixes::Union{Nothing, AbstractArray}, + flags::Union{Nothing, AbstractArray}, ) block_size, items_per_thread = alg.block_size, alg.items_per_thread @@ -372,7 +351,6 @@ function accumulate_1d_gpu!( elems_per_block = block_size * items_per_thread num_blocks = (length(v) + elems_per_block - 1) ÷ elems_per_block items = Val(items_per_thread) - prefixes = _scan_prefixes(v, neutral, num_blocks, temp) seed = _scan_first_seed(v, init, neutral) kernel1! = _accumulate_block!(backend, block_size) diff --git a/src/arithmetics.jl b/src/arithmetics.jl index 68b48f24..9d06ef29 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -44,24 +44,21 @@ prod(src::AbstractArray; kwargs...) = _reduce_or_empty(one, Base.mul_prod, src; # applying it as an `init`, which would change results such as the sign of a sum of `-0.0`s function _reduce_or_empty(empty, op, src; init=_NoInit(), dims=:, acctype=nothing, backend::Union{Nothing, Backend}=nothing, alg::Algorithm=Auto(), - kwargs...) - init isa _NoInit || return reduce(op, src; init, dims, acctype, backend, alg, kwargs...) + workspace=nothing, kwargs...) + kwargs = (; dims, acctype, backend, alg, kwargs...) + init isa _NoInit || return reduce(op, src; init, workspace, kwargs...) A = _acctype(op, Union{}, _mapped_eltype(identity, src), acctype) - if A !== Union{} && isempty(src) - # (the algorithm is checked as for any other input) - b = _resolve_backend(backend, src) - if _whole(dims) - _resolve_reduce(alg, b, A, dims) - return empty(A) - end + if A !== Union{} && isempty(src) && + (_whole(dims) || Base.any(d -> size(src, d) == 0, _reduced_dims(dims, ndims(src)))) + # (the algorithm and the workspace are checked as for any other input) + p = _plan(reduce, op, src; kwargs...) + _buffers(p, workspace, src) + _whole(dims) && return empty(A) dims_valid = _reduced_dims(dims, ndims(src)) - if Base.any(d -> size(src, d) == 0, dims_valid) - _resolve_reduce(alg, b, A, dims_valid) - dst_sizes = ntuple(d -> d in dims_valid ? 1 : size(src, d), ndims(src)) - return fill!(KernelAbstractions.allocate(b, A, dst_sizes), empty(A)) - end + dst_sizes = ntuple(d -> d in dims_valid ? 1 : size(src, d), ndims(src)) + return fill!(KernelAbstractions.allocate(p.backend, A, dst_sizes), empty(A)) end - return reduce(op, src; dims, acctype, backend, alg, kwargs...) + return reduce(op, src; workspace, kwargs...) end @@ -121,6 +118,14 @@ count(src::AbstractArray; kwargs...) = count(identity, src; kwargs...) count(f, src::AbstractArray; init=0, kwargs...) = mapreduce(_BoolValued(f), Base.add_sum, src; init, kwargs...) +_plan(::typeof(sum), src::AbstractArray; kwargs...) = _plan(reduce, Base.add_sum, src; kwargs...) +_plan(::typeof(prod), src::AbstractArray; kwargs...) = _plan(reduce, Base.mul_prod, src; kwargs...) +_plan(::typeof(maximum), src::AbstractArray; kwargs...) = _plan(reduce, max, src; kwargs...) +_plan(::typeof(minimum), src::AbstractArray; kwargs...) = _plan(reduce, min, src; kwargs...) +_plan(::typeof(count), src::AbstractArray; kwargs...) = _plan(count, identity, src; kwargs...) +_plan(::typeof(count), f, src::AbstractArray; init=0, kwargs...) = + _plan(mapreduce, _BoolValued(f), Base.add_sum, src; init, kwargs...) + """ cumsum(src::AbstractArray; kwargs...) @@ -172,3 +177,6 @@ p = AK.cumprod(m, dims=1) ``` """ cumprod(src::AbstractArray; kwargs...) = accumulate(Base.mul_prod, src; kwargs...) + +_plan(::typeof(cumsum), src::AbstractArray; kwargs...) = _plan(accumulate, Base.add_sum, src; kwargs...) +_plan(::typeof(cumprod), src::AbstractArray; kwargs...) = _plan(accumulate, Base.mul_prod, src; kwargs...) diff --git a/src/findall.jl b/src/findall.jl index 6fdbfc44..e1aa6e20 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -141,22 +141,9 @@ end end -function findall_temp(bools, backend, len, temp) - if isnothing(temp) - return KernelAbstractions.allocate(backend, Int, len) - end - - @argcheck get_backend(temp) === backend - @argcheck eltype(temp) === Int - @argcheck length(temp) >= len - @argcheck !Base.mightalias(temp, bools) - view(temp, 1:len) -end - - function findall_gpu( bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg::ScanScatter; - temp, + bufs, ) where I block_size = alg.block_size items_per_thread = alg.items_per_thread @@ -166,14 +153,14 @@ function findall_gpu( elems_per_block = block_size * items_per_thread num_blocks = cld(length(bools), elems_per_block) - block_counts = findall_temp(bools, backend, num_blocks, temp) + block_counts = bufs.counts input_indices = eachindex(bools) items = Val(items_per_thread) kernel! = findall_block!(backend, block_size) kernel!(nothing, bools, block_counts, input_indices, output_indices, items; ndrange=num_blocks * block_size) - accumulate!(+, block_counts; backend, init=0) + _accumulate_nested!(+, block_counts, bufs.scan; backend, init=0) n = @allowscalar block_counts[end] out = KernelAbstractions.allocate(backend, I, n) @@ -188,7 +175,7 @@ end function findall_cpu( bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg::CPUThreads.Partitioned; - temp, + bufs, ) where I input_indices = eachindex(bools) tp = TaskPartitioner(length(bools), alg.max_tasks, alg.min_elems) @@ -198,7 +185,7 @@ function findall_cpu( return out end - task_counts = findall_temp(bools, backend, tp.num_tasks, temp) + task_counts = bufs.counts itask_partition(tp) do itask, positions task_counts[itask] = Base.count( position -> @inbounds(bools[findall_index(input_indices, position)]), positions, @@ -227,13 +214,12 @@ end function findall_impl( - bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg; - temp::Union{Nothing, AbstractArray}, + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg; bufs, ) where I if alg isa ScanScatter - findall_gpu(bools, I, output_indices, backend, alg; temp) + findall_gpu(bools, I, output_indices, backend, alg; bufs) else - findall_cpu(bools, I, output_indices, backend, alg; temp) + findall_cpu(bools, I, output_indices, backend, alg; bufs) end end @@ -242,17 +228,7 @@ end _findall_launch(a::ScanScatter) = (; block_size=a.block_size) _findall_launch(a::CPUThreads.Partitioned) = (; max_tasks=a.max_tasks, min_elems=a.min_elems) -function findall_bools(pred, v::AbstractArray, backend::Backend, alg, temp_bools) - bools = if isnothing(temp_bools) - # On the resolved backend, which a range does not determine - KernelAbstractions.allocate(backend, Bool, size(v)) - else - @argcheck get_backend(temp_bools) === backend - @argcheck eltype(temp_bools) === Bool - @argcheck axes(temp_bools) == axes(v) - @argcheck !Base.mightalias(temp_bools, v) - temp_bools - end +function findall_bools!(bools, pred, v::AbstractArray, backend::Backend, alg) input_indices = eachindex(v) bool_indices = eachindex(bools) _foreachindex(Base.OneTo(length(v)), backend; _findall_launch(alg)...) do position @@ -265,10 +241,8 @@ end """ - findall(A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), temp=nothing, - temp_bools=nothing) - findall(pred, A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), temp=nothing, - temp_bools=nothing) + findall(A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), workspace=nothing) + findall(pred, A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), workspace=nothing) Stream compaction: select, in order, the elements of `items` at the positions where `A` is `true`, or where `pred` returns `true` for `A`'s elements. Values used as conditions, and `pred`'s results, @@ -289,13 +263,9 @@ The supported inputs are arrays. Dictionaries, other iterables, and scalar input `alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host and [`ScanScatter`](@ref) on GPUs, with -the device's settings. `backend` is derived from `A` and `items`. - -`temp` may provide the `Int` buffer used for block or task counts: at least -`cld(length(A), block_size * items_per_thread)` elements for `ScanScatter`, one per task on the -host. `temp_bools` may provide the Bool mask for the predicate form or for a mask whose element -type is not `Bool`; it must have the same axes as `A` and must not alias it. Omitted buffers are -allocated automatically. +the device's settings. `backend` is derived from `A` and `items`. `workspace` takes the scratch +memory of a [`workspace`](@ref) made for the same call (the mask of the predicate form and the +counts), so that `findall` allocates only its result. # Examples ```julia @@ -311,44 +281,55 @@ AK.findall(m) # [CartesianIndex(1, 1), CartesianIndex( AK.findall(m; items=LinearIndices(m)) # [1, 4] ``` """ -function findall( - values::AbstractArray; - items::AbstractArray=keys(values), - backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, - temp_bools::Union{Nothing, AbstractArray}=nothing, -) +function findall(values::AbstractArray; items::AbstractArray=keys(values), + backend::Union{Nothing, Backend}=nothing, alg::Algorithm=Auto(), + workspace=nothing) values isa AbstractArray{Bool} || _check_bool_result(identity, values) - _check_findall_items(values, items) - backend = _resolve_backend(backend, values, items) - a = _resolve_findall(alg, backend, eltype(values)) - bools = if values isa AbstractArray{Bool} - isnothing(temp_bools) || - throw(ArgumentError("temp_bools is not used for a Bool mask")) - values - else - findall_bools(identity, values, backend, a, temp_bools) - end - findall_impl(bools, eltype(items), items, backend, a; temp) + s = _findall_setup(values, items, false, backend, alg) + _findall_run(identity, values, items, s, _buffers(s.plan, workspace, values, items)) end - -function findall( - pred, v::AbstractArray; - items::AbstractArray=keys(v), - backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, - temp_bools::Union{Nothing, AbstractArray}=nothing, -) +function findall(pred, v::AbstractArray; items::AbstractArray=keys(v), + backend::Union{Nothing, Backend}=nothing, alg::Algorithm=Auto(), + workspace=nothing) _check_bool_result(pred, v) - _check_findall_items(v, items) + s = _findall_setup(v, items, true, backend, alg) + _findall_run(pred, v, items, s, _buffers(s.plan, workspace, v, items)) +end + +_plan(::typeof(findall), values::AbstractArray; items::AbstractArray=keys(values), + backend=nothing, alg::Algorithm=Auto()) = + _findall_setup(values, items, false, backend, alg).plan +_plan(::typeof(findall), pred, v::AbstractArray; items::AbstractArray=keys(v), backend=nothing, + alg::Algorithm=Auto()) = + _findall_setup(v, items, true, backend, alg).plan + +# The plan of `findall` over `v`: its scratch is a mask of `pred(v[i])` for the predicate form +# and for non-Bool values, the counts of the blocks or tasks, and the scan of the block counts +function _findall_setup(v, items, predicate::Bool, backend, alg) + length(items) == length(v) || throw(DimensionMismatch( + "`items` must have the length of the array, $(length(v)), got $(length(items))")) backend = _resolve_backend(backend, v, items) a = _resolve_findall(alg, backend, eltype(v)) - bools = findall_bools(pred, v, backend, a, temp_bools) - findall_impl(bools, eltype(items), items, backend, a; temp) + n = length(v) + mask = predicate || !(v isa AbstractArray{Bool}) ? (; mask=_buffer(Bool, size(v))) : (;) + counts, nested = if a isa ScanScatter + num_blocks = cld(n, a.block_size * a.items_per_thread) + scan = _accumulate_setup(+, Int, Int, (num_blocks,), backend; init=0).plan + (; counts=_buffer(Int, num_blocks), scan=scan.sizes), (; scan=scan.alg) + else + num_tasks = TaskPartitioner(n, a.max_tasks, a.min_elems).num_tasks + (; counts=_buffer(Int, num_tasks)), (;) + end + return (; plan=_Plan(backend, a, nested, merge(mask, counts))) end -_check_findall_items(v, items) = length(items) == length(v) || throw(DimensionMismatch( - "`items` must have the length of the array, $(length(v)), got $(length(items))")) +function _findall_run(pred, v, items, s, bufs) + backend, a = s.plan.backend, s.plan.alg + bools = if haskey(bufs, :mask) + findall_bools!(bufs.mask, pred, v, backend, a) + else + v + end + findall_impl(bools, eltype(items), items, backend, a; bufs) +end diff --git a/src/predicates.jl b/src/predicates.jl index 598f1db8..b4c02f74 100644 --- a/src/predicates.jl +++ b/src/predicates.jl @@ -118,7 +118,7 @@ end """ - any(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) + any(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto(), workspace=nothing) Check if any element of `v` satisfies the predicate `pred` (i.e. some `pred(v[i]) == true`); `pred` must return a `Bool`. Optimised differently to `mapreduce` due to shortcircuiting behaviour of @@ -130,7 +130,8 @@ reduction. `alg` is [`Auto()`](@ref Auto) by default: [`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned) on the host, and [`ConcurrentWrite`](@ref) on GPUs ([`ViaReduce`](@ref) on oneAPI). `backend` is derived from `v`; pass it for inputs that do not -determine it, such as index ranges. +determine it, such as index ranges. `workspace` takes +the scratch memory of a [`workspace`](@ref) made for the same call. On the host, multithreaded parallelisation is only worth it for large arrays, relatively expensive predicates, and/or rare occurrence of true; use `CPUThreads.Partitioned(; max_tasks, min_elems)` @@ -158,12 +159,14 @@ complex_any(CuArray(rand(Float32, 100)), CuArray(rand(Float32, 100))) ``` """ function any(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto()) - _check_bool_result(pred, v) - backend = _resolve_backend(backend, v) - _any(pred, v, backend, _resolve_predicate(alg, backend, eltype(v))) + alg::Algorithm=Auto(), workspace=nothing) + p = _predicate_plan(pred, |, v, backend, alg) + _any(pred, v, p.backend, p.alg, _buffers(p, workspace, v)) end +_plan(::typeof(any), pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) = + _predicate_plan(pred, |, v, backend, alg) + # The predicate must return a `Bool`, as in Base. Where inference shows it cannot, the call fails # on the host, before launching, with an `ArgumentError` instead of an exception from the kernel # (which GPU backends report without Base's `TypeError`). @@ -175,18 +178,36 @@ function _check_bool_result(pred, v) nothing end -function _any(pred, v, backend, alg::ConcurrentWrite) +# The plan of `any`/`all` (whose `ViaReduce` reduces with `op`): the flag of `ConcurrentWrite`, +# or the scratch of the reduction +function _predicate_plan(pred, op, v, backend, alg) + _check_bool_result(pred, v) + backend = _resolve_backend(backend, v) + a = _resolve_predicate(alg, backend, eltype(v)) + sizes = if a isa ConcurrentWrite + (; flag=_buffer(Int8, 1)) + elseif a isa ViaReduce + (; reduce=_mapreduce_setup(_BoolValued(pred), op, v, backend, op === (|) ? false : true, + nothing, nothing, :, a.reduce).plan.sizes) + else + (;) + end + return _Plan(backend, a, sizes) +end + +function _any(pred, v, backend, alg::ConcurrentWrite, bufs) # `ndrange` must not be zero isempty(v) && return false - out = KernelAbstractions.zeros(backend, Int8, 1) + out = bufs.flag + fill!(out, Int8(0)) _any_global!(backend, alg.block_size)(out, pred, v, ndrange=length(v)) return @allowscalar(out[1]) != 0 end -_any(pred, v, backend, alg::ViaReduce) = - mapreduce(_BoolValued(pred), |, v; backend, init=false, alg=alg.reduce) +_any(pred, v, backend, alg::ViaReduce, bufs) = + _mapreduce_nested(_BoolValued(pred), |, v, bufs.reduce; backend, init=false, alg=alg.reduce) -function _any(pred, v, backend, alg::CPUThreads.Partitioned) +function _any(pred, v, backend, alg::CPUThreads.Partitioned, bufs) overall = Ref(false) task_partition(length(v), alg.max_tasks, alg.min_elems) do irange for i in irange @@ -203,7 +224,7 @@ end """ - all(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) + all(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto(), workspace=nothing) Check if all elements of `v` satisfy the predicate `pred` (i.e. all `pred(v[i]) == true`); `pred` must return a `Bool`. The keywords are those of [`any`](@ref). @@ -229,13 +250,15 @@ complex_all(MtlArray(rand(Float32, 100)), MtlArray(rand(Float32, 100))) ``` """ function all(pred, v::AbstractArray; backend::Union{Nothing, Backend}=nothing, - alg::Algorithm=Auto()) - _check_bool_result(pred, v) - backend = _resolve_backend(backend, v) - _all(pred, v, backend, _resolve_predicate(alg, backend, eltype(v))) + alg::Algorithm=Auto(), workspace=nothing) + p = _predicate_plan(pred, &, v, backend, alg) + _all(pred, v, p.backend, p.alg, _buffers(p, workspace, v)) end -_all(pred, v, backend, alg::ConcurrentWrite) = !_any(!pred, v, backend, alg) -_all(pred, v, backend, alg::ViaReduce) = - mapreduce(_BoolValued(pred), &, v; backend, init=true, alg=alg.reduce) -_all(pred, v, backend, alg::CPUThreads.Partitioned) = !_any(!pred, v, backend, alg) +_plan(::typeof(all), pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) = + _predicate_plan(pred, &, v, backend, alg) + +_all(pred, v, backend, alg::ConcurrentWrite, bufs) = !_any(!pred, v, backend, alg, bufs) +_all(pred, v, backend, alg::ViaReduce, bufs) = + _mapreduce_nested(_BoolValued(pred), &, v, bufs.reduce; backend, init=true, alg=alg.reduce) +_all(pred, v, backend, alg::CPUThreads.Partitioned, bufs) = !_any(!pred, v, backend, alg, bufs) diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 09f5bc90..51d27883 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -39,14 +39,15 @@ end # Reduce the non-empty `src` to a host value; `neutral` is the partial-result seed of -# `_reduce_seed`, and `init` is a value or `_NoInit()`. +# `_reduce_seed`, `init` is a value or `_NoInit()`, and `partials` holds `_mapreduce_1d_partials` +# partial results (`nothing` if none are needed). function mapreduce_1d_gpu( f, op, src::MapReduceSource, backend::Backend; init, neutral, block_size::Int, items_per_thread::Int, - temp::Union{Nothing, AbstractArray}, + partials::Union{Nothing, AbstractArray}, switch_below::Int, ) @argcheck 1 <= block_size <= 1024 @@ -71,17 +72,7 @@ function mapreduce_1d_gpu( num_per_block = items_per_thread * block_size blocks = (len + num_per_block - 1) ÷ num_per_block - if !isnothing(temp) - neutral isa _Lane && throw(ArgumentError( - "`temp` needs a neutral element of the operator; pass `neutral`")) - @argcheck get_backend(temp) === backend - eltype(temp) === P || throw(ArgumentError( - "`temp` must have the reduction's accumulator element type $P, got $(eltype(temp))")) - @argcheck length(temp) >= blocks * 2 - dst = temp - else - dst = KernelAbstractions.allocate(backend, P, blocks * 2) - end + dst = partials # Later the kernel will be compiled for views anyways, so use same types for arrays. src_view = _mapreduce_1d_src_view(src) @@ -125,7 +116,22 @@ function mapreduce_1d_gpu( end _host_copy(src::AbstractArray) = Array(src) -_host_copy(src::Base.Broadcast.Broadcasted) = Array(Base.Broadcast.materialize(src)) +# A `Broadcasted` source is evaluated on the host from host copies of its arrays, so that it needs +# no device memory +_host_copy(src::Base.Broadcast.Broadcasted) = Base.Broadcast.materialize(_on_host(src)) +_on_host(bc::Base.Broadcast.Broadcasted) = + Base.Broadcast.Broadcasted(bc.f, Base.map(_on_host, bc.args), bc.axes) +_on_host(x::Base.Broadcast.Extruded) = _on_host(x.x) +_on_host(x::AbstractArray) = Array(x) +_on_host(x::AbstractRange) = x +_on_host(x) = x _mapreduce_1d_src_view(src::AbstractArray) = @view src[1:end] _mapreduce_1d_src_view(src::Base.Broadcast.Broadcasted) = src + +# The number of partial results of `mapreduce_1d_gpu` over `len` elements (two buffers of one per +# block of the first pass), or 0 when it needs none +function _mapreduce_1d_partials(len, block_size, items_per_thread, switch_below) + (len <= 1 || len < switch_below) && return 0 + return 2 * cld(len, block_size * items_per_thread) +end diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 165b9da4..0960ecc8 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -67,6 +67,101 @@ function _canonicalize_dims(src_sizes, src_strides, dims_valid) return Tuple(reduce_segs), Tuple(outer_segs) end +# The reduced and kept dimensions of a strided source, as `(strides, sizes)` of their merged +# segments. +function _mapreduce_segments(src_sizes, src_strides, dims_valid) + reduce_segs, outer_segs = _canonicalize_dims(src_sizes, src_strides, dims_valid) + return (Tuple(str for (str, _) in reduce_segs), Tuple(s for (_, s) in reduce_segs)), + (Tuple(str for (str, _) in outer_segs), Tuple(s for (_, s) in outer_segs)) +end + +# The launch shape of a strided reduction along `dims`, and the number of blocks per output of a +# `:multigroup` reduction: +# - square-like, contiguous output with strided input -> :tiled_strided +# - dst_size >= reduce_size -> :by_thread +# - dst_size == 1, or dst_size < GS_DST_CUTOFF +# (and dst_size < reduce_size) -> :multigroup (split one reduction +# across blocks, needs a second-pass +# combine and `partials` scratch), or +# :by_block_each when that is one block +# - otherwise (GS_DST_CUTOFF <= dst_size < reduce_size) -> :by_block, grid-striding over +# outputs, single pass +# +# Rationale: grid-striding by_block launches `min(dst_size, target_blocks)` blocks; +# for very small dst_size (e.g. 5 or 9) that under-fills an 84-SM GPU, so splitting +# the (large) reduction itself across many blocks via multigroup is still better. +# For dst_size==1 there is nothing to grid-stride over, so multigroup is the only +# option regardless of GS_DST_CUTOFF. +function _mapreduce_nd_shape( + dst_size, reduce_size, reduce_strides, reduce_sizes, outer_strides, outer_sizes, + block_size, target_blocks, +) + # by_block override 1: when the reduced dimension is the fastest-varying one + # (reduce_strides==(1,)), by_thread's per-thread strided access is badly + # uncoalesced, while by_block lets consecutive threads read consecutive elements. + use_by_block_for_coalescing = + dst_size >= reduce_size && reduce_size >= block_size && + length(reduce_sizes) == 1 && reduce_sizes[1] != 0 && + reduce_strides == (1,) + + # by_block override 2: when by_thread would launch too few blocks for a square-like + # output/reduction shape, split each output reduction across a full block. This + # fallback is mostly for layouts that do not satisfy the narrower tiled-strided + # pattern below; the common contiguous-output strided case uses tiled_strided. + # Avoid applying this to wide-output shapes, where by_thread's cross-output + # coalescing is better than a strided block reduction. + use_by_block_for_low_occupancy = + dst_size == reduce_size && reduce_size >= block_size && + cld(dst_size, block_size) < target_blocks && + length(reduce_sizes) == 1 && reduce_sizes[1] != 0 + + use_by_block = use_by_block_for_coalescing || use_by_block_for_low_occupancy + use_tiled_strided = + dst_size == reduce_size && reduce_size >= block_size && + length(outer_sizes) == 1 && outer_strides == (1,) && + length(reduce_sizes) == 1 && reduce_strides[1] > 1 && + block_size % TILED_STRIDED_ROWS_PER_BLOCK == 0 && + ispow2(block_size ÷ TILED_STRIDED_ROWS_PER_BLOCK) + + if use_tiled_strided + return :tiled_strided, 0 + elseif dst_size >= reduce_size && !use_by_block + return :by_thread, 0 + elseif use_by_block + return :by_block, 0 + elseif dst_size == 1 || dst_size < GS_DST_CUTOFF + reduce_groups = min( + cld(reduce_size, block_size), + block_size, + cld(target_blocks, dst_size), + cld(reduce_size, block_size * MIN_ITEMS_PER_THREAD), + ) + return reduce_groups > 1 ? (:multigroup, reduce_groups) : (:by_block_each, 0) + else + return :by_block, 0 + end +end + +# The scratch of a reduction along `dims` into `dst_size` outputs: the partial results of a +# `:multigroup` reduction, with the partial-result type of the seed `neutral` (see `_reduce_seed`) +function _mapreduce_nd_sizes(src, backend, alg, ::Type{A}, neutral, dims_valid, dst_size) where {A} + alg isa BlockReduce && A !== Union{} || return (;) + src_sizes = size(src) + len = Base.prod(src_sizes[d] for d in dims_valid; init=1) + (len <= 1 || dst_size == 0) && return (;) + layout = _mapreduce_strided_layout(src) + isnothing(layout) && return (;) + target_blocks = reduce_tuning(backend, A).target_blocks + (reduce_strides, reduce_sizes), (outer_strides, outer_sizes) = + _mapreduce_segments(src_sizes, layout[3], dims_valid) + shape, reduce_groups = _mapreduce_nd_shape( + dst_size, len, reduce_strides, reduce_sizes, outer_strides, outer_sizes, + alg.block_size, target_blocks, + ) + shape === :multigroup || return (;) + return (; partials=_buffer(typeof(neutral), dst_size, reduce_groups)) +end + # Main entry point # The reduced dimensions `dims` (an integer or a collection of integers), as Base validates them: @@ -92,6 +187,7 @@ function mapreduce_nd!( init, neutral, dims_valid::Tuple, + bufs::NamedTuple, ) where {A} # Launch settings of the elementwise passes, and of the reduction kernels launch = alg isa BlockReduce ? (; block_size=alg.block_size) : @@ -131,7 +227,6 @@ function mapreduce_nd!( # Degenerate cases end dst_size = length(dst) - neutral = _reduce_seed(op, A, neutral) if alg isa CPUThreads.Partitioned _mapreduce_nd_cpu_sections!(f, op, dst, src; init, neutral, launch...) @@ -158,61 +253,15 @@ function mapreduce_nd!( end buffer, base_offset, src_strides = layout - reduce_segs, outer_segs = _canonicalize_dims(src_sizes, src_strides, dims_valid) - - outer_strides = Tuple(str for (str, _) in outer_segs) - outer_sizes = Tuple(s for (_, s) in outer_segs) - reduce_strides = Tuple(str for (str, _) in reduce_segs) - reduce_sizes = Tuple(s for (_, s) in reduce_segs) - reduce_size = len - - # ───────────────────────────────────────────────────────────────────────── - # Dispatch decision (see header comment for the four paths): - # - # - square-like, contiguous output with strided input -> tiled_strided - # - dst_size >= reduce_size -> by_thread - # - dst_size == 1, or dst_size < GS_DST_CUTOFF - # (and dst_size < reduce_size) -> multigroup (split one - # reduction across blocks, - # needs a 2nd-pass combine) - # - otherwise (GS_DST_CUTOFF <= dst_size < reduce_size) -> by_block, grid-striding - # over outputs, single pass - # - # Rationale: grid-striding by_block launches `min(dst_size, target_blocks)` blocks; - # for very small dst_size (e.g. 5 or 9) that under-fills an 84-SM GPU, so splitting - # the (large) reduction itself across many blocks via multigroup is still better. - # For dst_size==1 there is nothing to grid-stride over, so multigroup is the only - # option regardless of GS_DST_CUTOFF. - # ───────────────────────────────────────────────────────────────────────── - - # by_block override 1: when the reduced dimension is the fastest-varying one - # (reduce_strides==(1,)), by_thread's per-thread strided access is badly - # uncoalesced, while by_block lets consecutive threads read consecutive elements. - use_by_block_for_coalescing = - dst_size >= reduce_size && reduce_size >= block_size && - length(reduce_sizes) == 1 && reduce_sizes[1] != 0 && - reduce_strides == (1,) - - # by_block override 2: when by_thread would launch too few blocks for a square-like - # output/reduction shape, split each output reduction across a full block. This - # fallback is mostly for layouts that do not satisfy the narrower tiled-strided - # pattern below; the common contiguous-output strided case uses tiled_strided. - # Avoid applying this to wide-output shapes, where by_thread's cross-output - # coalescing is better than a strided block reduction. - use_by_block_for_low_occupancy = - dst_size == reduce_size && reduce_size >= block_size && - cld(dst_size, block_size) < target_blocks && - length(reduce_sizes) == 1 && reduce_sizes[1] != 0 - - use_by_block = use_by_block_for_coalescing || use_by_block_for_low_occupancy - use_tiled_strided = - dst_size == reduce_size && reduce_size >= block_size && - length(outer_sizes) == 1 && outer_strides == (1,) && - length(reduce_sizes) == 1 && reduce_strides[1] > 1 && - block_size % TILED_STRIDED_ROWS_PER_BLOCK == 0 && - ispow2(block_size ÷ TILED_STRIDED_ROWS_PER_BLOCK) + (reduce_strides, reduce_sizes), (outer_strides, outer_sizes) = + _mapreduce_segments(src_sizes, src_strides, dims_valid) + reduce_size = len + shape, reduce_groups = _mapreduce_nd_shape( + dst_size, reduce_size, reduce_strides, reduce_sizes, outer_strides, outer_sizes, + block_size, target_blocks, + ) - if use_tiled_strided + if shape === :tiled_strided # Narrow layout-specific path: square-like row reductions with contiguous # outputs and strided input reads, e.g. size=(1024,1024), dims=2. Several # outputs share a block: small lane groups reduce one output each, preserving @@ -227,7 +276,7 @@ function mapreduce_nd!( base_offset, reduce_strides[1], dst_size, reduce_size, Val(rows_per_block), ndrange=(block_size * blocks,), ) - elseif dst_size >= reduce_size && !use_by_block + elseif shape === :by_thread # Many outputs, small reduction: one thread per output reduces sequentially blocks = cld(dst_size, block_size) kernel! = _mapreduce_nd_by_thread!(backend, block_size) @@ -237,62 +286,33 @@ function mapreduce_nd!( dst_size, reduce_size, ndrange=(block_size * blocks,), ) - elseif use_by_block - # Grid-stride by_block: one full block cooperates on each output, then - # grid-strides across remaining outputs if fewer blocks than outputs were - # launched. - launch_blocks = min(dst_size, target_blocks) - _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, neutral, init, + elseif shape === :multigroup + # Very few outputs, large reduction: split each output's reduction across + # `reduce_groups` blocks, combine the partials in a second pass + partial = bufs.partials + kernel! = _mapreduce_nd_multigroup!(backend, block_size) + kernel!( + buffer, partial, f, op, neutral, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - dst_size, reduce_size, launch_blocks, + Val(dst_size), reduce_size, reduce_groups, + ndrange=(block_size * dst_size * reduce_groups,), ) - elseif dst_size == 1 || dst_size < GS_DST_CUTOFF - # Very few outputs, large reduction: split each output's reduction across - # `reduce_groups` blocks (multigroup), combine partials in a second pass. - reduce_groups = min( - cld(reduce_size, block_size), - block_size, - cld(target_blocks, dst_size), - cld(reduce_size, block_size * MIN_ITEMS_PER_THREAD), + + # Second pass: reduce partial (dst_size × reduce_groups) → dst, one block + # per output. reduce_groups is small (<=block_size by construction), so use + # a small block size for this pass to avoid an oversubscribed tree-reduce. + pass2_block_size = _pass2_block_size(reduce_groups) + kernel2! = _mapreduce_partial_to_dst!(backend, pass2_block_size) + kernel2!( + partial, dst, op, neutral, init, + dst_size, reduce_groups, + ndrange=(pass2_block_size * dst_size,), ) - reduce_groups = max(reduce_groups, 1) - - if reduce_groups > 1 - partial = KernelAbstractions.allocate(backend, typeof(neutral), (dst_size, reduce_groups)) - - kernel! = _mapreduce_nd_multigroup!(backend, block_size) - kernel!( - buffer, partial, f, op, neutral, - base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - Val(dst_size), reduce_size, reduce_groups, - ndrange=(block_size * dst_size * reduce_groups,), - ) - - # Second pass: reduce partial (dst_size × reduce_groups) → dst, one block - # per output. reduce_groups is small (<=block_size by construction), so use - # a small block size for this pass to avoid an oversubscribed tree-reduce. - pass2_block_size = _pass2_block_size(reduce_groups) - kernel2! = _mapreduce_partial_to_dst!(backend, pass2_block_size) - kernel2!( - partial, dst, op, neutral, init, - dst_size, reduce_groups, - ndrange=(pass2_block_size * dst_size,), - ) - else - # reduce_groups collapsed to 1 (e.g. reduce_size <= block_size): just do a - # single-block-per-output reduction directly, no partial array needed. - _launch_mapreduce_nd_by_block!( - backend, block_size, buffer, dst, f, op, neutral, init, - base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, - dst_size, reduce_size, dst_size, - ) - end else - # GS_DST_CUTOFF <= dst_size < reduce_size: grid-stride over outputs, one pass. - # Cap launched blocks at target_blocks; each block handles - # ceil(dst_size / launch_blocks) outputs sequentially. - launch_blocks = min(dst_size, target_blocks) + # One block per output: grid-striding over the outputs with at most + # `target_blocks` blocks, or (a multigroup reduction that collapsed to one group, + # e.g. reduce_size <= block_size) one block for each of the few outputs + launch_blocks = shape === :by_block ? min(dst_size, target_blocks) : dst_size _launch_mapreduce_nd_by_block!( backend, block_size, buffer, dst, f, op, neutral, init, base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes, @@ -341,7 +361,10 @@ function _mapreduce_strided_layout(src::AbstractArray) return (p, base, s) end -_mapreduce_strided_layout(::Base.Broadcast.Broadcasted) = nothing +# (before Julia 1.12 a `Broadcasted` source reaches the kernels materialized, as a dense array: +# see `_mapreduce_source`) +_mapreduce_strided_layout(src::Base.Broadcast.Broadcasted) = + VERSION < v"1.12-" ? (nothing, 0, _mapreduce_dense_strides(size(src))) : nothing function _mapreduce_is_dense_buffer(p::AbstractArray) try diff --git a/src/reduce/reduce.jl b/src/reduce/reduce.jl index cc3908dd..9ab65927 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -49,7 +49,7 @@ reduce(op, src::AbstractArray; kwargs...) = mapreduce(identity, op, src; kwargs. acctype=nothing, dims=:, alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) Apply `f` to each element of `src` and reduce the results with the binary operator `op`. With @@ -86,8 +86,9 @@ device's settings. On the host, multithreading reductions only improves performa operations that hide the memory latency and thread launch overhead, e.g. reductions of tuples or structs, or expensive operators. -`temp` is scratch space for whole-array `BlockReduce` reductions: a vector with the accumulator -element type and at least `2 * cld(length(src), block_size * items_per_thread)` elements. +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call, so that +the reduction allocates none of its own (it still allocates its result along `dims`, and before +Julia 1.12 a materialized source). # Examples Computing a sum of squares, reducing down to a scalar that is copied to host: @@ -122,48 +123,103 @@ function mapreduce( acctype=nothing, dims=:, alg::Algorithm=Auto(), - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) - if !isempty(srcs) - src isa AbstractArray || - throw(ArgumentError("a Broadcasted source cannot be combined with more arrays")) - _mapreduce_check_map_axes(src, srcs...) - bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) - return _mapreduce(identity, op, bc, backend, init, neutral, acctype, dims, alg, temp) - end - return _mapreduce(f, op, src, backend, init, neutral, acctype, dims, alg, temp) + f, src = _mapreduce_fuse(f, src, srcs) + s = _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg) + bufs = _buffers(s.plan, workspace, src) + return _mapreduce_run(f, op, _mapreduce_source(src), s, init, bufs) +end + +function _plan( + ::typeof(mapreduce), f, op, src::MapReduceSource, srcs::AbstractArray...; + backend=nothing, init=_NoInit(), neutral=nothing, acctype=nothing, dims=:, + alg::Algorithm=Auto(), +) + f, src = _mapreduce_fuse(f, src, srcs) + return _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg).plan end -function _mapreduce(f, op, src, backend, init, neutral, acctype, dims, alg, temp) - backend = _resolve_backend(backend, src, temp) +_plan(::typeof(reduce), op, src::AbstractArray; kwargs...) = + _plan(mapreduce, identity, op, src; kwargs...) + +# The source of a reduction: several arrays become one `Broadcasted` object +_mapreduce_fuse(f, src, ::Tuple{}) = (f, src) +function _mapreduce_fuse(f, src, srcs::Tuple) + src isa AbstractArray || + throw(ArgumentError("a Broadcasted source cannot be combined with more arrays")) + _mapreduce_check_map_axes(src, srcs...) + bc = Base.Broadcast.instantiate(Base.Broadcast.broadcasted(f, src, srcs...)) + return identity, bc +end +# Scalar *linear* indexing into a multidimensional `Broadcasted` object is only available on +# Julia 1.12, so earlier versions materialize it when the reduction runs (after the plan and the +# workspace checks, which see the original arrays); the plan treats it as a dense array. +_mapreduce_source(src) = src +_mapreduce_source(src::Base.Broadcast.Broadcasted) = + VERSION < v"1.12-" ? _materialize_source(src) : src + +# The seed of the partial results (`_reduce_seed`), where an accumulator type exists +_mapreduce_seed(op, ::Type{A}, neutral) where {A} = + A === Union{} ? nothing : _reduce_seed(op, A, neutral) + +# Everything a reduction resolves before touching data: its backend, algorithm and scratch (the +# plan), accumulator type and partial-result seed, and for reductions along `dims` the result +function _mapreduce_setup(f, op, src, backend, init, neutral, acctype, dims, alg) + backend = _resolve_backend(backend, src) M = _mapped_eltype(f, src) + # The accumulator type, from `init`'s type + A = _acctype(op, init isa _NoInit ? Union{} : typeof(init), M, acctype) if _whole(dims) - # Scalar reductions accumulate from `init`'s type - A = _acctype(op, init isa _NoInit ? Union{} : typeof(init), M, acctype) a = _resolve_reduce(alg, backend, A, dims) - return _mapreduce_whole(f, op, src, backend, a, A; init, neutral, temp) + seed = _mapreduce_seed(op, A, neutral) + # (the buffer's name does not depend on the length, so that results infer) + sizes = if a isa BlockReduce && seed !== nothing + n = _mapreduce_1d_partials(length(src), a.block_size, a.items_per_thread, + a.switch_below) + (; partials=_buffer(typeof(seed), n)) + else + (;) + end + return (; plan=_Plan(backend, a, sizes), A=Val(A), seed) end - temp === nothing || throw(ArgumentError( - "`temp` only applies to whole-array reductions; to reduce into an existing array, " * - "use `mapreducedim!`")) dims_valid = _reduced_dims(dims, ndims(src)) + # The result has the accumulator type; where there is none (`op` always throws), only an + # empty reduction can succeed + R_type = A !== Union{} ? A : !(init isa _NoInit) ? typeof(init) : M === Union{} ? Nothing : M dst_sizes = ntuple(d -> d in dims_valid ? 1 : _srcsize(src, d), ndims(src)) - if init isa _NoInit && Base.any(d -> _srcsize(src, d) == 0, dims_valid) && - Base.prod(dst_sizes) > 0 + a = _resolve_reduce(alg, backend, A, dims_valid) + seed = _mapreduce_seed(op, A, neutral) + sizes = _mapreduce_nd_sizes(src, backend, a, A, seed, dims_valid, Base.prod(dst_sizes)) + return (; plan=_Plan(backend, a, sizes), A=Val(A), seed, dims_valid, R_type=Val(R_type), + dst_sizes) +end + +# Run the reduction set up by `_mapreduce_setup`, with the plan's scratch buffers `bufs` +function _mapreduce_run(f, op, src, s, init, bufs) + backend, a = s.plan.backend, s.plan.alg + if !haskey(s, :dims_valid) + return _mapreduce_whole(f, op, src, backend, a, _unval(s.A); init, neutral=s.seed, + partials=get(bufs, :partials, nothing)) + end + if init isa _NoInit && Base.any(d -> _srcsize(src, d) == 0, s.dims_valid) && + Base.prod(s.dst_sizes) > 0 throw(ArgumentError( "reducing over an empty dimension is not allowed without `init`; pass `init`")) end - # The result has the accumulator type; where there is none (`op` always throws), only an - # empty reduction can succeed - A = _acctype(op, init isa _NoInit ? Union{} : typeof(init), M, acctype) - R_type = A !== Union{} ? A : !(init isa _NoInit) ? typeof(init) : M === Union{} ? Nothing : M - dst = KernelAbstractions.allocate(backend, R_type, dst_sizes) - return _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, A; - init, neutral, overwrite=true) + dst = KernelAbstractions.allocate(backend, _unval(s.R_type), s.dst_sizes) + return mapreduce_nd!(f, op, dst, src, backend, a, _unval(s.A); + init, neutral=s.seed, dims_valid=s.dims_valid, bufs) end +# A reduction nested in another operation, with the scratch buffers of its plan in the outer one +_mapreduce_nested(f, op, src, bufs; backend, init=_NoInit(), neutral=nothing, dims=:, alg) = + _mapreduce_run(f, op, src, + _mapreduce_setup(f, op, src, backend, init, neutral, nothing, dims, alg), + init, bufs) + """ mapreducedim!( @@ -174,6 +230,7 @@ end overwrite::Bool=false, acctype=nothing, alg::Algorithm=Auto(), + workspace=nothing, ) -> R Reduce `src` (an array or a `Broadcasted` object) into `R`: the dimensions along which `R` has @@ -213,6 +270,8 @@ Base's (see [Differences from Base](@ref)): non-bits accumulator type for a kernel algorithm, and an algorithm that cannot run on the backend are `ArgumentError`s. +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call. + ```julia import AcceleratedKernels as AK using CUDA @@ -232,7 +291,21 @@ function mapreducedim!( overwrite::Bool=false, acctype=nothing, alg::Algorithm=Auto(), + workspace=nothing, ) + s = _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg) + bufs = _buffers(s.plan, workspace, R, src) + mapreduce_nd!(f, op, s.dst, _mapreduce_source(s.src), s.plan.backend, s.plan.alg, _unval(s.A); + init=s.init, neutral=s.seed, dims_valid=s.dims_valid, bufs) + return R +end + +_plan(::typeof(mapreducedim!), f, op, R::AbstractArray, src::MapReduceSource; + backend=nothing, init=_NoInit(), neutral=nothing, overwrite::Bool=false, acctype=nothing, + alg::Algorithm=Auto()) = + _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg).plan + +function _mapreducedim_setup(f, op, R, src, backend, init, neutral, overwrite, acctype, alg) backend = _resolve_backend(backend, R, src) nd = ndims(src) for d in 1:max(nd, ndims(R)) @@ -245,8 +318,11 @@ function mapreducedim!( dst = size(R) == dst_sizes ? R : reshape(R, dst_sizes) dims_valid = Tuple(d for d in 1:nd if dst_sizes[d] == 1 && _srcsize(src, d) != 1) A = _acctype(op, eltype(dst), _mapped_eltype(f, src), acctype) - _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, A; init, neutral, overwrite) - return R + a = _resolve_reduce(alg, backend, A, dims_valid) + seed = _mapreduce_seed(op, A, neutral) + sizes = _mapreduce_nd_sizes(src, backend, a, A, seed, dims_valid, length(dst)) + init = init isa _NoInit && !overwrite ? _Fold() : init + return (; plan=_Plan(backend, a, sizes), src, dst, A=Val(A), seed, dims_valid, init) end function _check_noalias(R, src::AbstractArray) @@ -258,23 +334,9 @@ _check_noalias(R, src::Base.Broadcast.Broadcasted) = foreach(a -> _check_noalias _check_noalias(R, src::Base.Broadcast.Extruded) = _check_noalias(R, src.x) _check_noalias(R, src) = nothing -# Reduce `src` into `dst`, which has `ndims(src)` dimensions and size 1 along `dims_valid`, with -# accumulator type `A`. -function _mapreducedim!(f, op, dst, src, backend, alg, dims_valid, ::Type{A}; - init, neutral, overwrite) where {A} - # scalar *linear* indexing into a multidimensional Broadcasted object is - # only available on Julia 1.12; on earlier version, materialize it first. - if VERSION < v"1.12-" && src isa Base.Broadcast.Broadcasted - src = _materialize_source(src) - end - a = _resolve_reduce(alg, backend, A, dims_valid) - init = init isa _NoInit && !overwrite ? _Fold() : init - mapreduce_nd!(f, op, dst, src, backend, a, A; init, neutral, dims_valid) - return dst -end - -# Reduce all of `src` to a host value, with accumulator type `A`. -function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, temp) where {A} +# Reduce all of `src` to a host value, with accumulator type `A` and partial-result seed +# `neutral`; `partials` is the plan's scratch +function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, partials) where {A} if length(src) == 0 init isa _NoInit || return init throw(ArgumentError( @@ -285,25 +347,18 @@ function _mapreduce_whole(f, op, src, backend, alg, ::Type{A}; init, neutral, te return @allowscalar Base.mapreduce_first(f, op, src[first(eachindex(src))]) end _check_acctype(op, f, A) - neutral = _reduce_seed(op, A, neutral) - - # scalar *linear* indexing into a multidimensional Broadcasted object is - # only available on Julia 1.12; on earlier version, materialize it first. - if VERSION < v"1.12-" && src isa Base.Broadcast.Broadcasted - src = _materialize_source(src) - end # The result, `op(init, partial)` or `partial`, has the accumulator type - return convert(A, _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, temp)) + return convert(A, _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, partials)) end -function _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, temp) +function _mapreduce_whole_run(f, op, src, backend, alg; init, neutral, partials) if alg isa BlockReduce mapreduce_1d_gpu( f, op, src, backend; init, neutral, block_size=alg.block_size, items_per_thread=alg.items_per_thread, - temp, switch_below=alg.switch_below, + partials, switch_below=alg.switch_below, ) else mapreduce_1d_cpu( diff --git a/src/sort/cpu_sample_sort.jl b/src/sort/cpu_sample_sort.jl index 8c3cb38d..5c8758ad 100644 --- a/src/sort/cpu_sample_sort.jl +++ b/src/sort/cpu_sample_sort.jl @@ -199,7 +199,9 @@ function _sample_sort!( end max_tasks = min(max_tasks, num_elements ÷ min_elems) if max_tasks <= 1 || num_elements < oversampling_factor * max_tasks - return Base.sort!(v; lt, by, rev, order) + # `temp` (a workspace's) is Base's scratch then + return temp isa Vector{eltype(v)} && length(temp) >= num_elements ? + Base.sort!(v; lt, by, rev, order, scratch=temp) : Base.sort!(v; lt, by, rev, order) end # Create a temporary buffer for the sorted output @@ -306,26 +308,22 @@ end # Key/value sort on Julia threads: find the stable sorting permutation of the keys (of each slice # along `dims`), then gather keys and values through it. -function _sample_sort_by_key!(keys, values, ord, dims; max_tasks, min_elems, temp_keys, temp_values) - for (v, temp) in ((keys, temp_keys), (values, temp_values)) - isnothing(temp) && continue - @argcheck length(temp) == length(v) - @argcheck eltype(temp) === eltype(v) - end - ix = similar(keys, Int) +function _sample_sort_by_key!(keys, values, ord, dims, bufs; max_tasks, min_elems) + ix = bufs.ix if dims isa Colon - _sample_sortperm!(vec(ix), vec(keys); order=ord, max_tasks, min_elems) + _sample_sortperm!(vec(ix), vec(keys); order=ord, max_tasks, min_elems, + temp=get(bufs.perm, :temp, nothing)) else _sample_sortperm_dims!(ix, keys, ord, dims; max_tasks, min_elems) end - _gather_through!(keys, ix, temp_keys; max_tasks, min_elems) - _gather_through!(values, ix, temp_values; max_tasks, min_elems) + _gather_through!(keys, ix, bufs.keys; max_tasks, min_elems) + _gather_through!(values, ix, bufs.values; max_tasks, min_elems) keys, values end # v[i] = v_old[ix[i]] for every linear index i -function _gather_through!(v, ix, temp; max_tasks, min_elems) - old = isnothing(temp) ? copy(v) : copyto!(temp, v) +function _gather_through!(v, ix, old; max_tasks, min_elems) + copyto!(old, v) task_partition(length(v), max_tasks, min_elems) do irange @inbounds for i in irange v[i] = old[ix[i]] @@ -333,3 +331,8 @@ function _gather_through!(v, ix, temp; max_tasks, min_elems) end v end + + +# The scratch of `_sample_sort!` over `n` elements of type `T`: the sorted output when the sort is +# parallel, `Base.sort!`'s scratch when it is not +_sample_sort_sizes(a::CPUThreads.SampleSort, ::Type{T}, n) where {T} = (; temp=_buffer(T, n)) diff --git a/src/sort/merge_sort.jl b/src/sort/merge_sort.jl index 83705328..b7603e72 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -128,9 +128,10 @@ end end -# GPU merge sort of `v` (or of each slice along `dims`), in place. +# GPU merge sort of `v` (or of each slice along `dims`), in place, with the scratch buffers of +# `_sort_sizes(::MergeSort, ...)`. function _merge_sort!( - v::AbstractArray, backend::Backend=get_backend(v); + v::AbstractArray, backend::Backend, bufs::NamedTuple; lt=isless, by=identity, @@ -138,29 +139,20 @@ function _merge_sort!( order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @argcheck block_size > 0 layout = slice_layout(v, dims) ord = Base.Order.ord(lt, by, rev, order) - if !isnothing(temp) - @argcheck length(temp) == length(v) - @argcheck eltype(temp) === eltype(v) - end (isempty(v) || layout.len <= 1) && return v # Compute keys once instead of evaluating `by` in every comparison, into an array of the # keys' type (broadcasting would make a `BitArray` of `Bool` keys on the host) if by !== identity - keys = similar(v, Base.promote_op(by, eltype(v))) + keys = bufs.keys map!(by, keys, v; backend, block_size) - _merge_sort_by_key!( - keys, v, backend; - lt, rev, order, block_size, dims, - temp_values=temp, # temp was for v swap buffer; maps to temp_values here - ) + _merge_sort_by_key!(keys, v, backend, bufs; lt, rev, order, block_size, dims) return v end @@ -180,7 +172,7 @@ function _merge_sort!( size_group = half_size_group * 2 if len > half_size_group p1 = v - p2 = isnothing(temp) ? similar(v) : temp + p2 = bufs.temp kernel! = _merge_sort_global!(backend, block_size) diff --git a/src/sort/merge_sort_by_key.jl b/src/sort/merge_sort_by_key.jl index 844c1ad1..6f3a3c8c 100644 --- a/src/sort/merge_sort_by_key.jl +++ b/src/sort/merge_sort_by_key.jl @@ -162,10 +162,13 @@ end # GPU merge sort of `keys` (or of each slice along `dims`), permuting `values` alike. +# GPU merge sort of `keys` (or of each slice along `dims`) carrying `values` along, in place, with +# the scratch buffers of `_merge_by_key_sizes`. function _merge_sort_by_key!( keys::AbstractArray, values::AbstractArray, - backend::Backend=get_backend(keys); + backend::Backend, + bufs::NamedTuple; lt=isless, by=identity, @@ -173,8 +176,6 @@ function _merge_sort_by_key!( order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, - temp_keys::Union{Nothing, AbstractArray}=nothing, - temp_values::Union{Nothing, AbstractArray}=nothing, dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @@ -184,14 +185,6 @@ function _merge_sort_by_key!( if !(dims isa Colon) @argcheck axes(keys) == axes(values) end - if !isnothing(temp_keys) - @argcheck length(temp_keys) == length(keys) - @argcheck eltype(temp_keys) === eltype(keys) - end - if !isnothing(temp_values) - @argcheck length(temp_values) == length(values) - @argcheck eltype(temp_values) === eltype(values) - end # Construct comparator ord = Base.Order.ord(lt, by, rev, order) @@ -211,10 +204,10 @@ function _merge_sort_by_key!( size_group = half_size_group * 2 if len > half_size_group pk1 = keys - pk2 = isnothing(temp_keys) ? similar(keys) : temp_keys + pk2 = bufs.temp_keys pv1 = values - pv2 = isnothing(temp_values) ? similar(values) : temp_values + pv2 = bufs.temp_values kernel! = _merge_sort_by_key_global!(backend, block_size) @@ -242,3 +235,10 @@ function _merge_sort_by_key!( keys, values end + + +# The scratch of `_merge_sort_by_key!` over slices `layout` of keys of size `ksize` and type `K`, +# with values of size `vsize` and type `V`: swap buffers, if the tiles need merging +_merge_by_key_sizes(a, layout, ksize, ::Type{K}, vsize, ::Type{V}) where {K, V} = + layout.len > 2 * a.block_size ? + (; temp_keys=_buffer(K, ksize), temp_values=_buffer(V, vsize)) : (;) diff --git a/src/sort/merge_sortperm.jl b/src/sort/merge_sortperm.jl index 88c8d570..71acb05e 100644 --- a/src/sort/merge_sortperm.jl +++ b/src/sort/merge_sortperm.jl @@ -2,17 +2,15 @@ function _merge_sortperm!( ix::AbstractArray, v::AbstractArray, - backend::Backend=get_backend(v); + backend::Backend, + bufs::NamedTuple; lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - inplace::Bool=false, block_size::Int=256, - temp_ix::Union{Nothing, AbstractArray}=nothing, - temp_v::Union{Nothing, AbstractArray}=nothing, dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @@ -21,28 +19,16 @@ function _merge_sortperm!( dims isa Colon || @argcheck axes(ix) == axes(v) layout = slice_layout(v, dims) Base.Order.ord(lt, by, rev, order) # validate the ordering keywords before touching ix - if !isnothing(temp_ix) - @argcheck length(temp_ix) == length(ix) - @argcheck eltype(temp_ix) === eltype(ix) - end - - if !isnothing(temp_v) - @argcheck length(temp_v) == length(v) - @argcheck eltype(temp_v) === eltype(v) - end # Initialise the linear indices that will be sorted by the keys in v _foreachindex(eachindex(ix), backend; block_size) do i @inbounds ix[i] = i end (isempty(v) || layout.len <= 1) && return ix - keys = inplace ? v : _copy(backend, v) + # (a range, for instance, has no backend: collected on the host first) + keys = copyto!(bufs.keys, _backend_free(v) ? collect(v) : v) - _merge_sort_by_key!( - keys, ix, backend; - lt, by, rev, order, block_size, dims, - temp_keys=temp_v, temp_values=temp_ix, - ) + _merge_sort_by_key!(keys, ix, backend, bufs; lt, by, rev, order, block_size, dims) ix end @@ -52,7 +38,8 @@ end function _merge_sortperm_lowmem!( ix::AbstractArray, v::AbstractArray, - backend::Backend=get_backend(v); + backend::Backend, + bufs::NamedTuple; lt=isless, by=identity, @@ -60,7 +47,6 @@ function _merge_sortperm_lowmem!( order::Base.Order.Ordering=Base.Order.Forward, block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, dims::Union{Colon, Integer}=Colon(), ) # Simple sanity checks @@ -69,10 +55,6 @@ function _merge_sortperm_lowmem!( dims isa Colon || @argcheck axes(ix) == axes(v) layout = slice_layout(v, dims) ord = Base.Order.ord(lt, by, rev, order) - if !isnothing(temp) - @argcheck length(temp) == length(ix) - @argcheck eltype(temp) === eltype(ix) - end # Initialise the linear indices that will be sorted by the keys in v _foreachindex(eachindex(ix), backend; block_size) do i @@ -96,7 +78,7 @@ function _merge_sortperm_lowmem!( size_group = half_size_group * 2 if len > half_size_group p1 = ix - p2 = isnothing(temp) ? similar(ix) : temp + p2 = bufs.temp kernel! = _merge_sort_global!(backend, block_size) diff --git a/src/sort/radix_sort.jl b/src/sort/radix_sort.jl index a77e6ed1..dc2c3380 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -366,18 +366,23 @@ end end -# Return the extrema of the transformed sort keys. -function _rs_key_range(v::AbstractArray{T}, backend::Backend, descending::Bool) where T +# The extrema of the transformed sort keys, as a reduction over `(key, key)` pairs +struct _RSKeyPair end +@inline (::_RSKeyPair)(x) = (k = _to_sort_key(x); (k, k)) +struct _RSMinMax end +@inline (::_RSMinMax)(a, b) = (min(a[1], b[1]), max(a[2], b[2])) + +function _rs_key_range_setup(v::AbstractArray{T}, backend::Backend) where T K = typeof(_to_sort_key(zero(T))) ident = (typemax(K), typemin(K)) - min_k, max_k = mapreduce( - x -> (k = _to_sort_key(x); (k, k)), - (a, b) -> (min(a[1], b[1]), max(a[2], b[2])), - v; - backend, - init=ident, - neutral=ident, - ) + return _mapreduce_setup(_RSKeyPair(), _RSMinMax(), v, backend, ident, ident, nothing, :, + Auto()), ident +end + +# Return the extrema of the transformed sort keys, with the scratch of `_rs_key_range_setup` +function _rs_key_range(v::AbstractArray{T}, backend::Backend, descending::Bool, bufs) where T + s, ident = _rs_key_range_setup(v, backend) + min_k, max_k = _mapreduce_run(_RSKeyPair(), _RSMinMax(), v, s, ident, bufs) if descending UInt64(~max_k), UInt64(~min_k) else @@ -386,17 +391,47 @@ function _rs_key_range(v::AbstractArray{T}, backend::Backend, descending::Bool) end +# The launch configuration of `_radix_sort!`: whether it uses the chunked kernels, the items per +# thread they get, and whether a single block sorts everything in local memory +function _rs_config(::Type{T}, n, backend, block_size, items_per_thread) where T + has_atomics = KernelAbstractions.supports_atomics(backend) + use_fast = has_atomics && block_size % _RS_CHUNK == 0 && + _rs_fast_local_memory(T, block_size, items_per_thread) <= LOCAL_MEMORY_BUDGET + items = use_fast ? items_per_thread : 1 + one_block = use_fast && _rs_block_local_memory(T, block_size) <= LOCAL_MEMORY_BUDGET && + n <= 2 * block_size + return (; has_atomics, use_fast, items, one_block) +end + +# The scratch of `_radix_sort!`, and the algorithms of the operations it calls: the output of +# every other pass, the digit histograms of every block, and the histograms' scan and the key +# range reduction +function _radix_plan(a, v::AbstractArray{T}, backend) where T + n = length(v) + n <= 1 && return (;), (;) + c = _rs_config(T, n, backend, a.block_size, a.items_per_thread) + c.one_block && return (;), (;) + num_blocks = cld(n, a.block_size * c.items) + hist_len = Int(_RS_SIZE) * num_blocks + scan = _accumulate_setup(+, UInt32, UInt32, (hist_len,), backend; init=UInt32(0), + inclusive=false).plan + key_range = first(_rs_key_range_setup(v, backend)).plan + return (; temp=_buffer(T, size(v)), hist=_buffer(UInt32, hist_len), scan=scan.sizes, + key_range=key_range.sizes), (; scan=scan.alg, key_range=key_range.alg) +end + + """ - _radix_sort!(v, backend; descending, block_size, temp) + _radix_sort!(v, backend, bufs; descending, block_size, items_per_thread) -In-place GPU radix sort for supported 32- and 64-bit integers and floats. +In-place GPU radix sort for supported 32- and 64-bit integers and floats, with the scratch +buffers of `_radix_plan`. """ function _radix_sort!( - v::AbstractArray{T}, backend::Backend=get_backend(v); + v::AbstractArray{T}, backend::Backend, bufs::NamedTuple; descending::Bool=false, block_size::Int=256, items_per_thread::Int=2, - temp::Union{Nothing, AbstractArray}=nothing, ) where T n = length(v) n <= 1 && return v @@ -405,15 +440,11 @@ function _radix_sort!( @argcheck items_per_thread >= 1 @argcheck _rs_portable_local_memory(T, block_size) <= LOCAL_MEMORY_BUDGET - has_atomics = KernelAbstractions.supports_atomics(backend) - use_fast = has_atomics && block_size % _RS_CHUNK == 0 && - _rs_fast_local_memory(T, block_size, items_per_thread) <= LOCAL_MEMORY_BUDGET - items = use_fast ? items_per_thread : 1 + (; has_atomics, use_fast, items, one_block) = _rs_config(T, n, backend, block_size, items_per_thread) n_passes = sizeof(T) * 8 ÷ Int(_RS_BITS) - if use_fast && _rs_block_local_memory(T, block_size) <= LOCAL_MEMORY_BUDGET && - n <= 2 * block_size + if one_block _radix_sort_block!(backend, block_size)( v, descending, Val(n_passes); ndrange=block_size) KernelAbstractions.synchronize(backend) @@ -421,22 +452,13 @@ function _radix_sort!( end num_blocks = cld(n, block_size * items) - - hist = similar(v, UInt32, Int(_RS_SIZE) * num_blocks) - - acc_temp = similar(v, UInt32, cld(length(hist), 512)) - + hist = bufs.hist p1 = v - p2 = if !isnothing(temp) - @argcheck length(temp) >= n && eltype(temp) === T - temp - else - similar(v) - end + p2 = bufs.temp ndrange = (block_size * num_blocks,) - min_key, max_key = _rs_key_range(p1, backend, descending) + min_key, max_key = _rs_key_range(p1, backend, descending, bufs.key_range) vitems = Val(items) hist_kern! = has_atomics ? @@ -455,7 +477,7 @@ function _radix_sort!( shift32 = UInt32(shift) hist_kern!(hist, p1, shift32, descending, vitems; ndrange) - accumulate!(+, hist; backend, init=UInt32(0), inclusive=false, temp=acc_temp) + _accumulate_nested!(+, hist, bufs.scan; backend, init=UInt32(0), inclusive=false) scat_kern!(p2, p1, hist, shift32, descending, vitems; ndrange) p1, p2 = p2, p1 diff --git a/src/sort/sort.jl b/src/sort/sort.jl index 9d347f7e..76302e4c 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -65,7 +65,7 @@ include("tuning.jl") dims::Union{Colon, Integer}=:, lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) -> v Sort `v` in place. `lt`, `by`, `rev` and `order` are those of `Base.sort!`, and so is the result: @@ -84,8 +84,8 @@ an `ArgumentError`. `backend` is derived from `v`; pass it only for arrays that do not determine their backend. -`temp` is an optional scratch array with the same length and element type as `v`, used by -`MergeSort`, `RadixSort` and, for `dims=:`, `CPUThreads.SampleSort`. +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call, so that the +sort allocates none of its own. # Examples ```julia @@ -112,33 +112,62 @@ function sort!( by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) - backend = _resolve_backend(backend, v, temp) + p, ord = _sort_setup(v; backend, alg, dims, lt, by, rev, order) + _sort_impl!(p.alg, v, p.backend, dims, ord, _buffers(p, workspace, v); lt, by, rev, order) + return v +end + +_plan(::typeof(sort!), v::AbstractArray; kwargs...) = first(_sort_setup(v; kwargs...)) + +function _sort_setup( + v; backend=nothing, alg::Algorithm=Auto(), dims::Union{Colon, Integer}=Colon(), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, +) + backend = _resolve_backend(backend, v) ord = Base.Order.ord(lt, by, rev, order) a = _resolve_sort(alg, backend, v, dims, ord) - _sort_impl!(a, v, backend, dims, ord; lt, by, rev, order, temp) - return v + sizes, nested = _sort_plan(a, v, backend, dims, by) + return _Plan(backend, a, nested, sizes), ord end -function _sort_impl!(a::MergeSort, v, backend, dims, ord; lt, by, rev, order, temp) - _merge_sort!(v, backend; lt, by, rev, order, block_size=a.block_size, temp, dims) +# The scratch of `sort!` with each algorithm, and the algorithms of the operations it calls +_sort_plan(a, v, backend, dims, by) = _sort_sizes(a, v, backend, dims, by), (;) +_sort_plan(a::RadixSort, v, backend, dims, by) = _radix_plan(a, v, backend) + +function _sort_sizes(a::MergeSort, v, backend, dims, by) + layout = slice_layout(v, dims) + (isempty(v) || layout.len <= 1) && return (;) + by === identity || return (; keys=_buffer(Base.promote_op(by, eltype(v)), size(v)), + _merge_by_key_sizes(a, layout, size(v), Base.promote_op(by, eltype(v)), + size(v), eltype(v))...) + return layout.len > 2 * a.block_size ? (; temp=_buffer(eltype(v), size(v))) : (;) end +_sort_sizes(a::BitonicSort, v, backend, dims, by) = (;) +_sort_sizes(a::CPUThreads.SampleSort, v, backend, dims, by) = + dims isa Colon ? _sample_sort_sizes(a, eltype(v), length(v)) : (;) -function _sort_impl!(a::RadixSort, v, backend, dims, ord; lt, by, rev, order, temp) - _radix_sort!(v, backend; descending=ord === Base.Order.Reverse, - block_size=a.block_size, items_per_thread=a.items_per_thread, temp) +function _sort_impl!(a::MergeSort, v, backend, dims, ord, bufs; lt, by, rev, order) + _merge_sort!(v, backend, bufs; lt, by, rev, order, block_size=a.block_size, dims) end -function _sort_impl!(a::BitonicSort, v, backend, dims, ord; lt, by, rev, order, temp) +function _sort_impl!(a::RadixSort, v, backend, dims, ord, bufs; lt, by, rev, order) + _radix_sort!(v, backend, bufs; descending=ord === Base.Order.Reverse, + block_size=a.block_size, items_per_thread=a.items_per_thread) +end + +function _sort_impl!(a::BitonicSort, v, backend, dims, ord, bufs; lt, by, rev, order) _bitonic_sort!(v, backend; lt, by, rev, order, dims, block_size=a.block_size, items_per_thread=a.items_per_thread) end -function _sort_impl!(a::CPUThreads.SampleSort, v, backend, dims, ord; lt, by, rev, order, temp) +function _sort_impl!(a::CPUThreads.SampleSort, v, backend, dims, ord, bufs; lt, by, rev, order) if dims isa Colon # `vec`: the local sorts use `Base.sort!`, which needs `dims` for other arrays - _sample_sort!(vec(v); lt, by, rev, order, max_tasks=a.max_tasks, min_elems=a.min_elems, temp) + _sample_sort!(vec(v); lt, by, rev, order, max_tasks=a.max_tasks, min_elems=a.min_elems, + temp=get(bufs, :temp, nothing)) else foreach_slice(v, dims; max_tasks=a.max_tasks, min_elems=a.min_elems) do slice Base.sort!(slice; order=ord) @@ -152,11 +181,16 @@ end Out-of-place [`sort!`](@ref): sort a copy of `v`, with the same keywords. """ -function sort(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kwargs...) +function sort(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, workspace=nothing, + kwargs...) backend = _resolve_backend(backend, v) - return sort!(_copy(backend, v); backend, kwargs...) + # The workspace must not alias the input either + workspace === nothing || _buffers(_plan(sort!, v; backend, kwargs...), workspace, v) + return sort!(_copy(backend, v); backend, workspace, kwargs...) end +_plan(::typeof(sort), v::AbstractArray; kwargs...) = _plan(sort!, v; kwargs...) + """ sortperm!( @@ -167,7 +201,7 @@ end dims::Union{Colon, Integer}=:, lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) -> ix Write into `ix` the stable permutation that sorts `v`, so that `v[ix]` is sorted; `ix` is always @@ -180,7 +214,6 @@ integer `dims`, `ix` must have the same axes as `v` and receives linear indices `Auto()` chooses [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort) on the host and [`MergeSort`](@ref) on GPUs; `MergeSort(lowmem=true)` avoids copying the keys. `RadixSort` and `BitonicSort` have no permutation path. `backend` is derived from `ix` and `v`. -`temp` is an optional scratch array like `ix`. """ function sortperm!( ix::AbstractArray, @@ -192,9 +225,23 @@ function sortperm!( by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, +) + p, ord = _sortperm_setup(ix, v; backend, alg, dims, lt, by, rev, order) + _sortperm_impl!(p.alg, ix, v, p.backend, dims, ord, _buffers(p, workspace, ix, v); + lt, by, rev, order) + return ix +end + +_plan(::typeof(sortperm!), ix::AbstractArray, v::AbstractArray; kwargs...) = + first(_sortperm_setup(ix, v; kwargs...)) + +function _sortperm_setup( + ix, v; backend=nothing, alg::Algorithm=Auto(), dims::Union{Colon, Integer}=Colon(), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, ) - backend = _resolve_backend(backend, ix, v, temp) + backend = _resolve_backend(backend, ix, v) ord = Base.Order.ord(lt, by, rev, order) if dims isa Colon length(ix) == length(v) || throw(ArgumentError( @@ -204,27 +251,38 @@ function sortperm!( "index array must have the same axes as the input, $(axes(ix)) != $(axes(v))")) end a = _resolve_sort(alg, backend, v, dims, ord; perm=true) - _sortperm_impl!(a, ix, v, backend, dims, ord; lt, by, rev, order, temp) - return ix + return _Plan(backend, a, _sortperm_sizes(a, ix, v, dims)), ord +end + +# The scratch of `sortperm!` with each algorithm +function _sortperm_sizes(a::MergeSort, ix, v, dims) + layout = slice_layout(v, dims) + (isempty(v) || layout.len <= 1) && return (;) + a.lowmem && return layout.len > 2 * a.block_size ? (; temp=_buffer(eltype(ix), size(ix))) : (;) + # The keys are sorted in a copy, carrying the indices along + return (; keys=_buffer(eltype(v), size(v)), + _merge_by_key_sizes(a, layout, size(v), eltype(v), size(ix), eltype(ix))...) end +_sortperm_sizes(a::CPUThreads.SampleSort, ix, v, dims) = + dims isa Colon ? _sample_sort_sizes(a, eltype(ix), length(ix)) : (;) -function _sortperm_impl!(a::MergeSort, ix, v, backend, dims, ord; lt, by, rev, order, temp) +function _sortperm_impl!(a::MergeSort, ix, v, backend, dims, ord, bufs; lt, by, rev, order) if a.lowmem - _merge_sortperm_lowmem!(ix, v, backend; lt, by, rev, order, - block_size=a.block_size, temp, dims) + _merge_sortperm_lowmem!(ix, v, backend, bufs; lt, by, rev, order, + block_size=a.block_size, dims) else # Copies keys alongside indices, so comparisons never read global memory; the low-memory # path does two global loads per comparison, O(n log²n) global traffic at large n. - _merge_sortperm!(ix, v, backend; lt, by, rev, order, - block_size=a.block_size, temp_ix=temp, dims) + _merge_sortperm!(ix, v, backend, bufs; lt, by, rev, order, block_size=a.block_size, dims) end end -function _sortperm_impl!(a::CPUThreads.SampleSort, ix, v, backend, dims, ord; - lt, by, rev, order, temp) +function _sortperm_impl!(a::CPUThreads.SampleSort, ix, v, backend, dims, ord, bufs; + lt, by, rev, order) if dims isa Colon _sample_sortperm!(vec(ix), vec(v); lt, by, rev, order, - max_tasks=a.max_tasks, min_elems=a.min_elems, temp) + max_tasks=a.max_tasks, min_elems=a.min_elems, + temp=get(bufs, :temp, nothing)) else _sample_sortperm_dims!(ix, v, ord, dims; max_tasks=a.max_tasks, min_elems=a.min_elems) end @@ -242,6 +300,19 @@ function sortperm(v::AbstractArray; backend::Union{Nothing, Backend}=nothing, kw return sortperm!(_similar(backend, v, Int), v; backend, kwargs...) end +# `ix` is only needed for its element type and shape, which `v` gives here +_plan(::typeof(sortperm), v::AbstractArray; kwargs...) = + first(_sortperm_setup(_IndexShape(v), v; kwargs...)) + +# The element type and shape of `similar(v, Int)`, without allocating it +struct _IndexShape{N, A <: AbstractArray} <: AbstractArray{Int, N} + v::A +end +_IndexShape(v::AbstractArray{T, N}) where {T, N} = _IndexShape{N, typeof(v)}(v) +Base.size(ix::_IndexShape) = size(ix.v) +Base.axes(ix::_IndexShape) = axes(ix.v) +_backend_vote(::_IndexShape) = nothing + """ sort_by_key!( @@ -252,8 +323,7 @@ end dims::Union{Colon, Integer}=:, lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp_keys::Union{Nothing, AbstractArray}=nothing, - temp_values::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, ) -> (keys, values) Sort `keys` in place and apply the same permutation to `values`, stably: values with equal keys @@ -262,8 +332,7 @@ must have as many elements as `keys` (the same axes with an integer `dims`). `Auto()` chooses [`CPUThreads.SampleSort`](@ref AcceleratedKernels.CPUThreads.SampleSort) on the host and [`MergeSort`](@ref) on GPUs, the algorithms that support key/value sorting. `backend` is -derived from `keys` and `values`. `temp_keys` and `temp_values` are optional scratch arrays like -`keys` and `values`. +derived from `keys` and `values`. Thrust, oneDPL and Kokkos call this operation `sort_by_key`, CUB `SortPairs`. @@ -287,10 +356,23 @@ function sort_by_key!( by=identity, rev::Union{Nothing, Bool}=nothing, order::Base.Order.Ordering=Base.Order.Forward, - temp_keys::Union{Nothing, AbstractArray}=nothing, - temp_values::Union{Nothing, AbstractArray}=nothing, + workspace=nothing, +) + p, ord = _sort_by_key_setup(keys, values; backend, alg, dims, lt, by, rev, order) + _sort_by_key_impl!(p.alg, keys, values, p.backend, dims, ord, + _buffers(p, workspace, keys, values); lt, by, rev, order) + return keys, values +end + +_plan(::typeof(sort_by_key!), keys::AbstractArray, values::AbstractArray; kwargs...) = + first(_sort_by_key_setup(keys, values; kwargs...)) + +function _sort_by_key_setup( + keys, values; backend=nothing, alg::Algorithm=Auto(), dims::Union{Colon, Integer}=Colon(), + lt=isless, by=identity, rev::Union{Nothing, Bool}=nothing, + order::Base.Order.Ordering=Base.Order.Forward, ) - backend = _resolve_backend(backend, keys, values, temp_keys, temp_values) + backend = _resolve_backend(backend, keys, values) ord = Base.Order.ord(lt, by, rev, order) if dims isa Colon length(keys) == length(values) || throw(ArgumentError( @@ -300,19 +382,30 @@ function sort_by_key!( "keys and values must have the same axes, $(axes(keys)) != $(axes(values))")) end a = _resolve_sort(alg, backend, keys, dims, ord; pairs=true) - _sort_by_key_impl!(a, keys, values, backend, dims, ord; - lt, by, rev, order, temp_keys, temp_values) - return keys, values + return _Plan(backend, a, _sort_by_key_sizes(a, keys, values, dims)), ord +end + +function _sort_by_key_sizes(a::MergeSort, keys, values, dims) + layout = slice_layout(keys, dims) + (isempty(keys) || layout.len <= 1) && return (;) + return _merge_by_key_sizes(a, layout, size(keys), eltype(keys), size(values), eltype(values)) +end +function _sort_by_key_sizes(a::CPUThreads.SampleSort, keys, values, dims) + # The sorting permutation (sorted by `_sample_sortperm!` for `dims=:`), and the old keys and + # values to gather from + perm = dims isa Colon ? (; perm=_sample_sort_sizes(a, Int, length(keys))) : (;) + return (; ix=_buffer(Int, size(keys)), perm..., keys=_buffer(eltype(keys), size(keys)), + values=_buffer(eltype(values), size(values))) end -function _sort_by_key_impl!(a::MergeSort, keys, values, backend, dims, ord; - lt, by, rev, order, temp_keys, temp_values) - _merge_sort_by_key!(keys, values, backend; lt, by, rev, order, dims, - block_size=a.block_size, temp_keys, temp_values) +function _sort_by_key_impl!(a::MergeSort, keys, values, backend, dims, ord, bufs; + lt, by, rev, order) + _merge_sort_by_key!(keys, values, backend, bufs; lt, by, rev, order, dims, + block_size=a.block_size) end -function _sort_by_key_impl!(a::CPUThreads.SampleSort, keys, values, backend, dims, ord; - lt, by, rev, order, temp_keys, temp_values) - _sample_sort_by_key!(keys, values, ord, dims; - max_tasks=a.max_tasks, min_elems=a.min_elems, temp_keys, temp_values) +function _sort_by_key_impl!(a::CPUThreads.SampleSort, keys, values, backend, dims, ord, bufs; + lt, by, rev, order) + _sample_sort_by_key!(keys, values, ord, dims, bufs; max_tasks=a.max_tasks, + min_elems=a.min_elems) end diff --git a/src/workspace.jl b/src/workspace.jl new file mode 100644 index 00000000..028f24e6 --- /dev/null +++ b/src/workspace.jl @@ -0,0 +1,153 @@ +# Scratch memory: every operation that needs scratch plans it with `_plan`, which resolves the +# algorithm and lists its buffers. The operation allocates the buffers, or takes them from a +# `Workspace` made from the same plan by `workspace`. + + +# A buffer requirement, with its element type in the type domain so that the allocated buffers, +# and the results computed with them, infer +struct _Buf{T, N} + dims::NTuple{N, Int} +end +_buffer(::Type{T}, dims::Integer...) where {T} = _Buf{T, length(dims)}(Int.(dims)) +_buffer(::Type{T}, dims::Tuple) where {T} = _Buf{T, length(dims)}(Int.(dims)) +Base.show(io::IO, b::_Buf{T}) where {T} = print(io, (T, b.dims)) + +# The public form of the sizes: `(eltype, dims)` pairs +_public_sizes(sizes::NamedTuple) = Base.map(_public_sizes, sizes) +_public_sizes(b::_Buf{T}) where {T} = (T, b.dims) + + +""" + Workspace + +Scratch memory for one operation, made by [`workspace`](@ref) and passed to the operation with +its `workspace` keyword. A workspace records the backend and device it was made for, the +algorithm the operation resolved to, and its buffers; a call whose arguments need a different +algorithm or other buffers throws an `ArgumentError` instead of using it. + +A workspace holds no state between calls, so it can be reused by any number of calls, but not by +calls that may run at the same time (on several tasks or streams). It covers the operations' +device memory, with two exceptions: on the host, the threaded algorithms still allocate small +per-task bookkeeping, and `Base.sort!` its own scratch for the per-task sorts of +`CPUThreads.SampleSort` and for the slices of a sort along `dims`; and before Julia 1.12, a +reduction of several arrays or of a `Broadcasted` object materializes it first. +""" +struct Workspace{B <: Backend, A, N <: NamedTuple, S <: NamedTuple, T <: NamedTuple} + backend::B + device::Int + alg::A + nested::N + sizes::S + buffers::T +end + +function Base.show(io::IO, ws::Workspace) + print(io, "Workspace(", _backend_name(ws.backend), ", ", _algname(ws.alg), ", ", + Base.format_bytes(_workspace_bytes(ws.sizes)), ")") +end + +_workspace_bytes(sizes::NamedTuple) = Base.sum(_workspace_bytes, values(sizes); init=0) +_workspace_bytes(b::_Buf{T}) where {T} = Base.prod(b.dims; init=1) * Base.elsize(Array{T}) + + +# An operation's scratch: its backend, its resolved algorithm (which the operation runs with), +# the resolved algorithms of the operations it calls, and its buffers, as a `NamedTuple` of +# `(eltype, dims)` requirements, or of the `NamedTuple`s of nested operations. +struct _Plan{B <: Backend, A, N <: NamedTuple, S <: NamedTuple} + backend::B + alg::A + nested::N + sizes::S +end +_Plan(backend, alg, sizes::NamedTuple) = _Plan(backend, alg, (;), sizes) + +# `_plan(op, args...; kwargs...)` has one method per operation, taking the operation's arguments +# (without `workspace`); operations without scratch have none. +function _plan end + + + +""" + workspace_size(op, args...; kwargs...) -> NamedTuple + +The scratch buffers the call `op(args...; kwargs...)` needs, as a `NamedTuple` of +`(eltype, dims)` pairs (nested for the operations it calls), without allocating anything. +An operation with an algorithm whose call needs no scratch gives an empty `NamedTuple`; the launch +wrappers (`foreachindex`, `map!`, `reverse!`, the searches) take no workspace. + +```julia +julia> AK.workspace_size(AK.sort!, CuArray(rand(Float32, 10_000)); alg=AK.RadixSort()) +(temp = (Float32, (10000,)), hist = (UInt32, (5120,)), scan = (prefixes = (UInt32, (3,)),), + key_range = (partials = (Tuple{UInt32, UInt32}, (40,)),)) + +julia> AK.workspace_size(AK.sum, CuArray(rand(Float32, 10_000))) +(partials = (Float32, (40,)),) +``` +""" +workspace_size(op, args...; kwargs...) = _public_sizes(_plan(op, args...; kwargs...).sizes) + +""" + workspace(op, args...; kwargs...) -> Workspace + +Allocate the scratch memory of the call `op(args...; kwargs...)`, to pass to that operation with +the same arguments, or others that resolve to the same algorithm and need the same buffers, +through its `workspace` keyword. The operation then allocates no scratch memory of its own; it +still allocates its result where it returns a new array (`sort`, `findall`, `mapreduce` along +`dims`, ...), and compiles kernels as usual. + +```julia +import AcceleratedKernels as AK +using CUDA + +v = CuArray(rand(Float32, 1_000_000)) +ws = AK.workspace(AK.sort!, v) +for _ in 1:10 + rand!(v) + AK.sort!(v; workspace=ws) # no scratch allocations +end +``` + +The workspace is checked on every call: a different backend or device, a different resolved +algorithm (`Auto()` may choose another one for another length), also for the operations it calls, +or different buffer sizes are an `ArgumentError`, and so is a workspace whose buffers alias the +operation's arrays. +""" +function workspace(op, args...; kwargs...) + p = _plan(op, args...; kwargs...) + return Workspace(p.backend, KernelAbstractions.device(p.backend), p.alg, p.nested, p.sizes, + _allocate(p.backend, p.sizes)) +end + +_allocate(backend, sizes::NamedTuple) = Base.map(s -> _allocate(backend, s), sizes) +_allocate(backend, b::_Buf{T}) where {T} = KernelAbstractions.allocate(backend, T, b.dims) + +_describe(sizes::NamedTuple) = isempty(sizes) ? "no buffers" : "the buffers $sizes" + +# The buffers of the plan: the workspace's, checked against the plan and the operation's arrays, +# or new ones +_buffers(p::_Plan, ::Nothing, arrays...) = _allocate(p.backend, p.sizes) +function _buffers(p::_Plan, ws::Workspace, arrays...) + ws.backend == p.backend && ws.device == KernelAbstractions.device(p.backend) || + throw(ArgumentError("the workspace was made for another backend or device")) + ws.alg == p.alg && ws.nested == p.nested || throw(ArgumentError( + "the workspace was made for $(ws.alg) $(ws.nested), but this call resolves to " * + "$(p.alg) $(p.nested)")) + ws.sizes == p.sizes || throw(ArgumentError( + "the workspace holds $(_describe(ws.sizes)), but this call needs $(_describe(p.sizes))")) + for b in _leaves(ws.buffers), a in arrays + _aliases(b, a) && + throw(ArgumentError("the workspace must not alias the operation's arrays")) + end + return ws.buffers +end + +# Whether the buffer `b` may alias the array or `Broadcasted` source `a` +_aliases(b, a::AbstractArray) = Base.mightalias(b, a) +_aliases(b, a::Base.Broadcast.Broadcasted) = Base.any(x -> _aliases(b, x), a.args) +_aliases(b, a::Base.Broadcast.Extruded) = _aliases(b, a.x) +_aliases(b, a) = false +_buffers(p::_Plan, ws, arrays...) = + throw(ArgumentError("`workspace` must be a `Workspace` from `workspace`, got $(typeof(ws))")) + +_leaves(buffers::NamedTuple) = Iterators.flatten(Base.map(_leaves, values(buffers))) +_leaves(buffer::AbstractArray) = (buffer,) diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index f404479a..dd72d436 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -136,16 +136,13 @@ end @test_throws MethodError AK.accumulate(+, y; init=10, dims=2, inclusive=false, bad=:kwarg, alg=slice_alg()) - # Oversized temporaries are allowed. + # A workspace, reused y = array_from_host(ones(Int32, 1000)) - AK.accumulate!(+, y; init=Int32(0), inclusive=false, temp=array_from_host(zeros(Int32, 1000)), - temp_flags=array_from_host(zeros(Int8, 1000)), alg=scan_alg(A; block_size=128)) + with_workspace(AK.accumulate!, +, y; init=Int32(0), inclusive=false, alg=scan_alg(A; block_size=128)) @test Array(y) == 0:999 - y = AK.accumulate(+, array_from_host(ones(Int32, 1000)); init=0, inclusive=false, - temp=array_from_host(zeros(Int64, 1000)), - temp_flags=array_from_host(zeros(Int8, 1000)), - alg=scan_alg(A; block_size=128)) + y = with_workspace(AK.accumulate, +, array_from_host(ones(Int32, 1000)); + init=0, inclusive=false, alg=scan_alg(A; block_size=128)) @test Array(y) == 0:999 # Cross-block coherence: small tiles (block_size 16-64, 1 item/thread) maximise the number of @@ -368,9 +365,9 @@ end neutral=Int32(0), dims=2, alg=slice_alg(; block_size=64)) AK.accumulate((x, y) -> x + 1, array_from_host(rand(Int32, 3, 4, 5)); init=Int32(0), neutral=Int32(0), dims=3, alg=slice_alg(; block_size=64)) - # The temporaries only apply to whole-array scans - @test_throws ArgumentError AK.accumulate(+, array_from_host(rand(Int32, 3, 4)); init=Int32(0), - dims=2, temp=array_from_host(zeros(Int32, 3))) + # The scratch keywords are gone: scratch is a workspace + @test_throws MethodError AK.accumulate(+, array_from_host(rand(Int32, 3, 4)); init=Int32(0), + dims=2, temp=array_from_host(zeros(Int32, 3))) end @testset "cumsum" begin diff --git a/test/generic/findall.jl b/test/generic/findall.jl index 1f92526e..e38a94a5 100644 --- a/test/generic/findall.jl +++ b/test/generic/findall.jl @@ -70,7 +70,7 @@ findall_alg(; block_size=nothing, items_per_thread=nothing, max_tasks=nothing, m if !TEST_KERNELS @test AK.findall(Any[true, false, true]) == findall(Any[true, false, true]) - @test AK.findall(Any[true, false, true]; temp_bools=Vector{Bool}(undef, 3)) == + @test with_workspace(AK.findall, Any[true, false, true]) == findall(Any[true, false, true]) @test_throws ArgumentError AK.findall([1]) @test_throws TypeError AK.findall(Any[true, missing]; @@ -204,24 +204,22 @@ findall_alg(; block_size=nothing, items_per_thread=nothing, max_tasks=nothing, m findall(x -> x > 0.5f0, h) end + # A workspace: the mask of the predicate form and the counts (and their scan) alg = findall_alg(block_size=64, items_per_thread=3, max_tasks=4) - temp = similar(v, Int, max(4, cld(length(v), 64 * 3))) - temp_bools = similar(v, Bool) - @test Array(AK.findall(x -> x > 0.5f0, v; alg, temp, temp_bools)) == + @test Array(with_workspace(AK.findall, x -> x > 0.5f0, v; alg)) == findall(x -> x > 0.5f0, h) + bools = array_from_host(rand(Bool, length(v))) + @test Array(with_workspace(AK.findall, bools; alg)) == findall(Array(bools)) + @test haskey(AK.workspace_size(AK.findall, x -> x > 0.5f0, v; alg), :mask) + @test !haskey(AK.workspace_size(AK.findall, bools; alg), :mask) + # ... made for another call + @test_throws ArgumentError AK.findall(bools; alg, + workspace=AK.workspace(AK.findall, bools[1:1]; alg)) @test_throws ArgumentError AK.findall(v; alg=OtherFindallAlgorithm()) - @test_throws ArgumentError AK.findall(identity, temp_bools; - temp_bools) - @test_throws ArgumentError AK.findall(identity, v; temp_bools=reshape(similar(v, Bool), :, 1)) - if TEST_KERNELS - bools = array_from_host(rand(Bool, length(v))) @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(block_size=192)) @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(items_per_thread=0)) - @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(), - temp=similar(v, Int32, 100)) - @test_throws ArgumentError AK.findall(bools; alg=AK.ScanScatter(), temp=similar(v, Int, 1)) end end end diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index cb7a081c..78b0f507 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -145,14 +145,8 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0) end # Testing different settings - AK.reduce( - (x, y) -> x + 1, - array_from_host(rand(Int32, 10_000)); - alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), - init=Int32(0), - neutral=Int64(0), - temp=array_from_host(zeros(Int64, 10_000)), - ) + with_workspace(AK.reduce, (x, y) -> x + 1, array_from_host(rand(Int32, 10_000)); + alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), init=Int32(0), neutral=Int64(0)) AK.reduce( (x, y) -> x + 1, array_from_host(rand(Int32, 10_000)); @@ -591,18 +585,17 @@ end end # Testing different settings, enforcing change of type between f and op - f(s, temp) = AK.mapreduce( + f(s) = with_workspace( + AK.mapreduce, p -> (p.x, p.y), (a, b) -> (a[1] < b[1] ? a[1] : b[1], a[2] < b[2] ? a[2] : b[2]), s; alg=reduce_alg(block_size=64, switch_below=50, max_tasks=10, min_elems=100), init=(typemax(Float32), typemax(Float32)), neutral=(typemax(Float32), typemax(Float32)), - temp=temp, ) v = array_from_host([Point(rand(Float32), rand(Float32)) for _ in 1:10_042]) - temp = similar(v, Tuple{Float32, Float32}) - f(v, temp) + f(v) # Test that undefined kwargs are not accepted @test_throws MethodError AK.mapreduce(-, +, v; alg=REDUCE_ALG, init=10, bad=:kwarg) diff --git a/test/generic/reduce_resolve.jl b/test/generic/reduce_resolve.jl index 0798a098..0d7f26ff 100644 --- a/test/generic/reduce_resolve.jl +++ b/test/generic/reduce_resolve.jl @@ -115,7 +115,7 @@ end zeros(Int32, 4, 4), RRB, AK.BlockReduce(256, 2, 0), Int32; init=AK._NoInit(), neutral=nothing, - dims_valid=(1,)) + dims_valid=(1,), bufs=(;)) end end diff --git a/test/generic/sort.jl b/test/generic/sort.jl index 28593fc0..8b1eaecd 100644 --- a/test/generic/sort.jl +++ b/test/generic/sort.jl @@ -36,23 +36,19 @@ if TEST_KERNELS # Testing different settings v = array_from_host(1:10_000, Float32) - AK.sort!(v, lt=(>), by=abs, rev=true, - alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Float32)) + with_workspace(AK.sort!, v; lt=(>), by=abs, rev=true, alg=AK.MergeSort(block_size=64)) @test issorted(Array(v)) v = array_from_host(1:10_000, Int32) - AK.sort!(v, lt=(>), rev=true, - alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sort!, v; lt=(>), rev=true, alg=AK.MergeSort(block_size=64)) @test issorted(Array(v)) v = array_from_host(1:10_000, Float32) - v = AK.sort(v, lt=(>), by=abs, rev=true, - alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Float32)) + v = with_workspace(AK.sort, v; lt=(>), by=abs, rev=true, alg=AK.MergeSort(block_size=64)) @test issorted(Array(v)) v = array_from_host(1:10_000, Int32) - v = AK.sort(v, lt=(>), by=abs, rev=true, - alg=AK.MergeSort(block_size=64), temp=array_from_host(1:10_000, Int32)) + v = with_workspace(AK.sort, v; lt=(>), by=abs, rev=true, alg=AK.MergeSort(block_size=64)) @test issorted(Array(v)) end @@ -94,13 +90,12 @@ end @test Array(tmp) == sort(v_h; by=abs) end - # temp kwarg still forwarded correctly through hoisting path + # a workspace also serves the path that caches the `by` keys n = 20_000 v_h = randn(Float32, n) v = array_from_host(v_h) tmp = copy(v) - temp = array_from_host(zeros(Float32, n)) - AK.sort!(tmp; alg=AK.MergeSort(), by=abs, temp) + with_workspace(AK.sort!, tmp; alg=AK.MergeSort(), by=abs) @test Array(tmp) == sort(v_h; by=abs) # sort! (public API) routes through the same hoisting path @@ -186,13 +181,12 @@ if AK._runs_threads(BACKEND) # Testing different settings v = array_from_host(rand(1:100_000, 10_000), Float32) - AK.sort!(v, lt=(>), by=abs, rev=true, - alg=AK.CPUThreads.SampleSort(max_tasks=64), temp=array_from_host(1:10_000, Float32)) + with_workspace(AK.sort!, v; + lt=(>), by=abs, rev=true, alg=AK.CPUThreads.SampleSort(max_tasks=64)) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - AK.sort!(v, lt=(>), rev=true, - alg=AK.CPUThreads.SampleSort(max_tasks=64), temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sort!, v; lt=(>), rev=true, alg=AK.CPUThreads.SampleSort(max_tasks=64)) @test issorted(Array(v)) end end @@ -228,23 +222,19 @@ end # Testing different settings v = array_from_host(rand(1:100_000, 10_000), Float32) - AK.sort!(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, - temp=array_from_host(1:10_000, Float32)) + with_workspace(AK.sort!, v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - AK.sort!(v; alg=SETTINGS_ALG, lt=(>), rev=true, - temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sort!, v; alg=SETTINGS_ALG, lt=(>), rev=true) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Float32) - v = AK.sort(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, - temp=array_from_host(1:10_000, Float32)) + v = with_workspace(AK.sort, v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true) @test issorted(Array(v)) v = array_from_host(rand(1:100_000, 10_000), Int32) - v = AK.sort(v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true, - temp=array_from_host(1:10_000, Int32)) + v = with_workspace(AK.sort, v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true) @test issorted(Array(v)) end @@ -279,8 +269,7 @@ end for alg in (AK.MergeSort(), AK.MergeSort(lowmem=true)) v = array_from_host(perm_h) ix = array_from_host(zeros(Int, length(perm_h))) - temp = array_from_host(zeros(Int, length(perm_h))) - AK.sortperm!(ix, v; alg, temp) + with_workspace(AK.sortperm!, ix, v; alg) @test is_valid_perm(perm_h, Int.(Array(ix))) end @@ -358,21 +347,13 @@ end # Testing different settings k = array_from_host(1:10_000, Float32) v = array_from_host(1:10_000, Int32) - AK.sort_by_key!(k, v, - lt=(>), by=abs, rev=true, - alg=SETTINGS_ALG, - temp_keys=array_from_host(1:10_000, Float32), - temp_values=array_from_host(1:10_000, Int32)) + with_workspace(AK.sort_by_key!, k, v; lt=(>), by=abs, rev=true, alg=SETTINGS_ALG) @test issorted(Array(k)) @test issorted(Array(v)) k = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sort_by_key!(k, v, - lt=(>), by=abs, rev=true, - alg=SETTINGS_ALG, - temp_keys=array_from_host(1:10_000, Int32), - temp_values=array_from_host(1:10_000, Float32)) + with_workspace(AK.sort_by_key!, k, v; lt=(>), by=abs, rev=true, alg=SETTINGS_ALG) @test issorted(Array(k)) @test issorted(Array(v)) @@ -417,7 +398,12 @@ end kh = Int32[2, 1] k = array_from_host(kh) v = array_from_host(Int32[20, 10]) - @test_throws ArgumentError AK.sort_by_key!(k, v; alg=SORT_ALG, temp_values=array_from_host(zeros(Int32, 1))) + # The scratch keywords are gone: scratch is a workspace + @test_throws MethodError AK.sort_by_key!(k, v; alg=SORT_ALG, temp_values=array_from_host(zeros(Int32, 1))) + k2, v2 = array_from_host(rand(Int32, 2000)), array_from_host(rand(Int32, 2000)) + @test_throws ArgumentError AK.sort_by_key!(k2, v2; alg=SORT_ALG, workspace=AK.workspace( + AK.sort_by_key!, array_from_host(rand(Int32, 4000)), array_from_host(rand(Int32, 4000)); + alg=SORT_ALG)) @test Array(k) == kh # Mismatched sizes @@ -464,20 +450,13 @@ if TEST_KERNELS # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sortperm!(ix, - v, - lt=(>), by=abs, rev=true, - alg=AK.MergeSort(block_size=64), - temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sortperm!, ix, v; lt=(>), by=abs, rev=true, alg=AK.MergeSort(block_size=64)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) v = array_from_host(1:10_000, Float32) - ix = AK.sortperm(v, - lt=(>), by=abs, rev=true, - alg=AK.MergeSort(block_size=64), - temp=array_from_host(1:10_000, Int)) + ix = with_workspace(AK.sortperm, v; lt=(>), by=abs, rev=true, alg=AK.MergeSort(block_size=64)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -524,11 +503,8 @@ if AK._runs_threads(BACKEND) # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sortperm!(ix, - v, - lt=(>), by=abs, rev=true, - alg=AK.CPUThreads.SampleSort(max_tasks=64), - temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sortperm!, ix, v; + lt=(>), by=abs, rev=true, alg=AK.CPUThreads.SampleSort(max_tasks=64)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -574,20 +550,15 @@ if TEST_KERNELS # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sortperm!(ix, - v, - lt=(>), by=abs, rev=true, - alg=AK.MergeSort(lowmem=true, block_size=64), - temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sortperm!, ix, v; + lt=(>), by=abs, rev=true, alg=AK.MergeSort(lowmem=true, block_size=64)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) v = array_from_host(1:10_000, Float32) - ix = AK.sortperm(v, - lt=(>), by=abs, rev=true, - alg=AK.MergeSort(lowmem=true, block_size=64), - temp=array_from_host(1:10_000, Int)) + ix = with_workspace(AK.sortperm, v; + lt=(>), by=abs, rev=true, alg=AK.MergeSort(lowmem=true, block_size=64)) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -632,20 +603,13 @@ end # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sortperm!(ix, - v; - alg=SETTINGS_ALG, - lt=(>), by=abs, rev=true, - temp=array_from_host(1:10_000, Int32)) + with_workspace(AK.sortperm!, ix, v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) v = array_from_host(1:10_000, Float32) - ix = AK.sortperm(v; - alg=SETTINGS_ALG, - lt=(>), by=abs, rev=true, - temp=array_from_host(1:10_000, Int)) + ix = with_workspace(AK.sortperm, v; alg=SETTINGS_ALG, lt=(>), by=abs, rev=true) ixh = Array(ix) vh = Array(v) @test issorted(vh[ixh]) @@ -721,16 +685,16 @@ if TEST_KERNELS @test res end - # ── temp kwarg: buffer reuse gives identical result ─────────────────────── + # ── workspace: buffer reuse gives identical result ─────────────────────── n = 20_000 Random.seed!(321) v1 = array_from_host(rand(Float32, n)) v2 = copy(v1) ix1 = array_from_host(zeros(Int, n)) ix2 = array_from_host(zeros(Int, n)) - temp = array_from_host(zeros(Int, n)) - AK.sortperm!(ix1, v1; alg=SORT_ALG, temp) - AK.sortperm!(ix2, v2; alg=SORT_ALG, temp) + ws = AK.workspace(AK.sortperm!, ix1, v1; alg=SORT_ALG) + AK.sortperm!(ix1, v1; alg=SORT_ALG, workspace=ws) + AK.sortperm!(ix2, v2; alg=SORT_ALG, workspace=ws) @test Array(ix1) == Array(ix2) # ── Exact match against Base.sortperm ──────────────────────────────────── @@ -845,12 +809,11 @@ end @test Array(AK.sort!(array_from_host(Int32[42]); alg=AK.RadixSort())) == Int32[42] @test Array(AK.sort!(array_from_host(Int32[2, 1]); alg=AK.RadixSort())) == Int32[1, 2] - # ── temp kwarg: preallocated buffer ─────────────────────────────────── + # ── workspace: preallocated buffers ─────────────────────────────────── n = 50_000 v_h = rand(Float32, n) v = array_from_host(v_h) - temp = similar(v) - AK.sort!(v; alg=AK.RadixSort(), temp) + with_workspace(AK.sort!, v; alg=AK.RadixSort()) @test Array(v) == sort(v_h) # ── Out-of-place ────────────────────────────────────────────────────── @@ -926,13 +889,13 @@ end @test Array(AK.sort(array_from_host(A_h); alg=SORT_ALG, dims=2)) == sort(A_h; dims=2) end - # by, lt, order and temp act on the values within each slice + # by, lt, order and a workspace act on the values within each slice A_h = rand(Float32, 300, 700) A = array_from_host(A_h) @test Array(AK.sort(A; alg=SORT_ALG, dims=1, by=x->-x)) == sort(A_h; dims=1, by=x->-x) @test Array(AK.sort(A; alg=SORT_ALG, dims=2, lt=(>))) == sort(A_h; dims=2, lt=(>)) @test Array(AK.sort(A; alg=SORT_ALG, dims=2, order=Base.Order.Reverse)) == sort(A_h; dims=2, order=Base.Order.Reverse) - @test Array(AK.sort(A; alg=SORT_ALG, dims=2, temp=similar(A))) == sort(A_h; dims=2) + @test Array(with_workspace(AK.sort, A; alg=SORT_ALG, dims=2)) == sort(A_h; dims=2) if TEST_KERNELS @test Array(AK.sort(A; alg=AK.MergeSort(block_size=64), dims=1)) == sort(A_h; dims=1) @test_throws ArgumentError AK.sort(A; dims=1, alg=AK.RadixSort()) @@ -1003,12 +966,12 @@ end @test Array(AK.sortperm(array_from_host(A_h); alg=SORT_ALG, dims=2)) == sortperm(A_h; dims=2) end - # by, order, temp and the low-memory GPU path + # by, order, a workspace and the low-memory GPU path A_h = rand(Float32, 300, 700) A = array_from_host(A_h) @test Array(AK.sortperm(A; alg=SORT_ALG, dims=1, by=x->-x)) == sortperm(A_h; dims=1, by=x->-x) @test Array(AK.sortperm(A; alg=SORT_ALG, dims=2, order=Base.Order.Reverse)) == sortperm(A_h; dims=2, order=Base.Order.Reverse) - @test Array(AK.sortperm(A; alg=SORT_ALG, dims=2, temp=similar(A, Int))) == sortperm(A_h; dims=2) + @test Array(with_workspace(AK.sortperm, A; alg=SORT_ALG, dims=2)) == sortperm(A_h; dims=2) if TEST_KERNELS @test Array(AK.sortperm(A; dims=2, alg=AK.MergeSort(lowmem=true))) == sortperm(A_h; dims=2) @test Array(AK.sortperm(A; dims=1, alg=AK.MergeSort(lowmem=true, block_size=64))) == sortperm(A_h; dims=1) diff --git a/test/generic/sort_resolve.jl b/test/generic/sort_resolve.jl index dc139f54..dbaaf1d9 100644 --- a/test/generic/sort_resolve.jl +++ b/test/generic/sort_resolve.jl @@ -297,7 +297,7 @@ end @test AK._backend_free(reshape(1:6, 2, 3)) && !AK._backend_free(NoBackendTestArray()) # The sorting entry points resolve from all their arrays, and accept only backends - @test_throws ArgumentError AK.sort!(d; temp=h) + @test_throws MethodError AK.sort!(d; temp=h) # no scratch keywords: see `workspace` @test_throws ArgumentError AK.sortperm!(zeros(Int, 4), d) @test_throws ArgumentError AK.sort_by_key!(d, h) @test_throws TypeError AK.sort!(h; backend=:cpu) diff --git a/test/generic/workspace.jl b/test/generic/workspace.jl new file mode 100644 index 00000000..52a18b51 --- /dev/null +++ b/test/generic/workspace.jl @@ -0,0 +1,240 @@ +# The workspace protocol: `workspace_size` and `workspace` plan the same buffers as the operation, +# which then allocates no scratch of its own; a workspace made for another call is rejected. + +@testset "workspace: sizes" begin + v = array_from_host(rand(Float32, 10_000)) + ks = array_from_host(rand(Int32(1):Int32(9), 10_000)) + + # Every operation that can need scratch has a plan; its sizes are `(eltype, dims)` pairs, or + # nested for the operations it calls + function check_sizes(sizes) + @test sizes isa NamedTuple + for s in values(sizes) + s isa NamedTuple ? check_sizes(s) : + @test s isa Tuple{Type, Tuple} && Base.all(d -> d isa Int && d >= 0, s[2]) + end + end + for (op, args, kw) in ( + (AK.sort!, (copy(v),), (;)), (AK.sort, (v,), (;)), (AK.sort!, (copy(v),), (; by=abs)), + (AK.sortperm!, (similar(v, Int), v), (;)), (AK.sortperm, (v,), (;)), + (AK.sort_by_key!, (copy(ks), copy(v)), (;)), + (AK.mapreduce, (abs, +, v), (;)), (AK.reduce, (+, v), (;)), (AK.sum, (v,), (;)), + (AK.prod, (v,), (;)), (AK.maximum, (v,), (;)), (AK.minimum, (v,), (;)), + (AK.count, (x -> x > 0.5f0, v), (;)), (AK.count, (array_from_host(rand(Bool, 100)),), (;)), + (AK.reduce, (+, reshape(v, 100, 100)), (; dims=1)), + (AK.mapreducedim!, (identity, +, array_from_host(zeros(Float32, 1, 100)), reshape(v, 100, 100)), (;)), + (AK.accumulate!, (+, copy(v)), (;)), (AK.accumulate!, (+, similar(v), v), (;)), + (AK.accumulate, (+, v), (;)), (AK.cumsum, (v,), (;)), (AK.cumprod, (v,), (;)), + (AK.accumulate, (+, reshape(v, 100, 100)), (; dims=2)), + (AK.findall, (x -> x > 0.5f0, v), (;)), (AK.findall, (array_from_host(rand(Bool, 100)),), (;)), + (AK.findall, (x -> x > 0.5f0, v), (; items=v)), + (AK.any, (x -> x > 2, v), (;)), (AK.all, (x -> x < 2, v), (;)), + # (the keywords that change the buffers' types) + (AK.sum, (v,), (; acctype=Float32)), (AK.reduce, (+, reshape(v, 100, 100)), (; dims=1, acctype=Float32)), + (AK.mapreducedim!, (identity, +, array_from_host(zeros(Float32, 1, 100)), reshape(v, 100, 100)), (; acctype=Float32)), + (AK.accumulate!, (+, similar(ks), ks), (; acctype=Int)), (AK.accumulate, (+, v), (; acctype=Float32)), + ) + sizes = AK.workspace_size(op, args...; kw...) + check_sizes(sizes) + ws = AK.workspace(op, args...; kw...) + @test ws isa AK.Workspace + @test AK._public_sizes(ws.sizes) == sizes + @test occursin("Workspace(", sprint(show, ws)) + # The operation runs with it, as without it + @test (op(args...; kw..., workspace=ws); true) + end + + if TEST_KERNELS + # Scratch mirrors the algorithm: radix sort's histograms and nested operations, none for + # the bitonic network + sizes = AK.workspace_size(AK.sort!, v; alg=AK.RadixSort()) + @test Base.all(k -> haskey(sizes, k), (:temp, :hist, :scan, :key_range)) + @test sizes.temp == (Float32, (10_000,)) + @test AK.workspace_size(AK.sort!, v; alg=AK.BitonicSort()) == (;) + # A whole-array reduction keeps its partial results in two halves of one buffer + @test haskey(AK.workspace_size(AK.sum, v; alg=AK.BlockReduce()), :partials) + # ... which a reduction of one element does not need + @test AK.workspace_size(AK.sum, v[1:1]; alg=AK.BlockReduce()) == (; partials=(Float32, (0,))) + # DecoupledLookback's flags only where it runs + if TEST_DL + @test haskey(AK.workspace_size(AK.accumulate!, +, copy(v); alg=AK.DecoupledLookback()), :flags) + end + @test !haskey(AK.workspace_size(AK.accumulate!, +, copy(v); alg=AK.ScanPrefixes()), :flags) + # `acctype` sets the partial results' type + m = reshape(v, 10, 1000) + R = array_from_host(zeros(Float32, 1, 1000)) + @test AK.workspace_size(AK.sum, v; alg=AK.BlockReduce(), acctype=Int32).partials[1] === Int32 + @test AK.workspace_size(AK.mapreducedim!, identity, +, array_from_host(zeros(Float32, 1, 1)), + reshape(v, 100, 100); alg=AK.BlockReduce(), + acctype=Int32).partials[1] === Int32 + end + # A scan into a destination of its running type needs no scratch array; one into another needs + # one of the running type, which the destination and `acctype` set + w = array_from_host(rand(Int32(1):Int32(9), 1000)) + @test !haskey(AK.workspace_size(AK.accumulate!, +, similar(w, Int), w), :work) + @test AK.workspace_size(AK.accumulate!, +, similar(w, Int8), w).work == (Int32, (1000,)) + @test AK.workspace_size(AK.accumulate!, +, similar(w), w; acctype=Int16).work == + (Int16, (1000,)) +end + + +@testset "workspace: reuse and checks" begin + Random.seed!(0) + vh = rand(Float32, 10_000) + v = array_from_host(vh) + + # One workspace serves any number of calls with the same plan + ws = AK.workspace(AK.sort!, v) + for _ in 1:3 + w = array_from_host(rand(Float32, 10_000)) + wh = Array(w) + AK.sort!(w; workspace=ws) + @test Array(w) == sort(wh) + end + ws = AK.workspace(AK.sum, v) + @test AK.sum(v; workspace=ws) ≈ sum(vh) + @test AK.sum(v; workspace=ws) ≈ sum(vh) + # ... also by an empty `sum`, and not by one with another accumulator type + e = array_from_host(zeros(Float32, 0, 3)) + @test Array(AK.sum(e; dims=1, workspace=AK.workspace(AK.sum, e; dims=1))) == zeros(1, 3) + @test_throws ArgumentError AK.sum(e; dims=1, workspace=:invalid) + @test_throws ArgumentError AK.prod(e[:, 1]; workspace=:invalid) + if TEST_KERNELS + @test_throws ArgumentError AK.sum(v; alg=AK.BlockReduce(), + workspace=AK.workspace(AK.sum, v; acctype=Int32, + alg=AK.BlockReduce())) + end + + # ... and is rejected by a call with another plan: other buffer sizes, another algorithm, + # another operation's buffers, or no `Workspace` at all + alg = TEST_KERNELS ? AK.MergeSort() : AK.CPUThreads.SampleSort(max_tasks=4) + @test_throws ArgumentError AK.sort!(array_from_host(rand(Float32, 20_000)); alg, + workspace=AK.workspace(AK.sort!, v; alg)) + if TEST_KERNELS + a = AK.BlockReduce() + @test_throws ArgumentError AK.sum(array_from_host(rand(Float32, 50_000)); alg=a, + workspace=AK.workspace(AK.sum, v; alg=a)) + end + @test_throws ArgumentError AK.sort!(copy(v); workspace=AK.workspace(AK.sum, v)) + @test_throws ArgumentError AK.sort!(copy(v); workspace=similar(v)) + if TEST_KERNELS + @test_throws ArgumentError AK.sort!(copy(v); alg=AK.MergeSort(), + workspace=AK.workspace(AK.sort!, v; alg=AK.RadixSort())) + end + + # A workspace must not alias the operation's arrays + alg = TEST_KERNELS ? AK.MergeSort() : AK.CPUThreads.SampleSort(max_tasks=4) + w = copy(v) + ws = AK.workspace(AK.sort!, w; alg) + @test haskey(ws.buffers, :temp) + aliased = AK.Workspace(ws.backend, ws.device, ws.alg, ws.nested, ws.sizes, + merge(ws.buffers, (; temp=w))) + @test_throws ArgumentError AK.sort!(w; alg, workspace=aliased) + + # The scratch keywords are gone + @test_throws MethodError AK.sort!(copy(v); temp=similar(v)) + @test_throws MethodError AK.sum(v; temp=similar(v)) + @test_throws MethodError AK.accumulate!(+, copy(v); temp=similar(v)) + @test_throws MethodError AK.findall(x -> x > 0.5f0, v; temp_bools=similar(v, Bool)) +end + + +@testset "workspace: aliasing, nested algorithms, inference" begin + v = array_from_host(rand(Float32, 10_000)) + + # An allocating sort checks the workspace against its input, not only against the copy + alg = TEST_KERNELS ? AK.MergeSort() : AK.CPUThreads.SampleSort(max_tasks=4) + ws = AK.workspace(AK.sort, v; alg) + w = copy(v) + aliased = AK.Workspace(ws.backend, ws.device, ws.alg, ws.nested, ws.sizes, + merge(ws.buffers, (; temp=w))) + @test_throws ArgumentError AK.sort(w; alg, workspace=aliased) + @test Array(AK.sort(w; alg, workspace=ws)) == sort(Array(w)) + + if TEST_KERNELS + # ... and a reduction against every array of a fused source + x = array_from_host(ones(Int32, 1000)) + y = copy(x) + a = AK.BlockReduce(block_size=2, items_per_thread=1) + ws = AK.workspace(AK.mapreduce, +, +, x, y; alg=a) + z = ws.buffers.partials + @test length(z) == 1000 + @test_throws ArgumentError AK.mapreduce(+, +, z, y; alg=a, workspace=ws) + # ... and a findall against its input + ws = AK.workspace(AK.findall, identity, array_from_host(rand(Bool, 1000))) + @test_throws ArgumentError AK.findall(identity, ws.buffers.mask; workspace=ws) + + # A `Broadcasted` source keeps its axes, on the host below `switch_below` too + b = Base.Broadcast.broadcasted(identity, array_from_host(Int32[2])) + bc = Base.Broadcast.Broadcasted(b.f, b.args, (Base.OneTo(1000),)) + for switch_below in (0, 2000) + a = AK.BlockReduce(; switch_below) + @test with_workspace(AK.mapreduce, identity, +, bc; alg=a) == 2000 + end + + # A fused source below `switch_below` finishes on the host without device scratch + a = AK.BlockReduce(switch_below=2000) + @test AK.workspace_size(AK.mapreduce, +, +, x, y; alg=a).partials[2] == (0,) + @test with_workspace(AK.mapreduce, +, +, x, y; alg=a) == 2000 + + # The nested operations' algorithms are part of the workspace + bools = array_from_host(rand(Bool, 10_000)) + alg = AK.ScanScatter() + ws = AK.workspace(AK.findall, bools; alg) + @test haskey(ws.nested, :scan) + other = AK.Workspace(ws.backend, ws.device, ws.alg, + merge(ws.nested, (; scan=AK.ScanPrefixes(16, 1))), ws.sizes, ws.buffers) + @test_throws ArgumentError AK.findall(bools; alg, workspace=other) + @test Array(AK.findall(bools; alg, workspace=ws)) == findall(Array(bools)) + ws = AK.workspace(AK.sort!, v; alg=AK.RadixSort()) + @test Base.all(k -> haskey(ws.nested, k), (:scan, :key_range)) + else + # A workspace of reference elements, on the host + ws = AK.workspace(AK.sort!, fill("a", 100); alg=AK.CPUThreads.SampleSort(max_tasks=4)) + @test occursin("Workspace(", sprint(show, ws)) + end + + # Results still infer + V = typeof(v) + @test only(Base.return_types(x -> AK.sum(x), (V,))) === Float32 + @test only(Base.return_types(x -> AK.reduce(+, x), (V,))) === Float32 + @test only(Base.return_types(x -> AK.sum(x; dims=1), (V,))) <: AbstractVector{Float32} + @test only(Base.return_types(x -> AK.accumulate(+, x), (V,))) <: AbstractVector{Float32} +end + + +# With a workspace, an operation allocates no device memory (measured where the backend counts it) +if @isdefined(CUDACore) && BACKEND isa CUDACore.CUDABackend + @testset "workspace: no device allocations" begin + # (`CUDACore.@allocated`, as a function: the macro is not defined when this file is loaded + # for other back-ends) + function device_bytes(f) + b0 = CUDACore.alloc_stats.alloc_bytes + f() + return CUDACore.alloc_stats.alloc_bytes - b0 + end + v = array_from_host(rand(Float32, 100_000)) + ks = array_from_host(rand(Int32(1):Int32(9), 100_000)) + m = reshape(v, 100, 1000) + for (op, args, kw) in ( + (AK.sort!, (copy(v),), (; alg=AK.MergeSort())), + (AK.sort!, (copy(v),), (; alg=AK.RadixSort())), + (AK.sort!, (copy(v),), (; alg=AK.MergeSort(), by=abs)), + (AK.sortperm!, (similar(v, Int), v), (; alg=AK.MergeSort())), + (AK.sortperm!, (similar(v, Int), v), (; alg=AK.MergeSort(lowmem=true))), + (AK.sort_by_key!, (copy(ks), copy(v)), (; alg=AK.MergeSort())), + (AK.sum, (v,), (;)), + # (before Julia 1.12 a reduction of several arrays materializes them, see `Workspace`) + (VERSION >= v"1.12-" ? ((AK.mapreduce, (*, +, v, v), (;)),) : ())..., + (AK.mapreducedim!, (identity, +, similar(v, 1, 1000), m), (;)), + (AK.mapreducedim!, (identity, +, similar(v, 100, 1), m), (;)), + (AK.accumulate!, (+, copy(v)), (;)), (AK.accumulate!, (+, similar(v, Float64), v), (;)), + (AK.any, (x -> x > 2, v), (; alg=AK.ConcurrentWrite())), + (AK.any, (x -> x > 2, v), (; alg=AK.ViaReduce())), + ) + ws = AK.workspace(op, args...; kw...) + op(args...; kw..., workspace=ws) + @test device_bytes(() -> op(args...; kw..., workspace=ws)) == 0 + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index e30eea89..8d9811cf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -36,6 +36,11 @@ const _testset_setup_code = quote d_arr end + global with_workspace + # Call `op` with the workspace made for the same call + with_workspace(op, args...; kwargs...) = + op(args...; kwargs..., workspace=AK.workspace(op, args...; kwargs...)) + global valid_backend_eltypes valid_backend_eltypes(backend, eltypes) = filter(T -> T !== Float64 || KernelAbstractions.supports_float64(backend), eltypes)