Add @groupreduce and @groupscan

docs/src/api.md

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

docs/src/index.md

@@ -208,6 +208,9 @@ end
   written against KernelInterface and for kernel languages built on it; capability queries
   say how a backend forms its sub-groups, whether they may diverge, and which types it can
   shuffle. See [Sub-groups](@ref ki_subgroups).
+- [`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan a value over the
+  workgroup, on every backend. `@groupreduce` uses sub-groups internally where the backend
+  can shuffle the values. See [Reductions and scans](@ref group_collectives).
 
 ## Semantic differences
 

docs/src/kernelinterface.md

@@ -213,7 +213,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,46 @@ 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](@id group_collectives)
+
+[`@groupreduce`](@ref) combines a value from every work-item of the workgroup and returns the
+result to all of them. [`@groupscan`](@ref) gives each work-item the result up to itself, in
+the order of `@index(Local, Linear)`; it is inclusive by default, and exclusive with
+`inclusive = false`. Both take an associative operator and its neutral element, and
+`@groupreduce` also needs the operator to be commutative.
+
+```julia
+@kernel function group_sums!(out, @Const(x))
+    i = @index(Global, Linear)
+    s = @groupreduce(+, x[i], zero(eltype(out)); groupsize = 256)
+    if @index(Local, Linear) == 1
+        out[@index(Group, Linear)] = s
+    end
+end
+```
+
+Like [`@synchronize`](@ref), they are collectives: every work-item of the workgroup has to
+reach them, not in a branch or loop that only some of them take, or after an early `return`.
+Use them as a statement of their own, optionally assigning the result, as in
+`res = @groupreduce(+, val, zero(T))`. Work-items that pad a partial workgroup take part as
+well: they contribute the neutral element, without evaluating `val`.
+
+Both go through local memory, whose size has to be known when the kernel is compiled: a
+kernel without a static workgroup size passes `groupsize`, a constant upper bound on the
+number of work-items per workgroup.
+
+They work on every backend. Where the backend can shuffle the values, `@groupreduce` first
+reduces each sub-group, and then combines those results through local memory.
+KernelAbstractions picks this automatically. The order in which values are combined can
+differ, so floating-point results may vary between backends and launch configurations.
+
+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
 
@@ -680,6 +682,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)
 
@@ -104,10 +116,18 @@ 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)
     args, tt = prepare_arguments(b, obj.f, ctx, args)
     kernel = KI.kernel_function(b, obj.f, tt; compiler_options(obj)...)
 
@@ -116,7 +136,7 @@ function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace,
         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