Support lazy arrays such as MappedArrays as read-only sources

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
 

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

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...)

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]

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)`

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

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"