Add @groupreduce and @groupscan

docs/src/api.md

@@ -15,6 +15,13 @@
 @ndrange
 ```
 
+### Reductions and scans
+
+```@docs
+@groupreduce
+@groupscan
+```
+
 ## Host language
 
 !!! note

docs/src/kernelinterface.md

@@ -206,7 +206,9 @@ used anywhere.
     kernel languages built on it. A kernel is either a KernelInterface kernel or a
     KernelAbstractions `@kernel`, never a mix: `@kernel` adds padding work-items to partial
     work-groups, which skip the kernel's body, so KernelInterface's sub-group functions called
-    in a `@kernel` aren't reached by all work-items. Host code may use the queries of both.
+    in a `@kernel` aren't reached by all work-items. A `@kernel` uses KernelAbstractions'
+    collectives instead, e.g. [`@groupreduce`](@ref KernelAbstractions.@groupreduce). Host
+    code may use the queries of both.
 
 ```@docs
 get_sub_group_size

docs/src/kernels.md

@@ -213,6 +213,22 @@ statements, and [`@uniform`](@ref) evaluates an expression outside the work-item
 can be reused across `@synchronize` statements. For scratch storage that does not need to
 survive across `@synchronize`, an `MArray` can be used instead.
 
+## Reductions and scans
+
+[`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan values over the workgroup.
+Like [`@synchronize`](@ref), they are collectives: all work-items of the workgroup have to
+reach them, not in a branch or loop that only some of them take, and they have to be used as
+statements of their own, e.g. `res = @groupreduce(+, val, zero(T))`. Work-items that pad a
+partial workgroup take part as well, contributing the neutral element. On backends with
+sub-groups, `@groupreduce` uses them where it can.
+
+A kernel is either a `@kernel` or a kernel written against
+[KernelInterface](@ref kernelinterface), never a mix: don't call KernelInterface's device
+functions, such as its sub-group shuffles and votes, in a `@kernel`. They have to be executed
+by all work-items of a sub-group, but in a `@kernel` with the default bounds checking, every
+statement other than the collectives above only runs on the work-items that are part of the
+`ndrange`. Host code can use KernelInterface's queries, e.g. to pick a kernel.
+
 ## Launching kernels
 
 Construct a kernel by calling the kernel function on a backend and optional static sizes, then

src/KernelAbstractions.jl

@@ -42,6 +42,8 @@ and then invoked on the arguments.
 - [`@uniform`](@ref)
 - [`@synchronize`](@ref)
 - [`@print`](@ref)
+- [`@groupreduce`](@ref)
+- [`@groupscan`](@ref)
 
 # Kernel constructor
 
@@ -670,6 +672,7 @@ function __workitems_iterspace end
 end
 
 include("macros.jl")
+include("groupreduction.jl")
 include("spawn.jl")
 include("foreach_index.jl")
 

src/backend_launch.jl

@@ -8,15 +8,27 @@
 ###
 
 """
-    mkcontext(kernel::Kernel, ndrange, iterspace, [launch])
+    mkcontext(kernel::Kernel, ndrange, iterspace, [launch]; subgroups = nothing)
 
 The hidden context argument for launching `kernel` over `ndrange`, partitioned as
-`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)).
+`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)) and the
+backend's sub-group capabilities `subgroups` (see [`kernel_subgroups`](@ref)).
 """
 mkcontext(kernel::Kernel, _ndrange, iterspace) =
     CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace)
-mkcontext(kernel::Kernel, _ndrange, iterspace, launch) =
-    CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch)
+mkcontext(kernel::Kernel, _ndrange, iterspace, launch; subgroups = nothing) =
+    CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch, subgroups)
+
+"""
+    kernel_subgroups(kernel::Kernel)
+
+The sub-group capabilities of the backend of `kernel` that its work-group collectives can
+use, a [`SubgroupCapabilities`](@ref), or `nothing` if the backend has none, or if the kernel
+doesn't use collectives (which doesn't depend on the backend, so that launching other
+kernels doesn't query it).
+"""
+kernel_subgroups(kernel::Kernel) =
+    uses_collectives(kernel.f) ? subgroup_capabilities(backend(kernel)) : nothing
 mkcontext(kernel::Kernel, I, _ndrange, iterspace, ::Dynamic) where {Dynamic} =
     CompilerMetadata{ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
 
@@ -103,18 +115,26 @@ function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize
 end
 
 function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace, args::Tuple)
+    # a constant `nothing` for kernels without collectives; for kernels with, the call below
+    # is a function barrier that makes the rest of the launch type stable
+    subgroups = kernel_subgroups(obj)
+    launch_kernel(obj, launch, subgroups, ndrange, _workgroupsize, iterspace, args)
+    return nothing
+end
+
+function launch_kernel(obj::Kernel, launch, subgroups, ndrange, _workgroupsize, iterspace, args::Tuple)
     b = backend(obj)
 
     # this might not be the final context, since we may tune the workgroupsize
-    ctx = mkcontext(obj, ndrange, iterspace, launch)
+    ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups)
     kernel = compile(obj, ctx, args)
 
     # tune the workgroup size, keeping the context type (and thus the kernel) the same
     if workgroupsize(obj) <: DynamicSize && _workgroupsize === nothing
         range = something(ndrange, static_ndrange(obj))
         threads = KI.launch_configuration(kernel; nitems = saturated_prod(extents(range))).workgroupsize
         iterspace, _ = partition(obj, ndrange, launch_workgroupsize(b, launch, threads, range))
-        ctx = mkcontext(obj, ndrange, iterspace, launch)
+        ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups)
     end
 
     # the geometry is valid by construction, except for the kernel's limits

src/compiler.jl

@@ -1,26 +1,29 @@
 """
-    CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch}
+    CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups}
 
 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).
+indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref). The `subgroups` field
+holds the backend's [`SubgroupCapabilities`](@ref) for kernels that use work-group
+collectives, or `nothing`.
 """
-struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch}
+struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups}
     groupindex::I
     ndrange::NDRange
     iterspace::Iterspace
     launch::Launch
+    subgroups::Subgroups
 
     # 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), Nothing}(idx, ndrange, iterspace, nothing)
+        return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing, Nothing}(idx, ndrange, iterspace, nothing, nothing)
     end
 
     # GPU variante: index is given implicit
-    function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing) where {NDRange, CB}
+    function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing, subgroups = nothing) where {NDRange, CB}
         ndrange = cartesian(ndrange)
-        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch)}(nothing, ndrange, iterspace, launch)
+        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch), typeof(subgroups)}(nothing, ndrange, iterspace, launch, subgroups)
     end
 end
 
@@ -34,6 +37,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 __subgroups(cm::CompilerMetadata) = cm.subgroups
 @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))
@@ -48,7 +52,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; cm.launch)
+        return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch, cm.subgroups)
     else
         return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace)
     end