diff --git a/docs/src/api.md b/docs/src/api.md index e7c6d6023..3313c5a88 100644 --- a/docs/src/api.md +++ b/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 diff --git a/docs/src/implementations.md b/docs/src/implementations.md index f75655648..011894d26 100644 --- a/docs/src/implementations.md +++ b/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. + diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index c022c7cef..3c56fd7c7 100644 --- a/src/KernelAbstractions.jl +++ b/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) diff --git a/src/compiler.jl b/src/compiler.jl index b7d388d62..926c9c7b3 100644 --- a/src/compiler.jl +++ b/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),)) diff --git a/src/nditeration.jl b/src/nditeration.jl index 1858f3857..8c9934b65 100644 --- a/src/nditeration.jl +++ b/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