diff --git a/docs/src/api/algorithms.md b/docs/src/api/algorithms.md index 723d1217..3692f663 100644 --- a/docs/src/api/algorithms.md +++ b/docs/src/api/algorithms.md @@ -40,14 +40,26 @@ element type, the ordering or the backend cannot run, or an invalid setting, is #### The `backend` keyword Operations run on the backend of their arrays: `backend` is derived from every array argument, -the destination first, and all of them must agree. Ranges, `CartesianIndices`, `LinearIndices` -(also through views and reshapes), numbers and other non-array arguments do not count. If no +the destination first, and all of them must agree. Ranges, `CartesianIndices`, `LinearIndices`, +numbers and other non-array arguments do not count. An array that wraps others (one whose `parent` +is another array, or a tuple of arrays) counts as the arrays it wraps: a view or a reshape, but +also a lazy array such as a [MappedArrays.jl](https://github.com/JuliaArrays/MappedArrays.jl) +`mappedarray(f, A, B)`, so that a lazy array computed from ranges does not count either. 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. +instance), for lazy arrays whose function reads device arrays it captures rather than wraps, 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. + +Results of operations on a lazy array (e.g. `AK.sort(mappedarray(f, A))`) are allocated with its +`similar` method, or on `backend` if it has none of its own (Base's fallback gives a host `Array`) +or wraps no memory. Running a GPU kernel over a lazy array needs a version of it that the kernel +can use: an `isbits` lazy array (`mappedarray(abs2, 1:n)`) is passed as it is, while one wrapping +device arrays needs an [Adapt.jl](https://github.com/JuliaGPU/Adapt.jl) rule for its type, as +Base's views and reshapes have. MappedArrays.jl does not define such rules yet, so a +`mappedarray` of GPU arrays cannot be processed on the GPU. #### The host backend diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index fe20d08c..e13ae612 100644 --- a/src/AcceleratedKernels.jl +++ b/src/AcceleratedKernels.jl @@ -30,7 +30,7 @@ end # Internal dependencies using ArgCheck: @argcheck import Adapt -using GPUArraysCore: @allowscalar, neutral_element +using GPUArraysCore: AnyGPUArray, @allowscalar, neutral_element using KernelAbstractions using KernelAbstractions: @context # Sub-group operations, with KernelAbstractions 0.10 diff --git a/src/algorithms.jl b/src/algorithms.jl index a9e66a6f..3d28f255 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -235,13 +235,23 @@ backendless(x::BaseWrapper) = backendless(parent(x)) # Values that never determine the backend: backendless arrays, 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 (and forwards Base's wrappers to the array they wrap). +# do not implement it (and forwards wrappers to the array they wrap). backend_vote(_) = nothing -backend_vote(x::AbstractArray) = backendless(x) ? nothing : get_backend(x) +backend_vote(x::AbstractArray) = backendless(x) ? nothing : wrapper_vote(x, 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) +# A wrapper whose `parent` is not an array with memory (a range, or a tuple of arrays for a lazy +# array over several, like `mappedarray(+, A, B)`), where `get_backend` throws, votes like the +# arrays it wraps instead, so that a lazy array over ranges does not vote either +wrapper_vote(x, p) = p === x || wraps_storage(p) ? get_backend(x) : backend_vote(p) + +# Whether `x` is an array with memory, or a chain of wrappers ending in one +wraps_storage(::IndexCollection) = false +wraps_storage(x::AbstractArray) = (p = parent(x); p === x || wraps_storage(p)) +wraps_storage(_) = false + backend_votes() = nothing backend_votes(x, xs...) = merge_backends(backend_vote(x), backend_votes(xs...)) @@ -259,10 +269,10 @@ end 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. +(destination first) agrees on, recursing into `Broadcasted` trees and into the arrays that +wrappers wrap; ranges, `CartesianIndices`, `LinearIndices`, wrappers of only those, and non-array +values do not count. Arguments on different backends are an `ArgumentError`. If no argument +determines it, the host backend is used. """ resolve_backend(backend::Backend, args...) = backend function resolve_backend(::Nothing, args...) diff --git a/src/reduce/mapreduce_1d_gpu.jl b/src/reduce/mapreduce_1d_gpu.jl index 6495f520..214ae31a 100644 --- a/src/reduce/mapreduce_1d_gpu.jl +++ b/src/reduce/mapreduce_1d_gpu.jl @@ -82,7 +82,7 @@ function mapreduce_1d_gpu( len == 1 && return finish(op, init, nothing, 0, @allowscalar(op_host(seed, f_host(src[1])))) if len < switch_below - h_src = host_copy(src) + h_src = host_copy(src, backend) return finish(op, init, nothing, 0, Base.mapreduce(f_host, op_host, h_src; init=seed)) end @@ -144,16 +144,17 @@ function mapreduce_1d_buffer(src) return layout.buffer, layout.offset end -host_copy(src::AbstractArray) = Array(src) +host_copy(src::AbstractArray, backend) = host_array(backend, 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 +host_copy(src::Base.Broadcast.Broadcasted, backend) = + Base.Broadcast.materialize(on_host(src, backend)) +on_host(bc::Base.Broadcast.Broadcasted, backend) = + Base.Broadcast.Broadcasted(bc.f, Base.map(x -> on_host(x, backend), bc.args), bc.axes) +on_host(x::Base.Broadcast.Extruded, backend) = on_host(x.x, backend) +on_host(x::AbstractArray, backend) = host_array(backend, x) +on_host(x::AbstractRange, _) = x +on_host(x, _) = x # (a view of all elements, indexed from 1 also when the source's indices are not) mapreduce_1d_view(src::AbstractArray) = @view src[begin:end] diff --git a/src/sort/merge_sort.jl b/src/sort/merge_sort.jl index b5606d67..7aaf8091 100644 --- a/src/sort/merge_sort.jl +++ b/src/sort/merge_sort.jl @@ -271,8 +271,7 @@ function merge_sortperm!( end (isempty(v) || layout.len <= 1) && return ix if by === identity - # (a range, for instance, has no backend: collected on the host first) - keys = copyto_on!(bufs.keys, backendless(v) ? collect(v) : v, backend) + keys = copyto_on!(bufs.keys, v, backend) merge_sort_by_key!(keys, ix, backend, bufs, ord; block_size, dims) else # Evaluated once instead of in every comparison; `ord` is `By(by, ord.order)` diff --git a/src/utils.jl b/src/utils.jl index 0a774696..06bb9b99 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -12,12 +12,57 @@ struct NoInit end @inline unval(::Val{V}) where {V} = V -# 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` +# The array that holds the memory behind `v`: `v` itself unless it wraps others, else the first of +# the arrays it wraps that holds memory; `nothing` if there is none (a range, or a lazy array over +# one) +storage(v::Tuple) = storage_first(v...) +function storage(v::AbstractArray) + backendless(v) && return nothing + p = parent(v) + p === v ? v : storage(p) +end +storage(_) = nothing +storage_first() = nothing +storage_first(x, xs...) = (s = storage(x); s === nothing ? storage_first(xs...) : s) + +# Whether `v` gets Base's `similar` fallback, a host `Array`, although it wraps the memory `s` of +# another array type (a lazy array without a `similar` method of its own, for instance). Checked on +# an empty result. +similar_falls_back(v, s) = + s !== v && !(s isa Array) && similar(v, eltype(v), Base.map(zero, size(v))) isa Array + +# Whether results of operations on `v` are allocated on `backend` rather than with `similar(v)`: +# where `v` holds no memory (a range, or a lazy array over one), or gets Base's fallback +function allocates_on_backend(v) + s = storage(v) + s === nothing || similar_falls_back(v, s) +end + +# An array for an operation's result similar_on(backend, v, ::Type{T}=eltype(v), dims=size(v)) where {T} = - backendless(v) ? KernelAbstractions.allocate(backend, T, dims) : similar(v, T, dims) -# A copy of `v`, likewise -copy_on(backend, v) = backendless(v) ? copyto!(similar_on(backend, v), collect(v)) : copy(v) + allocates_on_backend(v) ? KernelAbstractions.allocate(backend, T, dims) : similar(v, T, dims) +# A copy of `v`, likewise: its own `copy`, unless its results are allocated on `backend` (where +# Base's copy would give a host `Array`) +function copy_on(backend, v) + allocates_on_backend(v) && copies_by_kernel(backend, v) && + return copyto_on!(similar_on(backend, v), v, backend) + backendless(v) ? copyto!(similar_on(backend, v), collect(v)) : copy(v) +end +# A copy of `v` in host memory +host_array(backend, v) = copies_by_kernel(backend, v) ? + Array(copyto_on!(KernelAbstractions.allocate(backend, eltype(v), size(v)), v, backend)) : + Array(v) + +# Whether copies of `v` on `backend` use a kernel. Its own arrays, Base's wrappers of arrays in +# host memory and GPUArrays' wrappers of its arrays are copied by Base or GPUArrays, backendless +# arrays after collecting them on the host, and everything on the host backend. Base's copies of +# other wrappers index device memory element by element on the host, and lazy arrays over ranges +# may read device arrays their function captures. +function copies_by_kernel(backend, v) + (runs_threads(backend) || backendless(v) || parent(v) === v || v isa AnyGPUArray) && + return false + !(storage(v) isa Union{Array, BitArray}) +end # The destination of an operation (`what`, for the messages) that writes the elements of `src` # to `dst`: `src` itself where the operation runs in place, else an array that does not overlap @@ -55,9 +100,11 @@ has_offset_parent(A) = # `copyto!(dst, src)` of arrays of the same size function copyto_on!(dst, src, backend) - has_offset_parent(dst) || has_offset_parent(src) || return copyto!(dst, src) + has_offset_parent(dst) || has_offset_parent(src) || copies_by_kernel(backend, src) || + return copyto!(dst, backendless(src) ? collect(src) : src) + d0, s0 = firstindex(dst), firstindex(src) foreachindex_on(Base.OneTo(length(dst)), backend) do i - @inbounds dst[i] = src[i] + @inbounds dst[d0 + i - 1] = src[s0 + i - 1] end return dst end diff --git a/test/Project.toml b/test/Project.toml index 47157986..d930598f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" @@ -22,5 +23,6 @@ AcceleratedKernels = {path = ".."} [compat] Adapt = "4" +MappedArrays = "0.4" ParallelTestRunner = "2" oneAPI = "2.9.2" diff --git a/test/generic/lazy.jl b/test/generic/lazy.jl new file mode 100644 index 00000000..dd789454 --- /dev/null +++ b/test/generic/lazy.jl @@ -0,0 +1,277 @@ +using MappedArrays +using Adapt + + +# A lazy array computing its elements from one array, like MappedArrays' `mappedarray(f, data)`, +# but with an Adapt.jl rule, so that it can be passed to kernels when `data` is a device array +# (MappedArrays.jl does not define one). With `S`, `similar` gives an array like `data` instead of +# Base's host `Array`, and with `C`, `copy` evaluates it on `data`'s backend. +struct LazyMap{T, N, F, A <: AbstractArray, S, C} <: AbstractArray{T, N} + f::F + data::A +end +LazyMap(f, data::AbstractArray{<:Any, N}; similar_data::Bool=false, + copy_data::Bool=similar_data) where {N} = + LazyMap{Base.promote_op(f, eltype(data)), N, typeof(f), typeof(data), similar_data, + copy_data}(f, data) +Base.size(a::LazyMap) = size(a.data) +Base.parent(a::LazyMap) = a.data +Base.IndexStyle(::Type{<:LazyMap{T, N, F, A}}) where {T, N, F, A} = IndexStyle(A) +Base.@propagate_inbounds Base.getindex(a::LazyMap, i::Int) = a.f(a.data[i]) +Base.@propagate_inbounds Base.getindex(a::LazyMap{T, N}, I::Vararg{Int, N}) where {T, N} = + a.f(a.data[I...]) +Base.similar(a::LazyMap{<:Any, <:Any, <:Any, <:Any, true}, ::Type{T}, dims::Dims) where {T} = + similar(a.data, T, dims) +Base.copy(a::LazyMap{<:Any, <:Any, <:Any, <:Any, <:Any, true}) = map(a.f, a.data) +function Adapt.adapt_structure(to, a::LazyMap{T, N, <:Any, <:Any, S, C}) where {T, N, S, C} + f, data = adapt(to, a.f), adapt(to, a.data) + LazyMap{T, N, typeof(f), typeof(data), S, C}(f, data) +end + +# A wrapper of a host array which defines its own backend +struct TaggedTestBackend <: KernelAbstractions.Backend end +struct TaggedArray{T, N} <: AbstractArray{T, N} + data::Array{T, N} +end +Base.size(a::TaggedArray) = size(a.data) +Base.parent(a::TaggedArray) = a.data +Base.getindex(a::TaggedArray, i::Int...) = a.data[i...] +KernelAbstractions.get_backend(::TaggedArray) = TaggedTestBackend() + +# A wrapper that copies itself, but has no Adapt.jl rule to reach kernels +struct CopiedWrapper{T, A <: AbstractVector{T}} <: AbstractVector{T} + data::A +end +Base.size(a::CopiedWrapper) = size(a.data) +Base.parent(a::CopiedWrapper) = a.data +Base.getindex(a::CopiedWrapper, i::Int) = a.data[i] +Base.similar(a::CopiedWrapper, ::Type{T}, dims::Dims) where {T} = similar(a.data, T, dims) +Base.copy(a::CopiedWrapper) = copy(a.data) + +# A vector with offset axes that wraps another, like OffsetArrays' `OffsetVector` +struct OffsetWrapper{T, A <: AbstractVector{T}} <: AbstractVector{T} + data::A + offset::Int +end +Base.axes(a::OffsetWrapper) = (Base.IdentityUnitRange(axes(a.data, 1) .+ a.offset),) +Base.size(a::OffsetWrapper) = size(a.data) +Base.parent(a::OffsetWrapper) = a.data +Base.@propagate_inbounds Base.getindex(a::OffsetWrapper, i::Int) = a.data[i - a.offset] +Base.@propagate_inbounds Base.setindex!(a::OffsetWrapper, v, i::Int) = a.data[i - a.offset] = v +Base.similar(a::OffsetWrapper, ::Type{T}, dims::Dims) where {T} = similar(a.data, T, dims) +Base.copy(a::OffsetWrapper) = OffsetWrapper(copy(a.data), a.offset) +Adapt.adapt_structure(to, a::OffsetWrapper) = OffsetWrapper(adapt(to, a.data), a.offset) + +lazy_sq(x) = x * x +lazy_ci(I) = Int32(I[1] + 2 * I[2]) +lazy_big(x) = x > 50 +lazy_mod(x) = x % Int32(97) + +# Run every read-only operation on the source `s`, comparing with Base on its host equivalent +# `h`, with the keywords `kw` (`backend`, or none); `B` is the backend the results must be on +function test_lazy_ops(s, h, B; kw...) + T = eltype(h) + on_backend(x) = get_backend(x) == B + + # Reductions + @test AK.sum(s; alg=REDUCE_ALG, kw...) == sum(h) + @test AK.mapreduce(x -> 2x, +, s; alg=REDUCE_ALG, init=T(10), kw...) == + mapreduce(x -> 2x, +, h; init=T(10)) + @test AK.maximum(s; alg=REDUCE_ALG, kw...) == maximum(h) + @test AK.count(lazy_big, s; alg=REDUCE_ALG, kw...) == count(lazy_big, h) + @test AK.any(lazy_big, s; kw...) == any(lazy_big, h) + @test AK.all(x -> x > 0, s; kw...) == all(x -> x > 0, h) + # ... with another array + other = array_from_host(T.(eachindex(h))) + @test AK.mapreduce(*, +, s, other; alg=REDUCE_ALG, init=T(0), kw...) == + mapreduce(*, +, h, T.(eachindex(h)); init=T(0)) + + # findall, of the source and of `items` + r = AK.findall(lazy_big, s; kw...) + @test on_backend(r) && Array(r) == findall(lazy_big, h) + mask = array_from_host(isodd.(eachindex(h))) + r = AK.findall(mask; items=s, kw...) + @test on_backend(r) && Array(r) == h[isodd.(eachindex(h))] + + # Scans + r = AK.cumsum(s; kw...) + @test on_backend(r) && Array(r) == cumsum(h) + r = AK.accumulate(+, s; init=T(0), kw...) + @test on_backend(r) && Array(r) == accumulate(+, h; init=T(0)) + + # map, reverse + r = AK.map(x -> 2x, s; kw...) + @test on_backend(r) && Array(r) == map(x -> 2x, h) + dst = array_from_host(zeros(T, size(h))) + @test Array(AK.map!(x -> 2x, dst, s; kw...)) == map(x -> 2x, h) + r = AK.reverse(s; kw...) + @test on_backend(r) && Array(r) == reverse(h) + + # Sorting copies the source + r = AK.sort(s; alg=SORT_ALG, kw...) + @test on_backend(r) && Array(r) == sort(h) + r = AK.sortperm(s; alg=SORT_ALG, kw...) + @test on_backend(r) && h[Array(r)] == sort(h) + if TEST_KERNELS + r = AK.sortperm(s; alg=AK.MergeSort(), kw...) + @test h[Array(r)] == sort(h) + end + + # Binary search, in the source and for its values + sorted = sort(h) + xs = array_from_host(T.(round.(Int, range(0, maximum(h) + 1; length=200)))) + ix = array_from_host(zeros(Int, length(xs))) + AK.searchsortedfirst!(ix, issorted(h) ? s : array_from_host(sorted), xs; kw...) + @test Array(ix) == searchsortedfirst.(Ref(sorted), Array(xs)) + ix = array_from_host(zeros(Int, length(h))) + AK.searchsortedlast!(ix, array_from_host(sorted), s; kw...) + @test Array(ix) == searchsortedlast.(Ref(sorted), h) +end + +# Reductions and scans along dimensions of the matrix source `s` +function test_lazy_dims(s, h, B; kw...) + for dims in (1, 2, (1, 2)) + r = AK.mapreduce(x -> 2x, +, s; dims, init=Int32(10), alg=REDUCE_ALG, kw...) + @test get_backend(r) == B && Array(r) == mapreduce(x -> 2x, +, h; dims, init=Int32(10)) + end + R = array_from_host(ones(Int32, 1, size(h, 2))) + AK.mapreducedim!(identity, +, R, s; alg=REDUCE_ALG, kw...) + @test Array(R) == 1 .+ sum(h; dims=1) + for dims in (1, 2) + r = AK.accumulate(+, s; dims, init=Int32(0), kw...) + @test get_backend(r) == B && Array(r) == accumulate(+, h; dims, init=Int32(0)) + end + for dims in (:, 1, 2) + r = AK.reverse(s; dims, kw...) + @test get_backend(r) == B && Array(r) == reverse(h; dims) + end +end + + +@testset "lazy sources: backend resolution" begin + host = AK.HOST_BACKEND + d = array_from_host(Int32.(1:6)) + # A lazy array over ranges or Cartesian indices does not determine the backend... + @test AK.resolve_backend(nothing, mappedarray(lazy_sq, 1:6)) == host + @test AK.resolve_backend(nothing, mappedarray(lazy_ci, CartesianIndices((2, 3)))) == host + @test AK.resolve_backend(nothing, mappedarray(+, 1:6, 2:7)) == host + @test AK.resolve_backend(nothing, mappedarray(lazy_sq, reshape(1:6, 2, 3))) == host + @test AK.resolve_backend(nothing, LazyMap(lazy_sq, mappedarray(lazy_sq, 1:6))) == host + @test AK.resolve_backend(nothing, d, mappedarray(lazy_sq, 1:6)) == BACKEND + # ... while one over arrays votes with them, whatever it wraps them in + @test AK.resolve_backend(nothing, mappedarray(lazy_sq, d)) == BACKEND + @test AK.resolve_backend(nothing, mappedarray(+, d, 1:6)) == BACKEND + @test AK.resolve_backend(nothing, mappedarray(+, 1:6, d)) == BACKEND + @test AK.resolve_backend(nothing, LazyMap(lazy_sq, mappedarray(+, 1:6, d))) == BACKEND + if BACKEND != host + @test_throws ArgumentError AK.resolve_backend(nothing, mappedarray(+, d, Int32.(1:6))) + end + @test AK.resolve_backend(BACKEND, mappedarray(lazy_sq, 1:6)) == BACKEND + # A wrapper's own `get_backend` method is used + @test AK.resolve_backend(nothing, TaggedArray(zeros(2))) === TaggedTestBackend() + @test AK.resolve_backend(nothing, mappedarray(lazy_sq, TaggedArray(zeros(2)))) === + TaggedTestBackend() + + # Results are allocated like the array holding the source's memory, or on the backend + @test get_backend(AK.similar_on(BACKEND, mappedarray(lazy_sq, d))) == BACKEND + @test get_backend(AK.similar_on(BACKEND, mappedarray(lazy_sq, 1:6))) == BACKEND + @test AK.similar_on(host, mappedarray(+, 1:6, 1:6), Int8) isa Vector{Int8} +end + + +@testset "lazy sources over ranges" begin + # As in #23 + @test AK.sum(1234:100_000; backend=BACKEND, alg=REDUCE_ALG) == sum(1234:100_000) + @test AK.sum(mappedarray(x -> x * x, 1234:100_000); backend=BACKEND, alg=REDUCE_ALG) == + sum(x -> x * x, 1234:100_000) + @test AK.mapreduce(x -> 2x, +, mappedarray(x -> x * x, 1234:100_000); + backend=BACKEND, init=Int64(10), alg=REDUCE_ALG) == + sum(x -> 2 * x * x, 1234:100_000; init=Int64(10)) + + # Without a backend, on the host + @test AK.sum(mappedarray(lazy_sq, Int32(1):Int32(6))) == 91 + + n = 1000 + for (s, h) in ((mappedarray(lazy_sq, Int32(1):Int32(n)), map(lazy_sq, Int32(1):Int32(n))), + (mappedarray(+, Int32(1):Int32(n), Int32(2):Int32(n + 1)), collect(Int32(3):Int32(2):Int32(2n + 1))), + (mappedarray(lazy_mod, Int32(1):Int32(n)), map(lazy_mod, Int32(1):Int32(n)))) + # The backend given explicitly... + test_lazy_ops(s, h, BACKEND; backend=BACKEND) + # ... and on the host without it, where the host backend has no kernels + TEST_KERNELS || test_lazy_ops(s, h, BACKEND) + end + + C = CartesianIndices((37, 81)) + test_lazy_dims(mappedarray(lazy_ci, C), map(lazy_ci, C), BACKEND; backend=BACKEND) + R = reshape(Int32(1):Int32(37 * 81), 37, 81) + test_lazy_dims(mappedarray(lazy_mod, R), map(lazy_mod, R), BACKEND; backend=BACKEND) + TEST_KERNELS || test_lazy_dims(mappedarray(lazy_ci, C), map(lazy_ci, C), BACKEND) + # CartesianIndices themselves, with offset axes (as in #23) + C = CartesianIndices((12:345, 67:89)) + for dims in (:, 1, 2, (1, 2)) + r = AK.mapreduce(I -> I[1], +, C; dims, init=10, backend=BACKEND, alg=REDUCE_ALG) + @test (dims isa Colon ? r : Array(r)) == mapreduce(I -> I[1], +, C; dims, init=10) + end +end + + +@testset "lazy sources over arrays" begin + n = 3000 + h = Int32.(rand(1:100, n)) + d = array_from_host(h) + # A lazy array of the backend's arrays runs on their backend, given or not + test_lazy_ops(LazyMap(lazy_sq, d), map(lazy_sq, h), BACKEND) + test_lazy_ops(LazyMap(lazy_sq, d), map(lazy_sq, h), BACKEND; backend=BACKEND) + # ... also when its `similar` gives the backend's arrays + test_lazy_ops(LazyMap(lazy_sq, d; similar_data=true), map(lazy_sq, h), BACKEND) + # ... and nested + test_lazy_ops(LazyMap(x -> x + Int32(1), LazyMap(lazy_sq, d)), map(lazy_sq, h) .+ Int32(1), BACKEND) + + H = Int32.(rand(1:100, 37, 81)) + D = array_from_host(H) + test_lazy_dims(LazyMap(lazy_sq, D), map(lazy_sq, H), BACKEND) + test_lazy_dims(LazyMap(lazy_sq, view(D, 2:30, 3:70)), map(lazy_sq, H[2:30, 3:70]), BACKEND) + + if TEST_KERNELS + # Reductions finishing on the host evaluate the lazy array on the device first + alg = AK.BlockReduce(switch_below=100) + @test AK.sum(LazyMap(lazy_sq, d[1:50]); alg) == sum(lazy_sq, h[1:50]) + @test AK.mapreduce(*, +, LazyMap(lazy_sq, d[1:50]), d[1:50]; alg) == + sum(map(lazy_sq, h[1:50]) .* h[1:50]) + lazy = LazyMap(lazy_sq, d[1:50]; similar_data=true, copy_data=false) + @test AK.sum(lazy; alg) == sum(lazy_sq, h[1:50]) + else + # MappedArrays' own types can only be passed to kernels over ranges and other isbits + # arrays (MappedArrays.jl has no Adapt.jl rules), but work with the host's arrays + test_lazy_ops(mappedarray(lazy_sq, d), map(lazy_sq, h), BACKEND) + test_lazy_dims(mappedarray(lazy_sq, D), map(lazy_sq, H), BACKEND) + end +end + + +@testset "lazy sources: wrappers that copy themselves" begin + h = Int32.(rand(1:100, 1000)) + # Copied with their own `copy`, also when they cannot be passed to kernels + r = AK.sort(CopiedWrapper(array_from_host(h)); alg=SORT_ALG) + @test get_backend(r) == BACKEND && Array(r) == sort(h) + # ... which keeps their axes + r = AK.sort(OffsetWrapper(view(array_from_host(h), 11:20), -3); alg=SORT_ALG) + @test axes(r) == (-2:7,) && Array(parent(r)) == sort(h[11:20]) +end + + +@testset "lazy sources: copies of host wrappers" begin + # Base's copies, which do not race on packed storage, are used on the host backend (the + # sort itself runs on one task: sorting a `BitVector` on several races) + alg = AK.CPUThreads.SampleSort(max_tasks=1) + for _ in 1:20 + bits = view(BitVector(rand(Bool, 1000)), :) + @test AK.sort(bits; backend=AK.HOST_BACKEND, alg) == sort(collect(bits)) + @test bits[AK.sortperm(bits; backend=AK.HOST_BACKEND)] == sort(collect(bits)) + # ... and results are not allocated in packed storage, which `map!` would write to from + # several tasks + lazybits = LazyMap(identity, BitVector(rand(Bool, 1000))) + r = AK.map(identity, lazybits; backend=AK.HOST_BACKEND) + @test r isa Vector{Bool} && r == lazybits + end +end