Document the contract of custom NDRange mappings

docs/src/api.md

@@ -87,11 +87,13 @@ They are only documented here for developers and contributors of `KernelAbstract
 
 ```@docs
 KernelAbstractions.partition
+KernelAbstractions.cartesian
 KernelAbstractions.@context
 KernelAbstractions.argconvert
 KernelAbstractions.NDIteration.DynamicSize
 KernelAbstractions.NDIteration.StaticSize
 KernelAbstractions.NDIteration.NDRange
+KernelAbstractions.NDIteration.expand
 KernelAbstractions.NDIteration.StaticOffset
 KernelAbstractions.NDIteration.DynamicOffset
 KernelAbstractions.NDIteration.extents

docs/src/implementations.md

@@ -26,3 +26,25 @@ to the backend's array type, so that `adapt(backend, x)` and
 Adapt.adapt_storage(::CUDABackend, x) = adapt(CuArray, x)
 ```
 
+## Iteration spaces and the validity of work items
+
+The context a backend passes to a kernel carries the blocked iteration space,
+[`__iterspace(ctx)`](@ref KernelAbstractions.NDIteration.NDRange), and the
+`ndrange` of the launch, `__ndrange(ctx)`. A backend must derive everything
+about a work item from these two objects through three functions:
+
+- [`expand(iterspace, groupidx, idx)`](@ref KernelAbstractions.NDIteration.expand)
+  gives the `CartesianIndex` handled by work item `idx` of workgroup `groupidx`;
+- `expand(iterspace, groupidx, idx) in __ndrange(ctx)` tells whether that work
+  item has an index to handle, which is how `__validindex` must be implemented;
+- [`linear_index(__ndrange(ctx), I)`](@ref KernelAbstractions.NDIteration.linear_index)
+  gives the linear index of `I`.
+
+A backend must not assume that `__ndrange(ctx)` is a `CartesianIndices` or that
+`expand` is an affine map: the `mapping` field of the `NDRange` lets a package
+define its own iteration space, for example a list of indices to visit, by
+extending these functions for its mapping type. Overriding `__validindex` or
+`__index_Global_Linear` for a generic `ctx` would bypass such an extension.
+See [`NDRange`](@ref KernelAbstractions.NDIteration.NDRange) for the functions
+a custom mapping has to define.
+

src/KernelAbstractions.jl

@@ -526,6 +526,9 @@ Partition the iteration space of `kernel` into workgroups.
 
 Returns the blocked iteration space and whether dynamic bounds-checking is required for the
 last (possibly partial) workgroup. Primarily used by backend implementations and tests.
+
+Specialize on the type of `ndrange` to launch kernels over an iteration space of your own,
+see [`NDRange`](@ref KernelAbstractions.NDIteration.NDRange).
 """
 @inline function partition(kernel, ndrange, workgroupsize)
     static_ndrange = KernelAbstractions.ndrange(kernel)

src/compiler.jl

@@ -16,7 +16,14 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
     end
 end
 
-# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`).
+"""
+    cartesian(ndrange)
+
+The object stored as `ndrange` of a kernel context for a launch `ndrange` in any form accepted
+by [`partition`](@ref): `CartesianIndices` covering the range, or `nothing` for a static one.
+Specialize it together with `partition` for an iteration space of your own, returning an object
+that supports `Base.in` for a `CartesianIndex` and [`linear_index`](@ref KernelAbstractions.NDIteration.linear_index).
+"""
 cartesian(::Nothing) = nothing
 cartesian(ci::CartesianIndices) = ci
 cartesian(n::Integer) = CartesianIndices((Int(n),))

src/nditeration.jl

@@ -140,6 +140,24 @@ Encodes a blocked iteration space. The `mapping` field relates blocked indices t
 `ndrange` indices: `nothing` for the identity, or a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref)
 for an `ndrange` whose indices do not start at 1.
 
+# Custom mappings
+
+A package can iterate over a space of its own, for example a list of indices, by launching a
+kernel with an `ndrange` object of its own type and defining:
+
+- [`partition(kernel, ndrange, workgroupsize)`](@ref KernelAbstractions.partition) for that
+  type, returning an `NDRange` whose `mapping` describes the space, and whether the last
+  workgroup needs bounds-checking;
+- [`cartesian(ndrange)`](@ref KernelAbstractions.cartesian) for that type, returning the
+  object stored as `ndrange` of the kernel context, which supports `Base.in` for a
+  `CartesianIndex` and [`linear_index`](@ref);
+- [`expand`](@ref) for an `NDRange` with that mapping and `groupidx`, `idx` given as
+  `Integer` or `CartesianIndex`, returning the index handled by a work item, or an index that
+  is not `in` the `ndrange` object for a work item without one.
+
+Backends check the validity of a work item as `expand(iterspace, groupidx, idx) in ndrange`,
+so nothing else is needed for the kernel to see the mapped index through `@index`.
+
 # Example
 ```
 ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((256, 256)), CartesianIndices((32, 32)))
@@ -188,6 +206,13 @@ import Base.iterate
 
 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 a
+`CartesianIndex` or as a linear position in the blocked iteration space. The result follows
+the `mapping` of the `ndrange`.
+"""
 @inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where {N}
     offset = offsets(ndrange)
     nI = ntuple(Val(N)) do I