From 274b5e13920d6844d78309a092d1c6290dc787b6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 26 Oct 2026 20:13:47 +0100 Subject: [PATCH] Support lazy arrays such as MappedArrays as read-only sources Backend resolution asked every array that is not a range, CartesianIndices or LinearIndices (or a Base view, reshape or permutation of one) for its backend with KernelAbstractions' `get_backend`, which recurses through `parent` and throws when it reaches a range or a tuple of parents. So `AK.sum(mappedarray(abs2, 1:n))`, and every other operation on a MappedArrays.jl array over ranges or Cartesian indices, threw "Implement get_backend" without an explicit `backend`, and `mappedarray(+, A, B)` threw a MethodError. With `backend` given, reductions, predicates, findall and searchsorted already worked on such arrays, since an `isbits` mapped array is passed to kernels as it is. But the operations that allocate a result (map, accumulate, reverse, sort, sortperm) used `similar(v)`, which for a wrapper without its own `similar` is Base's host `Array`, so their kernels got a host array. Copies of such wrappers (`sort`, the keys of the merge sort permutation, the host finish of small reductions) went through `copy`, `copyto!` or `Array`, whose fallbacks scalar-index device arrays. Now a wrapper whose parents are not a single array with memory votes like the arrays it wraps, so a wrapper of only ranges does not vote, as Base's wrappers of ranges already did not; other arrays are still asked with `get_backend`. A wrapper that gets Base's `similar` fallback, or wraps no memory, has its results allocated on `backend`, and is copied with a kernel there; other wrappers keep their own `copy`, so an OffsetArray's copy keeps its axes. Copies into existing arrays and to the host also use a kernel for wrappers that GPUArrays (its wrappers of device arrays) and Base (wrappers of host memory) do not copy, since Base's fallbacks index device memory element by element. Lazy arrays over ranges are not evaluated on the host either, since their function may read device arrays it captures. A lazy array of device arrays also needs an Adapt.jl rule for its type to reach kernels. MappedArrays.jl has none, so `mappedarray(f, ::CuArray)` still only works on the host; that rule belongs in MappedArrays.jl. The tests use a lazy wrapper with such a rule for the device path. Supersedes #23, which allocated the 1D reduction's scratch on the given backend instead of with `similar(src)`, and whose range and mapped-array sums are among the tests. Co-authored-by: Dan Stahlke --- docs/src/api/algorithms.md | 22 ++- src/AcceleratedKernels.jl | 2 +- src/algorithms.jl | 22 ++- src/reduce/mapreduce_1d_gpu.jl | 19 +-- src/sort/merge_sort.jl | 3 +- src/utils.jl | 61 +++++++- test/Project.toml | 2 + test/generic/lazy.jl | 277 +++++++++++++++++++++++++++++++++ 8 files changed, 378 insertions(+), 30 deletions(-) create mode 100644 test/generic/lazy.jl 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