Launch @kernel kernels on N-d grids, indexing in 32 bits

docs/src/api.md

@@ -108,4 +108,8 @@ KernelAbstractions.NDIteration.DynamicOffset
 KernelAbstractions.NDIteration.extents
 KernelAbstractions.NDIteration.offsets
 KernelAbstractions.NDIteration.linear_index
+KernelAbstractions.LinearLaunch
+KernelAbstractions.NDLaunch
+KernelAbstractions.select_launch
+KernelAbstractions.launch_workgroupsize
 ```

docs/src/implementations.md

@@ -75,3 +75,73 @@ to the backend's array type, so that `adapt(backend, x)` and
 Adapt.adapt_storage(::CUDABackend, x) = adapt(CuArray, x)
 ```
 
+
+## Launching `@kernel` kernels
+
+A kernel written with [`@kernel`](@ref) receives a hidden context, a
+`KernelAbstractions.CompilerMetadata` built by the backend's `mkcontext`, from which
+[`@index`](@ref) computes its indices. By default (a context without a `launch`),
+`@index` assumes that the kernel was launched on a 1-D grid of
+`length(blocks(iterspace))` groups of `length(workitems(iterspace))` work-items. It then
+decomposes the linear hardware ids into Cartesian positions, which takes integer divisions
+when the `ndrange` is not known at compile time, and computes in `Int`.
+
+A backend **may** launch kernels differently, and pass the `launch` keyword to the
+`CompilerMetadata` constructor to tell `@index` how:
+
+- [`NDLaunch{T}`](@ref KernelAbstractions.NDLaunch): the grid has the shape of the
+  iteration space (for as many dimensions as the backend's grid has, i.e. up to 3), so
+  `@index` doesn't need any divisions;
+- [`LinearLaunch{T}`](@ref KernelAbstractions.LinearLaunch): the default 1-D grid.
+
+Either way `@index` computes in `T`, e.g. `Int32`, which is faster on GPUs. The backend
+has to implement the typed [`KI.get_group_id`](@ref KernelInterface.get_group_id) and
+[`KI.get_local_id`](@ref KernelInterface.get_local_id) queries such that they compute in
+`T` too, e.g. without checked conversions.
+
+[`select_launch`](@ref KernelAbstractions.select_launch) chooses the launch from the
+iteration space, whether the workgroup size will be tuned, and the limits of the backend
+([`KI.max_work_group_size`](@ref KernelInterface.max_work_group_size),
+[`KI.max_work_group_dims`](@ref KernelInterface.max_work_group_dims) and
+[`KI.max_num_groups`](@ref KernelInterface.max_num_groups)). It doesn't depend on the
+workgroup size a backend tunes afterwards, which keeps the context type (and thus the
+compiled kernel) the same before and after tuning, as long as the backend tunes with
+[`launch_workgroupsize`](@ref KernelAbstractions.launch_workgroupsize). A launch then
+looks like this:
+
+```julia
+function (obj::KA.Kernel{MyBackend})(args...; ndrange = nothing, workgroupsize = nothing)
+    ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize)
+    launch = KA.select_launch(obj, workgroupsize, iterspace)
+    ctx = KA.CompilerMetadata{KA.ndrange(obj), KA.DynamicCheck}(ndrange, iterspace; launch)
+    kernel = compile(obj.f, ctx, args...)
+
+    if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing
+        threads = max_threads(kernel)  # at most `KI.max_work_group_size(backend)`
+        workgroupsize = KA.launch_workgroupsize(backend, launch, threads, ndrange)
+        iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize)
+        ctx = KA.CompilerMetadata{KA.ndrange(obj), KA.DynamicCheck}(ndrange, iterspace; launch)
+    end
+
+    groups, items = size(KA.blocks(iterspace)), size(KA.workitems(iterspace))
+    prod(groups) == 0 && return
+    if launch isa KA.NDLaunch
+        run(kernel, ctx, args...; groups, items)            # padded to 3 dimensions
+    else
+        run(kernel, ctx, args...; groups = prod(groups), items = prod(items))
+    end
+end
+```
+
+The POCL backend is an example. Backends that launch on an N-d grid **must not** override
+`__validindex` or the `__index_*` functions, which dispatch on the launch.
+
+Packages that customize the iteration space (with a custom `partition` and `expand`)
+don't need to do anything for these launches: the index functions only compute the global
+index directly for the iteration spaces KernelAbstractions creates itself, and call
+`expand`, `in` and `linear_index` otherwise.
+
+Packages with an `Adapt` rule for `CompilerMetadata` **must** preserve its `launch`, e.g. by
+passing `launch = KernelAbstractions.__launch(ctx)` to the constructor. Otherwise the
+kernel computes its indices as if it had been launched on a 1-D grid, which gives wrong
+results for a kernel launched with an `NDLaunch`.

src/KernelAbstractions.jl

@@ -389,29 +389,16 @@ end
 # Internal kernel functions
 ###
 
-@inline function __index_Local_Linear(ctx)
-    return KI.get_local_id().x
-end
-
-@inline function __index_Group_Linear(ctx)
-    return KI.get_group_id().x
-end
+# The index functions dispatch on the launch configuration of the context (see
+# `launch.jl`). A context without one was launched on a 1-D grid, and is indexed in `Int`.
+@inline index_launch(ctx) = something(__launch(ctx), LinearLaunch{Int}())
 
-@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)
-end
-
-@inline function __index_Local_Cartesian(ctx)
-    return @inbounds workitems(__iterspace(ctx))[KI.get_local_id().x]
-end
-@inline function __index_Group_Cartesian(ctx)
-    return @inbounds blocks(__iterspace(ctx))[KI.get_group_id().x]
-end
-@inline function __index_Global_Cartesian(ctx)
-    return @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x)
-end
+@inline __index_Local_Linear(ctx) = local_linear(ctx, index_launch(ctx))
+@inline __index_Group_Linear(ctx) = group_linear(ctx, index_launch(ctx))
+@inline __index_Global_Linear(ctx) = global_linear(ctx, index_launch(ctx))
+@inline __index_Local_Cartesian(ctx) = local_cartesian(ctx, index_launch(ctx))
+@inline __index_Group_Cartesian(ctx) = group_cartesian(ctx, index_launch(ctx))
+@inline __index_Global_Cartesian(ctx) = global_cartesian(ctx, index_launch(ctx))
 
 @inline __index_Local_NTuple(ctx, I...) = Tuple(__index_Local_Cartesian(ctx, I...))
 @inline __index_Group_NTuple(ctx, I...) = Tuple(__index_Group_Cartesian(ctx, I...))
@@ -606,13 +593,23 @@ end
 ###
 
 include("compiler.jl")
+include("launch.jl")
 
 ###
 # Compiler/Frontend
 ###
 
 function __workitems_iterspace end
-function __validindex end
+
+# Whether the current work-item is part of the ndrange, or a padding lane of a partial
+# workgroup. Padding lanes still take part in `@synchronize`.
+@inline function __validindex(ctx)
+    if __dynamic_checkbounds(ctx)
+        return validindex(ctx, index_launch(ctx))
+    else
+        return true
+    end
+end
 
 # for reflection
 function mkcontext end

src/compiler.jl

@@ -1,18 +1,26 @@
-struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
+"""
+    CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch}
+
+The hidden context argument of kernels written with [`@kernel`](@ref). The `launch` field
+tells the index functions how the backend launched the kernel: `nothing` for a 1-D launch
+indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref).
+"""
+struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch}
     groupindex::I
     ndrange::NDRange
     iterspace::Iterspace
+    launch::Launch
 
     # CPU variant
     function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB}
         ndrange = cartesian(ndrange)
-        return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace)
+        return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing}(idx, ndrange, iterspace, nothing)
     end
 
     # GPU variante: index is given implicit
-    function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB}
+    function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing) where {NDRange, CB}
         ndrange = cartesian(ndrange)
-        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace)
+        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch)}(nothing, ndrange, iterspace, launch)
     end
 end
 
@@ -25,6 +33,7 @@ cartesian(t::Tuple) = CartesianIndices(t)
 
 @inline __iterspace(cm::CompilerMetadata) = cm.iterspace
 @inline __groupindex(cm::CompilerMetadata) = cm.groupindex
+@inline __launch(cm::CompilerMetadata) = cm.launch
 @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm)))
 @inline __dynamic_checkbounds(::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = CB <: DynamicCheck
 @inline __ndrange(::CompilerMetadata{NDRange}) where {NDRange <: StaticSize} = CartesianIndices(get(NDRange))
@@ -39,7 +48,7 @@ cartesian(t::Tuple) = CartesianIndices(t)
 function Adapt.adapt_structure(to, cm::CompilerMetadata{NDRange, CB, I}) where {NDRange, CB, I}
     iterspace = Adapt.adapt(to, cm.iterspace)
     if I === Nothing
-        return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace)
+        return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch)
     else
         return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace)
     end