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..637ccdcd 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" -GPUArraysCore = "0.2.0" -KernelAbstractions = "0.9.34, 0.10" +CUDACore = "6.4.1" +GPUArraysCore = "0.2.1" +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..ebf812e0 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", @@ -32,11 +33,14 @@ 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", + "Differences from Base" => "api/differences.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/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 new file mode 100644 index 00000000..5470e66e --- /dev/null +++ b/docs/src/api/algorithms.md @@ -0,0 +1,57 @@ +### 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 +``` + +The algorithm types of each operation are documented with it, for [sorting](sort.md), +[reductions](reduce.md), [scans](accumulate.md), [`findall`](findall.md) and +[`any`/`all`](predicates.md). + +#### 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); for reductions, `BlockReduce`; for +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: +```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` +(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 +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/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/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 new file mode 100644 index 00000000..53d5bcd9 --- /dev/null +++ b/docs/src/api/differences.md @@ -0,0 +1,25 @@ +### 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 | +| 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 | +| `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/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/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/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/api/reduce.md b/docs/src/api/reduce.md index b4d6a5ae..20f549b8 100644 --- a/docs/src/api/reduce.md +++ b/docs/src/api/reduce.md @@ -8,3 +8,37 @@ 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. `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 +[`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; 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/docs/src/api/sort.md b/docs/src/api/sort.md index e01d9c6d..b6fedcdd 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.: +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=temp) +ws = AK.workspace(AK.sort!, v) +AK.sort!(v; workspace=ws) ``` 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/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/docs/src/tuning.md b/docs/src/tuning.md new file mode 100644 index 00000000..ebf1cd2c --- /dev/null +++ b/docs/src/tuning.md @@ -0,0 +1,71 @@ +# 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 +``` + +The other families follow the same steps, with a one-line selection: reductions resolve with +`_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 +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.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 +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._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/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/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index c0205561..3790ca33 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 @@ -26,6 +26,8 @@ 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 fdd5acd0..ce2cb5ea 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -1,24 +1,4 @@ -# Available accumulation algorithms -abstract type AccumulateAlgorithm end -struct DecoupledLookback <: AccumulateAlgorithm end -struct ScanPrefixes <: AccumulateAlgorithm end - - -# 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 +include("tuning.jl") # Implementations, then interfaces @@ -28,94 +8,69 @@ 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); - init, - neutral=neutral_element(op, eltype(dst)), - dims::Union{Nothing, Int}=nothing, + op, dst, src; + backend=nothing, + init=, + neutral=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, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, + acctype=nothing, + alg::Algorithm=Auto(), + workspace=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 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). - -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. 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 +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. + +`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"): @@ -124,126 +79,215 @@ 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 -# Use a different algorithm -AK.accumulate!(+, v, alg=AK.DecoupledLookback()) +# Choose the algorithm and its settings +AK.accumulate!(+, v; alg=AK.ScanPrefixes(block_size=512)) ``` """ +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, v::AbstractArray, backend::Backend=get_backend(v); - init, + op, dst::AbstractArray, src::AbstractArray; + backend::Union{Nothing, Backend}=nothing, + workspace=nothing, kwargs... ) - _accumulate_impl!( - op, v, backend; - init, - 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 accumulate!( - op, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(dst); - init, - kwargs... -) - copyto!(dst, src) - _accumulate_impl!( - op, dst, backend; - init, - kwargs... - ) +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 -function _accumulate_impl!( - op, v::AbstractArray, backend::Backend; - init, - neutral=neutral_element(op, eltype(v)), - dims::Union{Nothing, Int}=nothing, +# 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(), +) where {D, T} + dims isa Integer && dims < 1 && + throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) + 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(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 - alg::AccumulateAlgorithm=ScanPrefixes(), - - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, - prefer_threads::Bool=true, +# 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 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 = haskey(bufs, :work) ? bufs.work : v + 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 - # GPU settings - block_size::Int=256, - items_per_thread::Union{Nothing, Int}=nothing, - 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, - ) + if dims === nothing + if a isa CPUThreads.Partitioned + accumulate_1d_cpu!(op, w, 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, 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 + # over + inclusive || fill!(w, _unlane(_scan_first_seed(w, init, neutral))) else - return accumulate_nd!( - op, v, backend; - init, neutral, dims, inclusive, - max_tasks, min_elems, prefer_threads, - block_size, - ) + 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 -""" - accumulate( - op, v::AbstractArray, backend::Backend=get_backend(v); - init, - neutral=neutral_element(op, eltype(v)), - dims::Union{Nothing, Int}=nothing, - inclusive::Bool=true, +# 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 - # CPU settings - max_tasks::Int=Threads.nthreads(), - min_elems::Int=2, +_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 - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), +# 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 - # 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). """ -function accumulate( - op, v::AbstractArray, backend::Backend=get_backend(v); - 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, v::AbstractArray; init=, kwargs...) + +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; backend::Union{Nothing, Backend}=nothing, + init=_NoInit(), dims=nothing, inclusive::Bool=true, acctype=nothing, + kwargs...) + backend = _resolve_backend(backend, v) + 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_cpu.jl b/src/accumulate/accumulate_1d_cpu.jl index 8b64de72..fb73f2b5 100644 --- a/src/accumulate/accumulate_1d_cpu.jl +++ b/src/accumulate/accumulate_1d_cpu.jl @@ -1,56 +1,58 @@ 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 + + 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 @@ -60,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 fb9da460..0c06a81e 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,34 +275,28 @@ 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 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, ::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}, + prefixes::Union{Nothing, AbstractArray}, + 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 @@ -301,39 +306,24 @@ 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 - - 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 + seed = _scan_first_seed(v, init, neutral) shift_to_exclusive = !inclusive && num_blocks > 1 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 @@ -344,25 +334,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}, + prefixes::Union{Nothing, AbstractArray}, + 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 @@ -372,22 +351,15 @@ 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 + 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) @@ -395,7 +367,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 5a44c34f..7cb2ea1d 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,12 @@ 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) + seed = _scan_first_seed(v, init, neutral) + if alg isa CPUThreads.Partitioned + _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 # 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 @@ -53,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 @@ -61,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 @@ -71,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) @@ -86,7 +77,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 @@ -100,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 @@ -122,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) @@ -163,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 @@ -182,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 @@ -191,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) @@ -200,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 @@ -226,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 @@ -246,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 @@ -327,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 new file mode 100644 index 00000000..c8756ff8 --- /dev/null +++ b/src/accumulate/tuning.jl @@ -0,0 +1,167 @@ +# 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 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, S=T) -> Algorithm + +Resolve `alg` for a scan of element type `T` along `dims` (`nothing` for the whole array in +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, + ::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, S) + _check_scan(a, backend, S, 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 new file mode 100644 index 00000000..a6df18af --- /dev/null +++ b/src/algorithms.jl @@ -0,0 +1,266 @@ +# 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 + +""" + ReduceAlgorithm <: Algorithm + +Supertype of the reduction algorithms: [`BlockReduce`](@ref). Reductions also accept +[`CPUThreads.Partitioned`](@ref AcceleratedKernels.CPUThreads.Partitioned). +""" +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 + +""" + 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 + +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: Algorithm, 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 + +""" + 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 + + +# 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 + +""" + _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 + +# 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 (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) + +_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`, 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...) + b = _backend_votes(args...) + b === nothing ? HOST_BACKEND : b +end diff --git a/src/arithmetics.jl b/src/arithmetics.jl index 16f98761..9d06ef29 100644 --- a/src/arithmetics.jl +++ b/src/arithmetics.jl @@ -1,322 +1,138 @@ """ - sum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - dims=nothing, + sum(src::AbstractArray; kwargs...) - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, +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. - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) - -Sum of elements of an array, with optional `init` and `dims`. Arguments are the same as for -[`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) -``` +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 -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 ``` """ -function sum( - src::AbstractArray, backend::Backend=get_backend(src); - init=zero(eltype(src)), - kwargs... -) - reduce( - +, src, backend; - init, - kwargs... - ) -end +sum(src::AbstractArray; kwargs...) = _reduce_or_empty(zero, Base.add_sum, src; kwargs...) """ - prod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, + prod(src::AbstractArray; kwargs...) - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) +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. -Product of elements of an array, with optional `init` and `dims`. Arguments are the same as for -[`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... - ) +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(), + 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) && + (_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)) + 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; workspace, kwargs...) end """ - maximum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemin(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, + maximum(src::AbstractArray; kwargs...) - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) +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). -Maximum of elements of an array, with optional `init` and `dims`. Arguments are the same as for -[`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; kwargs...) = reduce(max, src; kwargs...) """ - minimum( - src::AbstractArray, backend::Backend=get_backend(src); - init=typemax(eltype(src)), - dims=nothing, - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, + minimum(src::AbstractArray; kwargs...) - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, - ) +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). -Minimum of elements of an array, with optional `init` and `dims`. Arguments are the same as for -[`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; kwargs...) = reduce(min, src; 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`: `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). - # 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... - ) -end +count(src::AbstractArray; kwargs...) = count(identity, src; kwargs...) +count(f, src::AbstractArray; init=0, kwargs...) = + mapreduce(_BoolValued(f), Base.add_sum, src; init, kwargs...) - -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... - ) -end +_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, backend::Backend=get_backend(src); - init=zero(eltype(src)), - neutral=zero(eltype(src)), - dims::Union{Nothing, Int}=nothing, - - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), + cumsum(src::AbstractArray; kwargs...) - # 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, 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: @@ -334,39 +150,15 @@ 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... - ) -end +cumsum(src::AbstractArray; kwargs...) = accumulate(Base.add_sum, src; kwargs...) """ - cumprod( - src::AbstractArray, backend::Backend=get_backend(src); - init=one(eltype(src)), - neutral=one(eltype(src)), - dims::Union{Nothing, Int}=nothing, + cumprod(src::AbstractArray; kwargs...) - # Algorithm choice - alg::AccumulateAlgorithm=ScanPrefixes(), - - # GPU settings - block_size::Int=256, - temp::Union{Nothing, AbstractArray}=nothing, - temp_flags::Union{Nothing, AbstractArray}=nothing, - ) - -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, 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: @@ -384,16 +176,7 @@ 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... - ) -end +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 055a161d..e1aa6e20 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 @@ -81,44 +141,29 @@ 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 - @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) - 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 = 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,20 +173,19 @@ end function findall_cpu( - bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend; - max_tasks::Int, - min_elems::Int, - temp, + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, + alg::CPUThreads.Partitioned; + bufs, ) 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) 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, @@ -170,44 +214,24 @@ 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; bufs, ) where I - alg = findall_algorithm(alg) - - if use_gpu_algorithm(backend, prefer_threads) - findall_gpu(bools, I, output_indices, backend, alg; temp) + if alg isa ScanScatter + findall_gpu(bools, I, output_indices, backend, alg; bufs) else - findall_cpu(bools, I, output_indices, backend; max_tasks, min_elems, temp) + findall_cpu(bools, I, output_indices, backend, alg; bufs) 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, -) - bools = if isnothing(temp_bools) - similar(v, Bool) - 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; - max_tasks, min_elems, prefer_threads, 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 @@ -217,37 +241,31 @@ 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(), workspace=nothing) + findall(pred, A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), workspace=nothing) -The supported inputs are arrays. Dictionaries, other iterables, and scalar inputs accepted by -`Base.findall` are outside the scope of this package. +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. -## Settings +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: -- `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. +- `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. -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. +The supported inputs are arrays. Dictionaries, other iterables, and scalar inputs accepted by +`Base.findall` are outside the scope of this package. + +`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`. `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 @@ -256,50 +274,62 @@ 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, - temp::Union{Nothing, AbstractArray}=nothing, - temp_bools::Union{Nothing, AbstractArray}=nothing, -) - alg = findall_algorithm(alg) - bools = if values isa AbstractArray{Bool} - isnothing(temp_bools) || - throw(ArgumentError("temp_bools is not used for a Bool mask")) - values +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) + 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(), + workspace=nothing) + _check_bool_result(pred, v) + 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)) + 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 - findall_bools(identity, values, backend, temp_bools; - max_tasks, min_elems, prefer_threads, block_size=alg.block_size) + num_tasks = TaskPartitioner(n, a.max_tasks, a.min_elems).num_tasks + (; counts=_buffer(Int, num_tasks)), (;) 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) + return (; plan=_Plan(backend, a, nested, merge(mask, counts))) 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, - 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) +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/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/predicates.jl b/src/predicates.jl index cfd63f86..b4c02f74 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,24 @@ end """ - any( - pred, v::AbstractArray, backend::Backend=get_backend(v); + any(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto(), workspace=nothing) - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), - - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, - - # GPU settings - block_size::Int=256, - ) - -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. `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)` +to only use parallelism when worth it in your application. # Examples ```julia @@ -76,17 +144,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,114 +158,76 @@ 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(), 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`). +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, - prefer_threads=true, - block_size, - temp=alg.temp, - switch_below=alg.switch_below, - ) - end +# 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 - 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 - end - end - return overall[] + (;) 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 = 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, bufs) = + _mapreduce_nested(_BoolValued(pred), |, v, bufs.reduce; backend, init=false, alg=alg.reduce) + +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 + 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 + end + return overall[] +end """ - all( - pred, v::AbstractArray, backend::Backend=get_backend(v); + all(pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto(), workspace=nothing) - # Algorithm choice - alg::PredicatesAlgorithm=ConcurrentWrite(), - - # 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 @@ -212,82 +238,27 @@ 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(), workspace=nothing) + p = _predicate_plan(pred, &, v, backend, alg) + _all(pred, v, p.backend, p.alg, _buffers(p, workspace, v)) end +_plan(::typeof(all), pred, v::AbstractArray; backend=nothing, alg::Algorithm=Auto()) = + _predicate_plan(pred, &, v, backend, alg) -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, - prefer_threads=false, - block_size, - temp=alg.temp, - switch_below=alg.switch_below, - ) - 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, 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_cpu.jl b/src/reduce/mapreduce_1d_cpu.jl index 3a47b603..3d31ed10 100644 --- a/src/reduce/mapreduce_1d_cpu.jl +++ b/src/reduce/mapreduce_1d_cpu.jl @@ -1,37 +1,25 @@ +# 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, 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)) - 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 cfc81c0b..51d27883 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,19 +38,16 @@ end +# Reduce the non-empty `src` to a host value; `neutral` is the partial-result seed of +# `_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, - - # CPU settings - ignored here - max_tasks::Int, - min_elems::Int, - - # GPU settings block_size::Int, items_per_thread::Int, - temp::Union{Nothing, AbstractArray}, + partials::Union{Nothing, AbstractArray}, switch_below::Int, ) @argcheck 1 <= block_size <= 1024 @@ -56,29 +55,24 @@ 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. num_per_block = items_per_thread * block_size blocks = (len + num_per_block - 1) ÷ num_per_block - if !isnothing(temp) - @argcheck get_backend(temp) === backend - @argcheck eltype(temp) === typeof(init) - @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) - 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) @@ -93,7 +87,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 @@ -104,22 +98,40 @@ 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) +# 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 5304cce5..0960ecc8 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. @@ -71,26 +67,106 @@ function _canonicalize_dims(src_sizes, src_strides, dims_valid) return Tuple(reduce_segs), Tuple(outer_segs) end -# Main entry point +# 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 -function mapreduce_nd( - f, op, src::MapReduceSource, backend::Backend; - init, - neutral=neutral_element(op, typeof(init)), - dims, +# 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,) - # CPU settings - max_tasks::Int, - min_elems::Int, - prefer_threads::Bool=true, + # 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 - # GPU settings - block_size::Int, - temp::Union{Nothing, AbstractArray}, -) - @argcheck 1 <= block_size <= 1024 - @argcheck ispow2(block_size) + 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: +# 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 @@ -99,64 +175,61 @@ function mapreduce_nd( 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 - # Match Base: duplicate dims are ignored, e.g. dims=(2,2) behaves like dims=2. - dims_all = Tuple(Base.unique(dims_buf)) +# 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, + 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) : + (; 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, A).target_blocks + target_blocks >= 1 || throw(ArgumentError( + "the reduction tuning's `target_blocks` must be positive, got $target_blocks")) 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, max_tasks, min_elems, block_size) - 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) + len = Base.prod(src_sizes[d] for d in dims_valid; init=1) - # 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, max_tasks, min_elems, block_size) + _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) - 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, neutral, launch...) return dst end @@ -172,7 +245,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,), ) @@ -180,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 @@ -245,78 +272,49 @@ 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,), ) - 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) 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,), ) - 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, init, neutral, + 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(init), (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, init, neutral, - 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, init, neutral, - 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, 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, ) @@ -345,7 +343,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 @@ -362,13 +361,17 @@ 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 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 @@ -417,7 +420,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, ) @@ -427,7 +430,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,), @@ -435,7 +438,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,), @@ -443,7 +446,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,), @@ -469,33 +472,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(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 + 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 @@ -552,9 +547,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 @@ -562,25 +558,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 @@ -598,7 +595,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 @@ -611,14 +608,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 @@ -628,12 +625,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 @@ -646,13 +644,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 @@ -662,20 +660,21 @@ 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). @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 @@ -687,17 +686,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 @@ -718,7 +717,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, @@ -727,7 +726,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 @@ -741,21 +741,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 @@ -779,7 +779,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 @@ -794,14 +795,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] @@ -809,17 +810,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 @@ -828,17 +830,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 c18c44e1..9ab65927 100644 --- a/src/reduce/reduce.jl +++ b/src/reduce/reduce.jl @@ -1,43 +1,11 @@ # 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 +# 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) @@ -53,163 +21,83 @@ include("mapreduce_nd.jl") """ - reduce( - op, src::AbstractArray, backend::Backend=get_backend(src); - 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)), - 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`. - -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. - -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`. + reduce(op, src::AbstractArray; kwargs...) -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. - -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`). +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: ```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) +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)) ``` """ -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); - 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)), - temp::Union{Nothing, AbstractArray}=nothing, - switch_below::Int=0, + f, op, src, srcs::AbstractArray...; + backend=nothing, + init=, + neutral=nothing, + acctype=nothing, + dims=:, + alg::Algorithm=Auto(), + workspace=nothing, ) - 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. - -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`). +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. + +`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). -# Example +# Examples Computing a sum of squares, 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)) -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: @@ -219,123 +107,264 @@ 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) +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) -``` - -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) +rows = AK.mapreduce((x, y) -> x * y, +, a, b; dims=1) ``` """ function mapreduce( - f, op, src::MapReduceSource, backend::Backend=_mapreduce_backend(src); - init, - kwargs... + f, op, src::MapReduceSource, srcs::AbstractArray...; + backend::Union{Nothing, Backend}=nothing, + init=_NoInit(), + neutral=nothing, + acctype=nothing, + dims=:, + alg::Algorithm=Auto(), + workspace=nothing, ) - _mapreduce_impl( - f, op, src, backend; - init, - kwargs... - ) + 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 mapreduce( - f, op, src::AbstractArray, src2::AbstractArray, srcs::AbstractArray...; - init, - kwargs... +function _plan( + ::typeof(mapreduce), f, op, src::MapReduceSource, srcs::AbstractArray...; + backend=nothing, init=_NoInit(), neutral=nothing, acctype=nothing, dims=:, + alg::Algorithm=Auto(), ) - return _mapreduce_multi( - f, op, nothing, src, src2, srcs...; - init, - kwargs... - ) + 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::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 +_plan(::typeof(reduce), op, src::AbstractArray; kwargs...) = + _plan(mapreduce, identity, op, src; kwargs...) -function _mapreduce_multi( - f, op, backend::Union{Nothing, Backend}, src::AbstractArray, srcs::AbstractArray...; - init, - 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 mapreduce( - identity, op, bc, isnothing(backend) ? _mapreduce_backend(bc) : backend; - init, - kwargs... - ) + 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) + a = _resolve_reduce(alg, backend, A, dims) + 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 + + 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)) + 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 + 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 -function _mapreduce_impl( - f, op, src::MapReduceSource, backend::Backend; - init, - neutral=neutral_element(op, typeof(init)), - dims = nothing, +# 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) - # 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, +""" + mapreducedim!( + f, op, R::AbstractArray, src; + backend=get_backend(R), + init=, + neutral=nothing, + 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 +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. + +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call. + +```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(), + 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)) + 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) + 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 - # 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) +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 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( + "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) - if isnothing(dims) || dims isa Colon - if use_gpu_algorithm(backend, prefer_threads) - mapreduce_1d_gpu( - f, op, src, backend; - init, neutral, - max_tasks, min_elems, - block_size, items_per_thread, temp, - switch_below - ) - else - mapreduce_1d_cpu( - f, op, src, backend; - init, neutral, - max_tasks, min_elems, - block_size, temp, - switch_below - ) - 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, partials)) +end + +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, + partials, switch_below=alg.switch_below, + ) else - return mapreduce_nd( + mapreduce_1d_cpu( f, op, src, backend; - init, neutral, dims, - max_tasks, prefer_threads, - min_elems, block_size, - temp, + 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 new file mode 100644 index 00000000..25f081f8 --- /dev/null +++ b/src/reduce/tuning.jl @@ -0,0 +1,120 @@ +# 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, T, 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, ::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 + # 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, T, dims) + _checkdomain(a) + _require_threads(a, backend) +end diff --git a/src/reduce/utilities.jl b/src/reduce/utilities.jl index d3000532..d7cf0b6a 100644 --- a/src/reduce/utilities.jl +++ b/src/reduce/utilities.jl @@ -1,17 +1,188 @@ -# 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) +# `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. + +# 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 + +# 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 +# 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,16 +202,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, - max_tasks=Threads.nthreads(), - min_elems=1, - block_size=256, -) - foreachindex(dst, backend; max_tasks, min_elems, block_size) do i - dst[i] = op(init, f(src[i])) +# 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 + 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/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/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..5c8758ad 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 @@ -171,21 +172,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, @@ -211,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 @@ -257,21 +247,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 +265,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(eachindex(ix), HOST_BACKEND; max_tasks, min_elems) do i @inbounds ix[i] = i end @@ -305,7 +282,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 +292,47 @@ 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, bufs; max_tasks, min_elems) + ix = bufs.ix + if dims isa Colon + _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, 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, 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]] + end + 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 72a04dc3..b7603e72 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -128,24 +128,10 @@ 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!( - v::AbstractArray, backend::Backend=get_backend(v); +# 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, bufs::NamedTuple; lt=isless, by=identity, @@ -153,27 +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. + # 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, v, backend; - lt, rev, order, block_size, dims, - temp_values=temp, # temp was for v swap buffer; maps to temp_values here - ) + keys = bufs.keys + map!(by, keys, v; backend, block_size) + _merge_sort_by_key!(keys, v, backend, bufs; lt, rev, order, block_size, dims) return v end @@ -193,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) @@ -219,29 +198,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..6f3a3c8c 100644 --- a/src/sort/merge_sort_by_key.jl +++ b/src/sort/merge_sort_by_key.jl @@ -161,29 +161,14 @@ 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. +# 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, @@ -191,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 @@ -202,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) @@ -229,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) @@ -262,34 +237,8 @@ function merge_sort_by_key!( 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 +# 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 44b50f46..71acb05e 100644 --- a/src/sort/merge_sortperm.jl +++ b/src/sort/merge_sortperm.jl @@ -1,37 +1,16 @@ -""" - 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); + 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 @@ -40,85 +19,27 @@ 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(ix, block_size=block_size) do i + _foreachindex(eachindex(ix), backend; block_size) do i @inbounds ix[i] = i end (isempty(v) || layout.len <= 1) && return ix - keys = inplace ? v : copy(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 -""" - 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); + backend::Backend, + bufs::NamedTuple; lt=isless, by=identity, @@ -126,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 @@ -135,13 +55,9 @@ 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(ix, block_size=block_size) do i + _foreachindex(eachindex(ix), backend; block_size) do i @inbounds ix[i] = i end (isempty(ix) || layout.len <= 1) && return ix @@ -162,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) @@ -188,31 +104,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..dc2c3380 100644 --- a/src/sort/radix_sort.jl +++ b/src/sort/radix_sort.jl @@ -366,17 +366,23 @@ end end -# Return the extrema of the transformed sort keys. -function _rs_key_range(v::AbstractArray{T}, 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; - 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 @@ -385,17 +391,47 @@ function _rs_key_range(v::AbstractArray{T}, descending::Bool) where T 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 @@ -404,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) @@ -420,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, descending) + min_key, max_key = _rs_key_range(p1, backend, descending, bufs.key_range) vitems = Val(items) hist_kern! = has_atomics ? @@ -454,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 7876b8d9..76302e4c 100644 --- a/src/sort/sort.jl +++ b/src/sort/sort.jl @@ -8,456 +8,404 @@ 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, + workspace=nothing, + ) -> v - # 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, - ) - -Sorts the array `v` in-place using the specified backend. The `lt`, `by`, `rev`, and `order` -arguments are the same as for `Base.sort`. +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. -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. +`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`. -## 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. +`backend` is derived from `v`; pass it only for arrays that do not determine their backend. -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. +`workspace` takes the scratch memory of a [`workspace`](@ref) made for the same call, so that the +sort allocates none of its own. # 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... + 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.Order.Forward, + workspace=nothing, ) - _sort_impl!( - v, backend; - kwargs... - ) + 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_impl!( - v::AbstractArray, backend::Backend; - - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Forward, +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) + sizes, nested = _sort_plan(a, v, backend, dims, by) + return _Plan(backend, a, nested, sizes), ord +end - max_tasks=Threads.nthreads(), - min_elems=1, - prefer_threads::Bool=true, +# 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)) : (;) - alg::Union{Nothing, SortAlgorithm}=nothing, +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 - # 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, 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 - # 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, 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 - # 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, 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=get(bufs, :temp, nothing)) 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, workspace=nothing, + kwargs...) + backend = _resolve_backend(backend, v) + # 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!( 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, + workspace=nothing, + ) -> ix - # 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, - ) +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). -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. +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)`. -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. - -## 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`. """ 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, + order::Base.Order.Ordering=Base.Order.Forward, + 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 - alg::Union{Nothing, SortAlgorithm}=nothing, +_plan(::typeof(sortperm!), ix::AbstractArray, v::AbstractArray; kwargs...) = + first(_sortperm_setup(ix, v; kwargs...)) - # Permute each 1D slice along this dimension; `:` permutes the whole array as one vector - dims::Union{Colon, Integer}=Colon(), +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) + 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) + return _Plan(backend, a, _sortperm_sizes(a, ix, v, dims)), ord +end - # GPU settings; nothing => merge sort's tuned default (sortperm is merge-only) - block_size::Union{Nothing, Int}=nothing, +# 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)) : (;) - # Temporary buffer, same size as `v` - 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))")) +function _sortperm_impl!(a::MergeSort, ix, v, backend, dims, ord, bufs; lt, by, rev, order) + if a.lowmem + _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, bufs; lt, by, rev, order, block_size=a.block_size, dims) end +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::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=get(bufs, :temp, nothing)) 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 + _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...) - lt=isless, - by=identity, - rev::Union{Nothing, Bool}=nothing, - order::Base.Order.Ordering=Base.Order.Forward, +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 - # CPU settings - max_tasks=Threads.nthreads(), - min_elems=1, +# `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...)) - # Algorithm choice - alg::Union{Nothing, SortAlgorithm}=nothing, +# 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 - # 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, +""" + 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, + workspace=nothing, + ) -> (keys, values) + +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`). + +`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`. + +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, + workspace=nothing, ) - ix = similar(v, Int) - sortperm!( - ix, v, backend; - kwargs... - ) + 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) + 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) + 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, 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, 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/sort/tuning.jl b/src/sort/tuning.jl new file mode 100644 index 00000000..4b53bfb4 --- /dev/null +++ b/src/sort/tuning.jl @@ -0,0 +1,185 @@ +# 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() + + +# 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(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 + +_checkdomain(a::CPUThreads.SampleSort) = _check_threads(a) + + +# 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) = _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 _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) + _require_threads(a, backend) +end diff --git a/src/utils.jl b/src/utils.jl index 46fd96cc..82a2a3e4 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -2,28 +2,13 @@ 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 - -# 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 -# 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 +# 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 """ struct TypeWrap{T} 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/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 94adb396..dd72d436 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,34 @@ 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. + # A workspace, reused 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))) + 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, - prefer_threads, block_size=128, alg, - temp=array_from_host(zeros(Int64, 1000)), - temp_flags=array_from_host(zeros(Int8, 1000))) + 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 @@ -133,14 +155,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,20 +207,23 @@ 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 # 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; prefer_threads, max_tasks=4, init, neutral=SCAN_I2, - inclusive, block_size, items_per_thread, alg) + 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 end @@ -210,7 +235,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), @@ -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, 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 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 @@ -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,174 @@ 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 + + +@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/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/findall.jl b/test/generic/findall.jl index 3adbdb25..e38a94a5 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)) == + @test with_workspace(AK.findall, Any[true, false, true]) == 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,70 @@ 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))) - temp_bools = similar(v, Bool) - @test Array(AK.findall(x -> x > 0.5f0, v; - prefer_threads, max_tasks=4, alg, temp, temp_bools)) == + # 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) + @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; prefer_threads, 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)) - - if !prefer_threads - 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, - temp=similar(v, Int32, 100)) - @test_throws ArgumentError AK.findall(bools; prefer_threads, - temp=similar(v, Int, 1)) + @test_throws ArgumentError AK.findall(v; alg=OtherFindallAlgorithm()) + if TEST_KERNELS + @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)) 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/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/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 c004cf03..78b0f507 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,35 @@ 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 + 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)); - prefer_threads, - 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 +167,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 +184,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 +197,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 +210,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 +225,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 +233,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 +241,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 +250,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 +260,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 +272,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 +347,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 +355,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( + AK.mapreducedim!( + identity, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 1, 5)), 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( + AK.mapreducedim!( + identity, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 4, 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 +396,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 +422,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 +440,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 +455,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 +466,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 +475,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,69 +537,70 @@ 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 # 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; - 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) + f(v) # 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 +615,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 +632,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 +669,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 +692,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 +700,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 +709,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 +731,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 +740,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 +750,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,48 +776,38 @@ 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 - AK.mapreduce( + AK.mapreducedim!( -, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 1, 5)), 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( + AK.mapreducedim!( -, (x, y) -> x + 1, + array_from_host(zeros(Int32, 3, 4, 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 +816,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 +834,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 +846,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 +862,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 +880,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 +892,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 +908,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 +926,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 +938,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 +954,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 +972,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 +984,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 +1000,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 +1018,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,32 +1032,20 @@ 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 -@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; @@ -1078,11 +1057,251 @@ 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 + + +@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 new file mode 100644 index 00000000..0d7f26ff --- /dev/null +++ b/test/generic/reduce_resolve.jl @@ -0,0 +1,206 @@ +# 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, 1, 4), + zeros(Int32, 4, 4), RRB, + AK.BlockReduce(256, 2, 0), Int32; + init=AK._NoInit(), neutral=nothing, + dims_valid=(1,), bufs=(;)) + 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 + + +@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 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 diff --git a/test/generic/scan_resolve.jl b/test/generic/scan_resolve.jl new file mode 100644 index 00000000..1723cbd3 --- /dev/null +++ b/test/generic/scan_resolve.jl @@ -0,0 +1,152 @@ +# 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 + + # 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) + + # 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 diff --git a/test/generic/sort.jl b/test/generic/sort.jl index de7bd002..8b1eaecd 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,31 @@ 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)) + 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.merge_sort!(v, lt=(>), rev=true, - 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.merge_sort(v, lt=(>), by=abs, rev=true, - 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.merge_sort(v, lt=(>), by=abs, rev=true, - 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 @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 +69,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 +78,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,17 +86,16 @@ 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 - # 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.merge_sort!(tmp; 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 @@ -109,27 +111,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 +166,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 +174,19 @@ 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)) + 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.sample_sort!(v, lt=(>), rev=true, - 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 @@ -171,7 +199,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 +207,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,34 +215,26 @@ 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, - 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; prefer_threads, lt=(>), rev=true, - max_tasks=64, min_elems=8, block_size=64, - 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; prefer_threads, lt=(>), by=abs, rev=true, - max_tasks=64, min_elems=8, block_size=64, - 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; prefer_threads, lt=(>), by=abs, rev=true, - max_tasks=64, min_elems=8, block_size=64, - 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 @@ -229,56 +249,65 @@ 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) 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; prefer_threads, alg, temp) + with_workspace(AK.sortperm!, ix, v; alg) @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 +315,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 +326,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 +337,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,48 +347,72 @@ 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, - 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)) + 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.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)) + 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, 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]) + # 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 + @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 +421,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 +431,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 +441,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 +450,23 @@ 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)) + 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.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 = 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]) 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 +474,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 +484,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 +494,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,11 +503,8 @@ 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, - 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]) @@ -468,7 +512,7 @@ end end -if !prefer_threads +if TEST_KERNELS @testset "merge_sortperm_lowmem" begin Random.seed!(0) @@ -477,7 +521,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 +531,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 +541,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,20 +550,15 @@ 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, - 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.merge_sortperm_lowmem(v, - lt=(>), by=abs, rev=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]) @@ -535,7 +574,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 +584,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 +594,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]) @@ -564,29 +603,20 @@ end # Testing different settings ix = array_from_host(1:10_000, Int32) v = array_from_host(1:10_000, Float32) - AK.sortperm!(ix, - v; - prefer_threads, - lt=(>), by=abs, rev=true, - block_size=64, - 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; - prefer_threads, - lt=(>), by=abs, rev=true, - block_size=64, - 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]) 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 +634,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 +644,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 +662,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,22 +679,22 @@ 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 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; temp) - AK.sortperm!(ix2, v2; 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 ──────────────────────────────────── @@ -674,7 +704,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 +714,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 +724,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 +737,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 +747,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 +758,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 +768,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 +780,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 +793,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 +801,26 @@ 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 ─────────────────────────────────── + # ── workspace: preallocated buffers ─────────────────────────────────── 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) + with_workspace(AK.sort!, v; alg=AK.RadixSort()) @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 +828,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 +843,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 +874,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 +883,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 + # 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; 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(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()) 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 +950,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 +960,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 + # 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; 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(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) 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 +1000,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 +1010,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 +1048,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 +1058,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 +1067,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 +1079,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 +1088,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 +1100,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 +1144,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 +1156,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 +1168,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..dbaaf1d9 --- /dev/null +++ b/test/generic/sort_resolve.jl @@ -0,0 +1,318 @@ +# 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 + # ... 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) + @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 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) + # ... 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/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 e901ccc6..8d9811cf 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)) @@ -36,8 +36,17 @@ 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) + + # 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 @@ -52,7 +61,8 @@ 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 $_testset_setup_code end) @@ -67,7 +77,8 @@ 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 $_testset_setup_code end) @@ -81,7 +92,8 @@ 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 $_testset_setup_code end) @@ -95,7 +107,8 @@ 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 $_testset_setup_code end) @@ -110,18 +123,23 @@ 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 $_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) @@ -132,7 +150,8 @@ 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 $_testset_setup_code end)