Launch kernels over a list of indices

docs/src/api.md

@@ -96,5 +96,11 @@ KernelAbstractions.NDIteration.StaticOffset
 KernelAbstractions.NDIteration.DynamicOffset
 KernelAbstractions.NDIteration.extents
 KernelAbstractions.NDIteration.offsets
+KernelAbstractions.NDIteration.expand
 KernelAbstractions.NDIteration.linear_index
+KernelAbstractions.NDIteration.IndexMap
+KernelAbstractions.NDIteration.MappedNDRange
+KernelAbstractions.NDIteration.MappedIndices
+KernelAbstractions.NDIteration.invalid_index
+KernelAbstractions.__validindex
 ```

docs/src/index.md

@@ -136,6 +136,10 @@ end
 - `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
 

docs/src/kernels.md

@@ -253,5 +253,23 @@ Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` retur
 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 get an index of `typemin(Int)` along every axis.
+
 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

@@ -398,9 +398,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 linear_index(__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)
@@ -533,6 +531,10 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
     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
         errmsg = """
@@ -588,6 +590,39 @@ last (possibly partial) workgroup. Primarily used by backend implementations and
     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
+
 function construct(backend::Backend, ::S, ::NDRange, xpu_name::XPUName) where {Backend <: GPU, S <: _Size, NDRange <: _Size, XPUName}
     return Kernel{Backend, S, NDRange, XPUName}(backend, xpu_name)
 end
@@ -603,7 +638,39 @@ 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
+
+# Backends may override `__validindex(ctx)` with this same check, so `expand` yields an
+# index `in` `__ndrange(ctx)` exactly for the valid work items of every kind of `ndrange`.
+@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

@@ -16,12 +16,18 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
     end
 end
 
-# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`).
+# `CartesianIndices` or `MappedIndices` covering a launch `ndrange` (any form accepted by `partition`).
 cartesian(::Nothing) = nothing
 cartesian(ci::CartesianIndices) = ci
+cartesian(r::MappedIndices) = r
 cartesian(n::Integer) = CartesianIndices((Int(n),))
 cartesian(r::AbstractUnitRange) = CartesianIndices((r,))
 cartesian(t::Tuple) = CartesianIndices(t)
+cartesian(v::AbstractVector) = MappedIndices(v)
+cartesian(m::IndexMap) = MappedIndices(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

src/nditeration.jl

@@ -3,9 +3,11 @@ module NDIteration
 export _Size, StaticSize, DynamicSize, get
 export NDRange, blocks, workitems, expand
 export StaticOffset, DynamicOffset, offsets, extents, linear_index
+export IndexMap, MappedNDRange, MappedIndices, invalid_index
 export DynamicCheck, NoDynamicCheck
 
 import Base.@pure
+import Adapt
 
 struct DynamicCheck end
 struct NoDynamicCheck end
@@ -31,6 +33,7 @@ extents(t::Tuple) = map(extent, t)
 extents(ci::CartesianIndices) = size(ci)
 extents(r::AbstractUnitRange) = (length(r),)
 extents(n::Integer) = (Int(n),)
+extents(v::AbstractVector) = (length(v),)
 
 """
     offsets(ndrange)
@@ -39,26 +42,84 @@ Offset of the first index along each axis of `ndrange` relative to 1.
 """
 offsets(t::Tuple) = map(axis_offset, t)
 
+"""
+    IndexMap{N}(map::AbstractVector)
+
+Iteration space given by the indices listed in `map`, whose elements are `CartesianIndex{N}`
+or `NTuple{N, <:Integer}`. Work item `p` handles the index `map[p]`.
+"""
+struct IndexMap{N, A <: AbstractVector}
+    map::A
+    IndexMap{N}(map::AbstractVector) where {N} = new{N, typeof(map)}(map)
+end
+IndexMap(map::AbstractVector) = IndexMap{mapdims(eltype(map))}(map)
+
+mapdims(::Type{CartesianIndex{N}}) where {N} = N
+mapdims(::Type{<:NTuple{N, Integer}}) where {N} = N
+mapdims(::Type{T}) where {T} = throw(ArgumentError("an index map must have elements of type `CartesianIndex{N}` or `NTuple{N, Integer}`, got `$T`"))
+
+Base.length(m::IndexMap) = length(m.map)
+extents(m::IndexMap) = (length(m),)
+Base.@propagate_inbounds Base.getindex(m::IndexMap{N}, i::Integer) where {N} = mapindex(Val(N), m.map[i])
+mapindex(::Val{N}, I::CartesianIndex{N}) where {N} = I
+mapindex(::Val{N}, I::Tuple) where {N} = CartesianIndex{N}(I)
+
+Adapt.adapt_structure(to, m::IndexMap{N}) where {N} = IndexMap{N}(Adapt.adapt(to, m.map))
+
+"""
+    invalid_index(Val(N))
+
+`CartesianIndex{N}` returned by [`expand`](@ref) for a work item past the end of an
+[`IndexMap`](@ref): `typemin(Int)` along every axis.
+"""
+@inline invalid_index(::Val{N}) where {N} = CartesianIndex(ntuple(_ -> typemin(Int), Val(N)))
+
+# Entry `p` of the map, or `invalid_index` for a position past its end.
+@inline function mapped_index(m::IndexMap{N}, p::Integer) where {N}
+    return p <= length(m) ? (@inbounds m[p]) : invalid_index(Val(N))
+end
+
+"""
+    MappedIndices{N}
+
+The `ndrange` of a launch over an [`IndexMap`](@ref) with `CartesianIndex{N}` entries.
+`size` and `length` give the number of listed indices, and a `CartesianIndex{N}` is `in`
+it unless it is the [`invalid_index`](@ref), so that backends can check the validity of a
+work item as `expand(iterspace, groupidx, idx) in ndrange` for every kind of `ndrange`.
+"""
+struct MappedIndices{N}
+    length::Int
+end
+MappedIndices(m::IndexMap{N}) where {N} = MappedIndices{N}(length(m))
+MappedIndices(v::AbstractVector) = MappedIndices{mapdims(eltype(v))}(length(v))
+Base.length(r::MappedIndices) = r.length
+Base.size(r::MappedIndices) = (r.length,)
+@inline Base.in(I::CartesianIndex{N}, r::MappedIndices{N}) where {N} = I != invalid_index(Val(N))
+
 """
     normalize_ndrange(ndrange)
 
-Canonical form of a launch `ndrange`: `nothing`, or a tuple of `Int` extents and
-`UnitRange{Int}` axes.
+Canonical form of a launch `ndrange`: `nothing`, a tuple of `Int` extents and
+`UnitRange{Int}` axes, or an [`IndexMap`](@ref).
 """
 normalize_ndrange(::Nothing) = nothing
 normalize_ndrange(n::Integer) = (Int(n),)
 normalize_ndrange(r::AbstractUnitRange) = (axis(r),)
 normalize_ndrange(ci::CartesianIndices) = map(axis, ci.indices)
 normalize_ndrange(t::Tuple) = map(axis, t)
+normalize_ndrange(m::IndexMap) = m
+normalize_ndrange(v::AbstractVector) = IndexMap(v)
 
 """
     normalize_workgroupsize(workgroupsize)
 
 Canonical form of a launch `workgroupsize`: `nothing`, or a tuple of `Int` extents.
+Anything else is passed through to be rejected by `partition`.
 """
 normalize_workgroupsize(::Nothing) = nothing
 normalize_workgroupsize(n::Integer) = (Int(n),)
 normalize_workgroupsize(t::Tuple) = extents(t)
+normalize_workgroupsize(x) = x
 
 # Two ndranges denote the same indices.
 same_axes(a::Tuple, b::Tuple) = extents(a) == extents(b) && offsets(a) == offsets(b)
@@ -137,8 +198,9 @@ end
     NDRange
 
 Encodes a blocked iteration space. The `mapping` field relates blocked indices to
-`ndrange` indices: `nothing` for the identity, or a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref)
-for an `ndrange` whose indices do not start at 1.
+`ndrange` indices: `nothing` for the identity, a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref)
+for an `ndrange` whose indices do not start at 1, or an [`IndexMap`](@ref) for a 1-D
+blocked space whose work items look up their index in a list.
 
 # Example
 ```
@@ -182,12 +244,32 @@ static_mapping(t::Tuple) = StaticOffset{offsets(t)}()
 dynamic_mapping(::Tuple{Vararg{Int}}) = nothing
 dynamic_mapping(t::Tuple) = DynamicOffset(offsets(t))
 
+"""
+    MappedNDRange
+
+A 1-D blocked iteration space whose `mapping` is an [`IndexMap`](@ref). [`expand`](@ref)
+looks up the index listed for a work item, and the `ndrange` of such a launch is a
+[`MappedIndices`](@ref).
+"""
+const MappedNDRange = NDRange{1, <:Any, <:Any, <:Any, <:Any, <:IndexMap}
+
+Adapt.adapt_structure(to, range::NDRange{N, B, W}) where {N, B, W} =
+    NDRange{N, B, W}(Adapt.adapt(to, range.blocks), Adapt.adapt(to, range.workitems), Adapt.adapt(to, range.mapping))
+
 import Base.iterate
 @inline iterate(range::NDRange) = iterate(blocks(range))
 @inline iterate(range::NDRange, state) = iterate(blocks(range), state)
 
 Base.length(range::NDRange) = length(blocks(range))
 
+"""
+    expand(ndrange::NDRange, groupidx, idx)
+
+Index of the `ndrange` handled by work item `idx` of workgroup `groupidx`, both given as
+`CartesianIndex` or linear positions in the blocked iteration space. For a
+[`MappedNDRange`](@ref) this is the entry of the index map for that work item, or the
+[`invalid_index`](@ref) for a work item past its end.
+"""
 @inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where {N}
     offset = offsets(ndrange)
     nI = ntuple(Val(N)) do I
@@ -251,6 +333,22 @@ Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::Integer,
     return expand(ndrange, blocks(ndrange)[groupidx], idx)
 end
 
+"""
+    linear_index(ndrange::MappedNDRange, groupidx, idx)
+
+Position in the index map of work item `idx` of workgroup `groupidx`.
+"""
+@inline linear_index(ndrange::MappedNDRange, groupidx::Integer, idx::Integer) = (groupidx - 1) * length(workitems(ndrange)) + idx
+@inline linear_index(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1}) = linear_index(ndrange, groupidx.I[1], idx.I[1])
+@inline linear_index(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::Integer) = linear_index(ndrange, groupidx.I[1], idx)
+@inline linear_index(ndrange::MappedNDRange, groupidx::Integer, idx::CartesianIndex{1}) = linear_index(ndrange, groupidx, idx.I[1])
+
+# The listed index of a work item, or `invalid_index` for a work item past the end of the map.
+@inline expand(ndrange::MappedNDRange, groupidx::Integer, idx::Integer) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx))
+@inline expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1}) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx))
+@inline expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::Integer) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx))
+@inline expand(ndrange::MappedNDRange, groupidx::Integer, idx::CartesianIndex{1}) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx))
+
 """
     partition(ndrange, workgroupsize)
 

src/pocl/backend.jl

@@ -319,15 +319,6 @@ end
 
 @device_override KI.get_sub_group_local_id() = get_sub_group_local_id() % UInt32
 
-@device_override @inline function KA.__validindex(ctx)
-    if KA.__dynamic_checkbounds(ctx)
-        I = @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(1), get_local_id(1))
-        return I in KA.__ndrange(ctx)
-    else
-        return true
-    end
-end
-
 
 ## Shared and Scratch Memory