Launch kernels over a list of indices

docs/src/api.md

@@ -91,4 +91,13 @@ KernelAbstractions.@context
 KernelAbstractions.argconvert
 KernelAbstractions.NDIteration.DynamicSize
 KernelAbstractions.NDIteration.StaticSize
+KernelAbstractions.NDIteration.NDRange
+KernelAbstractions.NDIteration.StaticOffset
+KernelAbstractions.NDIteration.DynamicOffset
+KernelAbstractions.NDIteration.extents
+KernelAbstractions.NDIteration.offsets
+KernelAbstractions.NDIteration.linear_index
+KernelAbstractions.NDIteration.IndexMap
+KernelAbstractions.NDIteration.MappedNDRange
+KernelAbstractions.__validindex
 ```

docs/src/index.md

@@ -132,6 +132,15 @@ but users must avoid the use of `@index(Global)` and instead use their own deriv
 end
 ```
 
+### 0.10
+- `ndrange` entries may be index ranges, given statically (`kernel(backend, workgroupsize, (-2:N+3, 0:M+1))`)
+  or at launch (`ndrange=(-2:N+3, 0:M+1)`, a single range, or a `CartesianIndices`).
+  `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted indices.
+- `ndrange` may be a device vector of `CartesianIndex`/`NTuple` indices, running one work item per
+  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.
+
 ## Semantic differences
 
 ### To CUDA.jl/AMDGPU.jl

docs/src/kernels.md

@@ -232,5 +232,44 @@ kernel = my_kernel(backend, 32, size(A))
 kernel(A)
 ```
 
+### Index ranges
+
+Each entry of `ndrange` is either an extent (indices `1:n`) or a range of indices, so a kernel
+can iterate over a region whose indices do not start at 1. `ndrange` can also be given as a
+single range or as a `CartesianIndices`:
+
+```julia
+# static ndrange over the indices -2:N+3 along x and 0:M+1 along y
+kernel = my_kernel(backend, (16, 16), (-2:N+3, 0:M+1))
+kernel(A)
+
+# dynamic ndrange
+kernel = my_kernel(backend, (16, 16))
+kernel(A, ndrange=(-2:N+3, 0:M+1))
+kernel(A, ndrange=CartesianIndices(A))   # e.g. for an OffsetArray
+```
+
+Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted
+indices, `@index(Global, Linear)` counts the indices of the region from 1 in column-major order,
+and `@ndrange()` returns the extents.
+
+### Index maps
+
+`ndrange` can be a vector of indices (`CartesianIndex{N}` or `NTuple{N, <:Integer}` elements,
+stored on the backend's device) to run one work item per listed index, for example over the
+active cells of a masked domain:
+
+```julia
+active = CuArray([CartesianIndex(i, j, k) for (i, j, k) in cells if mask[i, j, k]])
+kernel = my_kernel(backend, 256)          # 1-D workgroup size is required
+kernel(A, ndrange=active)
+```
+
+Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the listed
+index, `@index(Global, Linear)` its position in the vector, and `@ndrange()` the length of
+the vector. The kernel must be constructed with a dynamic `ndrange`, and the workgroup size
+must be static or given with `workgroupsize`. With `@kernel unsafe_indices=true` the work
+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.

src/KernelAbstractions.jl

@@ -412,9 +412,7 @@ end
 end
 
 @inline function __index_Global_Linear(ctx)
-    I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x)
-    # TODO: This is unfortunate, can we get the linear index cheaper
-    return @inbounds LinearIndices(__ndrange(ctx))[I]
+    return __global_linear(__iterspace(ctx), __ndrange(ctx), KI.get_group_id().x, KI.get_local_id().x)
 end
 
 @inline function __index_Local_Cartesian(ctx)
@@ -544,6 +542,12 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
 @inline function partition(kernel, ndrange, workgroupsize)
     static_ndrange = KernelAbstractions.ndrange(kernel)
     static_workgroupsize = KernelAbstractions.workgroupsize(kernel)
+    ndrange = NDIteration.normalize_ndrange(ndrange)
+    workgroupsize = NDIteration.normalize_workgroupsize(workgroupsize)
+
+    if ndrange isa IndexMap
+        return mapped_partition(kernel, ndrange, workgroupsize)
+    end
 
     if ndrange === nothing && static_ndrange <: DynamicSize ||
             workgroupsize === nothing && static_workgroupsize <: DynamicSize
@@ -562,7 +566,7 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
     end
 
     if static_ndrange <: StaticSize
-        if ndrange !== nothing && ndrange != get(static_ndrange)
+        if ndrange !== nothing && !NDIteration.same_axes(ndrange, get(static_ndrange))
             error("Static NDRange ($static_ndrange) and launch NDRange ($ndrange) differ")
         end
         ndrange = get(static_ndrange)
@@ -577,14 +581,16 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
 
     @assert workgroupsize !== nothing
     @assert ndrange !== nothing
-    blocks, workgroupsize, dynamic = NDIteration.partition(ndrange, workgroupsize)
+    blocks, workgroupsize, dynamic = NDIteration.partition(extents(ndrange), workgroupsize)
 
     if static_ndrange <: StaticSize
         static_blocks = StaticSize{blocks}
         blocks = nothing
+        mapping = NDIteration.static_mapping(ndrange)
     else
         static_blocks = DynamicSize
         blocks = CartesianIndices(blocks)
+        mapping = NDIteration.dynamic_mapping(ndrange)
     end
 
     if static_workgroupsize <: StaticSize
@@ -594,7 +600,40 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
         workgroupsize = CartesianIndices(workgroupsize)
     end
 
-    iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize)
+    iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize, mapping)
+    return iterspace, dynamic
+end
+
+# Partition of an index map: a 1-D blocked space over the positions in the map.
+@inline function mapped_partition(kernel, map::IndexMap, workgroupsize)
+    static_ndrange = KernelAbstractions.ndrange(kernel)
+    static_workgroupsize = KernelAbstractions.workgroupsize(kernel)
+
+    if static_ndrange <: StaticSize
+        error("An index map is a runtime iteration space; construct the kernel with a dynamic ndrange")
+    end
+    if static_workgroupsize <: StaticSize
+        if workgroupsize !== nothing && workgroupsize != get(static_workgroupsize)
+            error("Static WorkgroupSize ($static_workgroupsize) and launch WorkgroupSize $(workgroupsize) differ")
+        end
+        workgroupsize = get(static_workgroupsize)
+    elseif !(workgroupsize isa Tuple)
+        error("An index map requires a workgroup size, either static or given with `workgroupsize`")
+    end
+    if length(workgroupsize) != 1
+        error("An index map requires a 1-D workgroup size, got $(workgroupsize)")
+    end
+
+    blocks, workgroupsize, dynamic = NDIteration.partition((length(map),), workgroupsize)
+
+    if static_workgroupsize <: StaticSize
+        static_workgroupsize = StaticSize{workgroupsize}
+        workgroupsize = nothing
+    else
+        workgroupsize = CartesianIndices(workgroupsize)
+    end
+
+    iterspace = NDRange{1, DynamicSize, static_workgroupsize}(CartesianIndices(blocks), workgroupsize, map)
     return iterspace, dynamic
 end
 
@@ -613,7 +652,37 @@ include("compiler.jl")
 ###
 
 function __workitems_iterspace end
-function __validindex end
+
+"""
+    __validindex(ctx, groupidx, idx)
+
+Whether work item `idx` of workgroup `groupidx` has an index within the `ndrange`.
+Both indices are linear or `CartesianIndex` positions within the blocked iteration space.
+"""
+@inline function __validindex(ctx, groupidx, idx)
+    if __dynamic_checkbounds(ctx)
+        return __inrange(__iterspace(ctx), __ndrange(ctx), groupidx, idx)
+    else
+        return true
+    end
+end
+
+@inline function __validindex(ctx)
+    return __validindex(ctx, KI.get_group_id().x, KI.get_local_id().x)
+end
+
+@inline function __inrange(iterspace::NDRange, ndrange, groupidx, idx)
+    I = @inbounds expand(iterspace, groupidx, idx)
+    return I in ndrange
+end
+@inline __inrange(iterspace::MappedNDRange, ndrange, groupidx, idx) = linear_index(iterspace, groupidx, idx) <= length(iterspace.mapping)
+
+# Global linear index of work item `idx` of workgroup `groupidx`.
+@inline function __global_linear(iterspace::NDRange, ndrange, groupidx, idx)
+    I = @inbounds expand(iterspace, groupidx, idx)
+    return linear_index(ndrange, I)
+end
+@inline __global_linear(iterspace::MappedNDRange, ndrange, groupidx, idx) = linear_index(iterspace, groupidx, idx)
 
 # for reflection
 function mkcontext end

src/compiler.jl

@@ -5,21 +5,29 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
 
     # CPU variant
     function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB}
-        if ndrange !== nothing
-            ndrange = CartesianIndices(ndrange)
-        end
+        ndrange = cartesian(ndrange)
         return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace)
     end
 
     # GPU variante: index is given implicit
     function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB}
-        if ndrange !== nothing
-            ndrange = CartesianIndices(ndrange)
-        end
+        ndrange = cartesian(ndrange)
         return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace)
     end
 end
 
+# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`).
+cartesian(::Nothing) = nothing
+cartesian(ci::CartesianIndices) = ci
+cartesian(n::Integer) = CartesianIndices((Int(n),))
+cartesian(r::AbstractUnitRange) = CartesianIndices((r,))
+cartesian(t::Tuple) = CartesianIndices(t)
+cartesian(v::AbstractVector) = CartesianIndices((length(v),))
+cartesian(m::IndexMap) = CartesianIndices((length(m),))
+
+Adapt.adapt_structure(to, cm::CompilerMetadata{NDRange, CB}) where {NDRange, CB} =
+    CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, Adapt.adapt(to, cm.iterspace))
+
 @inline __iterspace(cm::CompilerMetadata) = cm.iterspace
 @inline __groupindex(cm::CompilerMetadata) = cm.groupindex
 @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm)))