Combine the values of a block with sub-group reductions

src/block.jl

@@ -17,19 +17,54 @@
 end
 
 
-# The sub-group width `W` for a block scan of `T` in blocks of `block_size` threads, as a `Val`,
-# or `Val(0)` for the local-memory scan. Sub-groups must be enabled by the scan tuning, linear
-# (from KernelAbstractions 0.10), able to shuffle `T`, and `W ≤ block_size ≤ W²`, as one
-# sub-group scans the sub-groups' totals.
+# The sub-group width `W` for a block scan or reduction of `T` in blocks of `block_size` threads,
+# as a `Val`, or `Val(0)` for the local-memory one. Sub-groups must be enabled by the tuning,
+# linear (from KernelAbstractions 0.10), able to shuffle `T`, and `W ≤ block_size ≤ W²`, as one
+# sub-group combines the sub-groups' values.
 @static if isdefined(KernelAbstractions, :KernelInterface)
-    function scan_subgroup_width(backend, ::Type{T}, block_size) where {T}
-        tuning(scan_tuning, backend, T).subgroups || return Val(0)
+    function subgroup_width(backend, ::Type{T}, block_size) where {T}
         KI.supports_linear_subgroups(backend) && KI.supports_shuffle(backend, T) || return Val(0)
         W = KI.sub_group_size(backend)
         return ispow2(W) && W <= block_size <= W * W ? Val(W) : Val(0)
     end
 else
-    scan_subgroup_width(backend, T, block_size) = Val(0)
+    subgroup_width(backend, T, block_size) = Val(0)
+end
+scan_subgroup_width(backend, ::Type{T}, block_size) where {T} =
+    tuning(scan_tuning, backend, T).subgroups ? subgroup_width(backend, T, block_size) : Val(0)
+reduce_subgroup_width(backend, ::Type{T}, block_size) where {T} =
+    tuning(reduce_tuning, backend, T).subgroups ? subgroup_width(backend, T, block_size) : Val(0)
+
+# The local memory of a block reduction of `N` values (see `block_reduce`): a value per
+# sub-group, or per thread without sub-groups
+@inline block_slots(N, ::Val{W}) where {W} = W == 0 ? N : N ÷ W
+
+# Reduce the values `x` of the block's threads with `op` into the result of thread 0, with
+# `sdata` of `block_slots(@groupsize()[1], subgroups)` elements. All threads in the block must
+# call it. With a sub-group width `W` (see `reduce_subgroup_width`), each sub-group reduces its
+# values, and then the sub-groups' values (as CUB's `BLOCK_REDUCE_WARP_REDUCTIONS`), in the order
+# of the threads; else a tree in local memory, which needs `op` to be commutative. It does
+# not synchronize after reading `sdata`, so a block that reduces again must first.
+@inline function block_reduce(@context, op, sdata, x, ithread, ::Val{0})
+    sdata[ithread + 0x1] = x
+    @synchronize()
+    @inline reduce_group!(@context, op, sdata, ithread)
+    return sdata[0x1]
+end
+@static if isdefined(KernelAbstractions, :KernelInterface)
+@inline function block_reduce(@context, op, sdata, x, ithread, ::Val{W}) where {W}
+    N = @groupsize()[1]
+    x = KI.sub_group_reduce(op, x)
+    N == W && return x
+    nsg = N ÷ W
+    if ithread % W == 0x0
+        sdata[ithread ÷ W + 0x1] = x
+    end
+    @synchronize()
+    # (every sub-group combines the sub-groups' values, as some backends (PoCL) need the
+    # sub-groups of a block to run the same sub-group operations; thread 0's is the result)
+    return KI.sub_group_reduce(op, sdata[ithread % nsg + 0x1], nsg)
+end
 end
 
 

src/reduce/kernels_nd.jl

@@ -84,17 +84,19 @@ end
 
 # `ByLanes`: each output reduces `reduce_size` consecutive elements (from the offset
 # of its outer index), with `lanes` threads that read them in an interleaved order; a block of
-# `N` threads covers `N ÷ lanes` outputs, which are combined in local memory with log2(lanes)
-# steps.
+# `N` threads covers `N ÷ lanes` outputs, whose lanes are combined with sub-group operations
+# where they fit a sub-group of width `W` (see `reduce_subgroup_width`), else in local memory with
+# log2(lanes) steps.
 @kernel inbounds=true cpu=false unsafe_indices=true function by_lanes_kernel!(
     @Const(src), dst,
     f::F, op::OP, seed, init,
     base_offset, outer_strides, outer_sizes,
-    output_size, reduce_size, ::Val{lanes},
-) where {F, OP, lanes}
+    output_size, reduce_size, ::Val{lanes}, ::Val{W},
+) where {F, OP, lanes, W}
     seed = kernel_value(seed)
     @uniform N = @groupsize()[1]
-    sdata = @localmem typeof(seed) (N,)
+    # (no local memory where the lanes of an output are a segment of a sub-group)
+    sdata = @localmem typeof(seed) (lanes <= W ? 1 : N,)
     mapf, accop = seeded_ops(f, op, seed)
 
     iblock  = @index(Group, Linear) - 0x1
@@ -115,20 +117,25 @@ end
         end
     end
 
-    sdata[ithread + 0x1] = acc
-    @synchronize()
+    total = if lanes <= W
+        KI.sub_group_reduce(accop, acc, lanes)
+    else
+        sdata[ithread + 0x1] = acc
+        @synchronize()
 
-    step = lanes ÷ 2
-    while step >= 1
-        if lane < step
-            sdata[ithread + 0x1] = accop(sdata[ithread + 0x1], sdata[ithread + step + 0x1])
+        step = lanes ÷ 2
+        while step >= 1
+            if lane < step
+                sdata[ithread + 0x1] = accop(sdata[ithread + 0x1], sdata[ithread + step + 0x1])
+            end
+            @synchronize()
+            step ÷= 2
         end
-        @synchronize()
-        step ÷= 2
+        sdata[ithread + 0x1]
     end
 
     if lane == 0x0 && iout < output_size
-        dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[ithread + 0x1])
+        dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total)
     end
 end
 
@@ -277,11 +284,11 @@ end
     base_offset,
     outer_strides, outer_sizes,
     reduce_strides, reduce_sizes,
-    output_size, reduce_size, num_blocks,
+    output_size, reduce_size, num_blocks, subgroups::Val,
 ) where {F, OP}
     seed = kernel_value(seed)
     @uniform N = @groupsize()[1]
-    sdata = @localmem typeof(seed) (N,)
+    sdata = @localmem typeof(seed) (block_slots(N, subgroups),)
     mapf, accop = seeded_ops(f, op, seed)
 
     iblock  = @index(Group, Linear) - 0x1
@@ -298,13 +305,10 @@ end
             j  += N
         end
 
-        sdata[ithread + 0x1] = acc
-        @synchronize()
-
-        @inline reduce_group!(@context, accop, sdata, ithread)
+        total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups)
 
         if ithread == 0x0
-            dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[0x1])
+            dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total)
         end
 
         next_iout = iout + num_blocks
@@ -323,11 +327,11 @@ end
     base_offset,
     outer_strides, outer_sizes,
     output_size, reduce_size, num_blocks,
-    ::Val{W},
+    ::Val{W}, subgroups::Val,
 ) where {F, OP, W}
     seed = kernel_value(seed)
     @uniform N = @groupsize()[1]
-    sdata = @localmem typeof(seed) (N,)
+    sdata = @localmem typeof(seed) (block_slots(N, subgroups),)
     mapf, accop = seeded_ops(f, op, seed)
 
     iblock  = @index(Group, Linear) - 0x1
@@ -350,13 +354,10 @@ end
             jw += UInt32(1)
         end
 
-        sdata[ithread + 0x1] = acc
-        @synchronize()
-
-        @inline reduce_group!(@context, accop, sdata, ithread)
+        total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups)
 
         if ithread == 0x0
-            dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, sdata[0x1])
+            dst[iout + 0x1] = finish(op, init, dst, iout + 0x1, total)
         end
 
         next_iout = iout + num_blocks
@@ -377,11 +378,11 @@ end
     base_offset,
     outer_strides, outer_sizes,
     reduce_strides, reduce_sizes,
-    ::Val{output_size}, reduce_size, groups,
+    ::Val{output_size}, reduce_size, groups, subgroups::Val,
 ) where {F, OP, output_size}
     seed = kernel_value(seed)
     @uniform N = @groupsize()[1]
-    sdata = @localmem typeof(seed) (N,)
+    sdata = @localmem typeof(seed) (block_slots(N, subgroups),)
     mapf, accop = seeded_ops(f, op, seed)
 
     iblock  = @index(Group, Linear) - 0x1
@@ -401,12 +402,9 @@ end
         j  += N * groups
     end
 
-    sdata[ithread + 0x1] = acc
-    @synchronize()
-
-    @inline reduce_group!(@context, accop, sdata, ithread)
+    total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups)
 
     if ithread == 0x0
-        partial[iout + igroup * output_size + 0x1] = sdata[0x1]
+        partial[iout + igroup * output_size + 0x1] = total
     end
 end

src/reduce/mapreduce_1d_gpu.jl

@@ -10,11 +10,11 @@
 # `iblock` reduces the tiles `iblock`, `iblock + nblocks`, ... of `src`, and stores
 # `finish(op, init, dst, iblock + 1, partial)` in `dst[iblock + 1]`.
 @kernel inbounds=true cpu=false unsafe_indices=true function mapreduce_1d_kernel!(
-    @Const(src), dst, f::F, op::OP, seed, init, nblocks, ::Val{NI}, ::Val{K},
+    @Const(src), dst, f::F, op::OP, seed, init, nblocks, ::Val{NI}, ::Val{K}, subgroups::Val,
 ) where {F, OP, NI, K}
     seed = kernel_value(seed)
 
-    sdata = @localmem typeof(seed) (NI,)
+    sdata = @localmem typeof(seed) (block_slots(NI, subgroups),)
     mapf, accop = seeded_ops(f, op, seed)
     len = length(src)
 
@@ -42,14 +42,10 @@
             end
         end
     end
-    sdata[ithread + 0x1] = acc
-
-    @synchronize()
-
-    @inline reduce_group!(@context, accop, sdata, ithread)
+    total = @inline block_reduce(@context, accop, sdata, acc, ithread, subgroups)
 
     if ithread == 0x0
-        dst[iblock + 0x1] = finish(op, init, dst, iblock + 0x1, sdata[0x1])
+        dst[iblock + 0x1] = finish(op, init, dst, iblock + 0x1, total)
     end
 end
 
@@ -126,7 +122,9 @@ function mapreduce_1d_pass!(commutative::Bool, backend, block_size, items_per_th
     if commutative
         kernel! = mapreduce_1d_kernel!(backend, block_size)
         Base.inferencebarrier(kernel!)(src, dst, f, op, seed, init, blocks, Val(block_size),
-                                       Val(items_per_thread); ndrange=(block_size * blocks,))
+                                       Val(items_per_thread),
+                                       reduce_subgroup_width(backend, typeof(seed), block_size);
+                                       ndrange=(block_size * blocks,))
     else
         len = length(src)
         src, base_offset = something(buffer, (src, 0))

src/reduce/mapreduce_nd.jl

@@ -135,6 +135,7 @@ function launch!(shape::ByLanes, f::F, op::OP, dst, src, buffer, backend, alg, b
     by_lanes_kernel!(backend, block_size)(
         buffer, dst, f, op, seed, init,
         base_offset, outer_strides, outer_sizes, dst_size, reduce_size, Val(shape.lanes),
+        reduce_subgroup_width(backend, typeof(seed), block_size),
         ndrange=(block_size * blocks,),
     )
 end
@@ -152,13 +153,15 @@ function launch!(shape::ByBlock, f::F, op::OP, dst, src, buffer, backend, alg, b
         Base.inferencebarrier(kernel!)(
             buffer, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes,
             dst_size, reduce_size, blocks, Val(W),
+            reduce_subgroup_width(backend, typeof(seed), block_size),
             ndrange=(block_size * blocks,),
         )
     else
         kernel! = by_block_kernel!(backend, block_size)
         Base.inferencebarrier(kernel!)(
             buffer, dst, f, op, seed, init, base_offset, outer_strides, outer_sizes,
             reduce_strides, reduce_sizes, dst_size, reduce_size, blocks,
+            reduce_subgroup_width(backend, typeof(seed), block_size),
             ndrange=(block_size * blocks,),
         )
     end
@@ -174,6 +177,7 @@ function launch!(shape::ByBlocks, f::F, op::OP, dst, src, buffer, backend, alg,
         buffer, partials, f, op, seed,
         base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes,
         Val(dst_size), reduce_size, groups,
+        reduce_subgroup_width(backend, typeof(seed), block_size),
         ndrange=(block_size * dst_size * groups,),
     )
     combine_groups!(backend, partials, dst, op, seed, init, dst_size, groups, false,
@@ -238,7 +242,9 @@ function combine_groups!(backend, partials, dst, op::OP, seed, init, dst_size, g
     else
         by_block_kernel!(backend, block_size)(
             partials, dst, PartialInputs(), op, seed, init, 0, (1,), (dst_size,), (dst_size,),
-            (groups,), dst_size, groups, dst_size, ndrange=(block_size * dst_size,))
+            (groups,), dst_size, groups, dst_size,
+            reduce_subgroup_width(backend, typeof(seed), block_size),
+            ndrange=(block_size * dst_size,))
     end
 end
 

src/reduce/mapreduce_ordered.jl

@@ -6,9 +6,9 @@
 # Block `iblock` reduces group `g` of output `iout` (`iblock = iout + g * output_size`): the
 # outputs' slices are split into tiles of `N * K` elements, and each of the `groups` groups takes
 # a contiguous range of them (an even share, as CUB's `GridEvenShare`). In each tile, every thread
-# reduces a run of `K` consecutive elements, and the block combines the runs in order (with the
-# scans' block scan, which uses sub-groups where `subgroups` says so) and appends the tile's total
-# to the group's. Indices and `f`, `seed` and `init` are those of the other reduction kernels (see
+# reduces a run of `K` consecutive elements, and the block combines the runs in order (with
+# sub-group reductions where `subgroups` says so, else the scans' block scan) and appends the
+# tile's total to the group's. Indices and `f`, `seed` and `init` are those of the other reduction kernels (see
 # `by_block_kernel!` and `mapreduce_1d_kernel!`); with several groups, `dst` holds
 # `output_size × groups` partial results, for a second pass of this kernel with
 # `f === PartialInputs()`. With a vector width `W`, the runs of full tiles are read with `K ÷ W`
@@ -24,7 +24,7 @@
 ) where {F, OP, K, W}
     seed = kernel_value(seed)
     @uniform N = @groupsize()[1]
-    totals = @localmem typeof(seed) (N,)
+    totals = @localmem typeof(seed) (block_slots(N, subgroups),)
     mapf, accop = seeded_ops(f, op, seed)
 
     iblock  = @index(Group, Linear) - 0x1
@@ -63,10 +63,8 @@
                 end
             end
         end
-        totals[ithread + 0x1] = acc
-        _, tile_total = block_exclusive_scan!(@context, accop, totals, seed, N, ithread,
-                                              subgroups)
-        total = accop(total, tile_total)
+        total = accop(total, ordered_block_reduce(@context, accop, totals, acc, seed, N, ithread,
+                                                  subgroups))
         tile += 1
     end
 
@@ -76,6 +74,21 @@
     end
 end
 
+# The reduction of the values `x` of the block's threads in their order, into the result of thread
+# 0: with sub-groups (`block_reduce`, which is in order), or with the local-memory scan
+@inline function ordered_block_reduce(@context, op, totals, x, seed, N, ithread, ::Val{0})
+    totals[ithread + 0x1] = x
+    _, total = block_exclusive_scan!(@context, op, totals, seed, N, ithread, Val(0))
+    return total
+end
+@inline function ordered_block_reduce(@context, op, totals, x, seed, N, ithread,
+                                      subgroups::Val)
+    total = @inline block_reduce(@context, op, totals, x, ithread, subgroups)
+    # (before the next tile's values overwrite `totals`)
+    @synchronize()
+    return total
+end
+
 # One pass of `ordered_kernel!`, of `groups` blocks per output; `dense_buffer` is whether
 # `src` is a dense buffer, which it can read with vector loads
 function ordered_pass!(
@@ -90,7 +103,7 @@ function ordered_pass!(
         src, dst, f, op, seed, init,
         base_offset, outer_strides, outer_sizes, reduce_strides, reduce_sizes,
         output_size, reduce_size, groups, Val(items_per_thread), Val(W),
-        scan_subgroup_width(backend, typeof(seed), block_size),
+        reduce_subgroup_width(backend, typeof(seed), block_size),
         ndrange=(block_size * output_size * groups,),
     )
     return nothing