Add foreach_index

docs/src/api.md

@@ -70,6 +70,12 @@ KernelAbstractions.ndrange
 KernelAbstractions.backend
 ```
 
+### Index loops
+
+```@docs
+foreach_index
+```
+
 ## Reflection
 
 These macros help inspect the generated kernel code. LLVM IR reflection via

docs/src/index.md

@@ -140,6 +140,7 @@ end
   listed index (`kernel(A, ndrange=active_cells)`).
 - Index validity is decided by the generic `__validindex(ctx, groupidx, idx)`, so backends only
   supply the hardware indices.
+- `foreach_index(f, itr)` runs `f` once per index of `itr` without writing a kernel out.
 
 ## Semantic differences
 

docs/src/kernels.md

@@ -273,3 +273,27 @@ items of a partial last workgroup have no valid index.
 
 Obtain the backend from an array with [`get_backend`](@ref) and always call [`synchronize`](@ref) before reading results on the host.
 See the [Quickstart](@ref) for a full walkthrough and the Examples section of the manual for larger patterns.
+
+## Loops without a kernel
+
+A kernel whose body is a loop over the indices of an array needs none of the kernel language
+beyond the index itself. [`foreach_index`](@ref) launches such a loop directly, with one work item
+per index of `eachindex(itr)`:
+
+```julia
+function scale!(y, x)
+    foreach_index(x) do i
+        @inbounds y[i] = 2 * x[i] + 1
+    end
+    return y
+end
+```
+
+The body is an ordinary Julia function, and it receives the index a `for i in eachindex(itr)` loop
+would: a linear index for an array with `IndexLinear` style, a `CartesianIndex` otherwise. The
+launch is asynchronous like any other, and bounds checks are not elided.
+
+Because the body becomes a kernel, every value it captures must have a known type — which is why
+the example wraps the loop in a function, and why the body must not assign to a captured variable.
+Write the kernel out with [`@kernel`](@ref) when it needs more than an index: workgroup-level
+indices, local memory, or synchronization.

src/KernelAbstractions.jl

@@ -6,6 +6,7 @@ export @index, @groupsize, @ndrange
 export @print
 export Backend, GPU, CPU
 export synchronize, get_backend, allocate
+export foreach_index
 
 import PrecompileTools
 
@@ -689,6 +690,7 @@ function mkcontext end
 function launch_config end
 
 include("macros.jl")
+include("foreach_index.jl")
 
 ###
 # Backends/Interface

src/foreach_index.jl

@@ -0,0 +1,75 @@
+# One work item per index of a collection.
+#
+# The index space is carried by the `ndrange`, so the kernels take no argument besides the
+# function: `@index(Global, Cartesian)` already returns the index of the `ndrange`, including
+# the offset of an index space that does not start at 1.
+#
+# `f` is inlined on purpose: an out-of-line call to a closure can spill its captures to local
+# memory on some backends.
+
+@kernel function foreach_index_linear_kernel(f)
+    I = @index(Global, Cartesian)
+    @inline f(I.I[1])
+end
+
+@kernel function foreach_index_cartesian_kernel(f)
+    I = @index(Global, Cartesian)
+    @inline f(I)
+end
+
+# `eachindex` of an `AbstractArray` is either a range of linear indices or a `CartesianIndices`;
+# hand `f` the index type that `itr` is indexed with in either case.
+foreach_index_kernel(::AbstractUnitRange, backend) = foreach_index_linear_kernel(backend)
+foreach_index_kernel(::CartesianIndices, backend) = foreach_index_cartesian_kernel(backend)
+foreach_index_kernel(indices, backend) = throw(
+    ArgumentError(
+        "`foreach_index` needs an index space that is a range of linear indices or a `CartesianIndices`, got `$(typeof(indices))`"
+    )
+)
+
+"""
+    foreach_index(f, itr, backend = get_backend(itr); workgroupsize = nothing)
+
+Call `f(i)` once for every `i` in `eachindex(itr)`, with one work item per index.
+
+This is the kernel-free spelling of a `for` loop over the indices of a collection: the body is
+an ordinary Julia function, so the same code runs on every backend.
+
+```julia
+function scale!(y, x)
+    foreach_index(x) do i
+        @inbounds y[i] = 2 * x[i] + 1
+    end
+    return y
+end
+```
+
+`f` receives the same index that a `for i in eachindex(itr)` loop would: a linear index for an
+array with `IndexLinear` style, a `CartesianIndex` otherwise.
+
+Like any other kernel launch this is asynchronous; call [`synchronize`](@ref) before reading the
+result on the host. Bounds checks are not elided, so write `@inbounds` in the body where it is
+warranted, as in a hand-written kernel. `workgroupsize` is passed on to the launch, and by
+default the backend chooses it.
+
+# Extended help
+
+`f` becomes the body of a kernel, so it is subject to the same restrictions: every value it
+captures must be of a known type. Closing over a variable of the enclosing *global* scope leaves
+its type unknown and fails to compile (typically with `unsupported dynamic function invocation`),
+which is why the example above wraps the loop in a function.
+
+For the same reason `f` must not assign to a captured variable, as that makes Julia box the
+capture; accumulate into a one-element array, with an atomic update if the indices race.
+
+See also [`@kernel`](@ref) to write the kernel out, which is what to reach for when the body
+needs more of the kernel language than an index (workgroup-level indices, local memory, or
+synchronization).
+"""
+function foreach_index(f::F, itr, backend::Backend = get_backend(itr); workgroupsize = nothing) where {F}
+    indices = eachindex(itr)
+    isempty(indices) && return nothing
+    kernel = foreach_index_kernel(indices, backend)
+    kernel(f; ndrange = indices, workgroupsize)
+    return nothing
+end

test/foreach_index.jl

@@ -0,0 +1,141 @@
+using KernelAbstractions
+using KernelAbstractions: @atomic
+using Adapt
+using Test
+
+# An `IndexLinear` array whose `axes1` does not start at 1, so that `eachindex` returns a
+# range of linear indices with an offset.
+struct OffsetVec{T, A <: AbstractVector{T}} <: AbstractVector{T}
+    data::A
+    r::Base.IdentityUnitRange{UnitRange{Int}}
+end
+OffsetVec(data::AbstractVector, off::Int) = OffsetVec(data, Base.IdentityUnitRange((1 + off):(length(data) + off)))
+Base.size(v::OffsetVec) = size(v.data)
+Base.axes(v::OffsetVec) = (v.r,)
+Base.IndexStyle(::Type{<:OffsetVec}) = IndexLinear()
+Base.@propagate_inbounds Base.getindex(v::OffsetVec, i::Int) = v.data[i - first(v.r) + 1]
+Base.@propagate_inbounds Base.setindex!(v::OffsetVec, x, i::Int) = (v.data[i - first(v.r) + 1] = x)
+KernelAbstractions.get_backend(v::OffsetVec) = KernelAbstractions.get_backend(v.data)
+Adapt.adapt_structure(to, v::OffsetVec) = OffsetVec(adapt(to, v.data), v.r)
+
+# The loop bodies live in functions so that their captures have known types, as the
+# `foreach_index` docstring requires.
+
+function foreach_index_copy!(dst, src)
+    foreach_index(src) do i
+        @inbounds dst[i] = src[i]
+    end
+    return dst
+end
+
+function foreach_index_copy_backend!(dst, src, backend)
+    foreach_index(src, backend) do i
+        @inbounds dst[i] = src[i]
+    end
+    return dst
+end
+
+function foreach_index_copy_workgroupsize!(dst, src, workgroupsize)
+    foreach_index(src; workgroupsize) do i
+        @inbounds dst[i] = src[i]
+    end
+    return dst
+end
+
+function foreach_index_fill_index!(v)
+    foreach_index(v) do i
+        @inbounds v[i] = i
+    end
+    return v
+end
+
+function foreach_index_mark!(out, itr)
+    foreach_index(itr) do I
+        @inbounds out[I] += 1
+    end
+    return out
+end
+
+function foreach_index_sum_range!(out, range, backend)
+    foreach_index(range, backend) do i
+        @inbounds @atomic out[1] += i
+    end
+    return out
+end
+
+function foreach_index_testsuite(Backend, AT)
+    backend = Backend()
+
+    @testset "linear index space" begin
+        @testset "$(length(dims))-D" for dims in ((1024,), (32, 33), (4, 5, 6))
+            src = AT(reshape(collect(1:prod(dims)), dims))
+            dst = AT(zeros(Int, dims))
+            @test foreach_index_copy!(dst, src) === dst
+            synchronize(backend)
+            @test Array(dst) == Array(src)
+
+            # every index is visited exactly once
+            out = AT(zeros(Int, dims))
+            foreach_index_mark!(out, src)
+            synchronize(backend)
+            @test all(isone, Array(out))
+        end
+    end
+
+    @testset "cartesian index space" begin
+        A = AT(zeros(Int, 8, 10))
+        v = view(A, 1:2:8, 2:2:10)
+        @test eachindex(v) isa CartesianIndices
+        foreach_index_mark!(v, v)
+        synchronize(backend)
+        ref = zeros(Int, 8, 10)
+        ref[1:2:8, 2:2:10] .= 1
+        @test Array(A) == ref
+    end
+
+    @testset "offset linear index space" begin
+        v = OffsetVec(AT(zeros(Int, 10)), -3)
+        @test eachindex(v) == -2:7
+        foreach_index_fill_index!(v)
+        synchronize(backend)
+        @test Array(v.data) == collect(-2:7)
+    end
+
+    @testset "index space without device memory" begin
+        # a range carries no backend of its own, so it has to be given
+        out = AT(zeros(Int, 1))
+        foreach_index_sum_range!(out, 1:100, backend)
+        synchronize(backend)
+        @test Array(out)[1] == sum(1:100)
+    end
+
+    @testset "workgroupsize" begin
+        src = AT(collect(1:1000))
+        dst = AT(zeros(Int, 1000))
+        foreach_index_copy_workgroupsize!(dst, src, 8)
+        synchronize(backend)
+        @test Array(dst) == Array(src)
+    end
+
+    @testset "explicit backend" begin
+        src = AT(collect(1:100))
+        dst = AT(zeros(Int, 100))
+        foreach_index_copy_backend!(dst, src, backend)
+        synchronize(backend)
+        @test Array(dst) == Array(src)
+    end
+
+    @testset "empty index space" begin
+        src = AT(Int[])
+        dst = AT(Int[])
+        @test foreach_index_copy!(dst, src) === dst
+        synchronize(backend)
+        @test isempty(Array(dst))
+    end
+
+    @testset "errors" begin
+        # an index space that is neither linear nor cartesian
+        @test_throws ArgumentError foreach_index(identity, Dict(1 => 2), backend)
+    end
+    return
+end

test/testsuite.jl

@@ -35,6 +35,7 @@ include("unroll.jl")
 include("nditeration.jl")
 include("offsets.jl")
 include("indexmap.jl")
+include("foreach_index.jl")
 include("copyto.jl")
 include("devices.jl")
 include("print_test.jl")
@@ -85,6 +86,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{
         indexmap_testsuite(backend, AT)
     end
 
+    @conditional_testset "foreach_index" skip_tests begin
+        foreach_index_testsuite(backend, AT)
+    end
+
     @conditional_testset "copyto!" skip_tests begin
         copyto_testsuite(backend, AT)
     end