From 0436c659e8b41bdde3d5f917ba658c26f0fc1f6e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 24 Oct 2026 12:54:04 +0200 Subject: [PATCH 1/2] Reduce strided slices with the tiled kernel where one thread per output underfills MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A strided reduction into contiguous outputs with at least as many outputs as elements in each took one thread per output (the tiled kernel only exactly square ones), whose `dst_size` threads leave a GPU mostly idle below tens of thousands of outputs: summing a 3072×2048 BFloat16 matrix along `dims=2` took 90 µs on an RTX 5080. These reductions now take the tiled kernel, 8 outputs per block in one pass, where tiles of 32 outputs would not fill the device, and the columns kernel where they would; one thread per output only for rows of at least 64 KiB of outputs, or reductions too short to share. That matrix takes 16.9 µs, and an 8192×8192 Float32 one 295 instead of 416 µs. --- src/reduce/mapreduce_nd.jl | 84 ++++++++++++++++++++++++-------------- test/generic/reduce.jl | 15 +++++++ 2 files changed, 68 insertions(+), 31 deletions(-) diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 59b8cf47..e8fed420 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -49,6 +49,15 @@ const COLUMNS_MIN_OUTPUTS = 8 # The most threads per block of the columns kernel: its unrolled loads take more registers than # the other kernels, and Metal allows it fewer than 1024 threads per block const COLUMNS_MAX_THREADS = 256 +# A strided reduction into contiguous outputs takes one thread per output when a row of outputs, +# which those threads load together at each step, is at least this many bytes, or when the +# reduction is shorter than COLUMNS_MIN_REDUCE. Measured on an RTX 5080 and an RTX 6000 Ada: one +# thread per output is as fast as the other kernels from 64 KiB of Float32 or BFloat16 outputs, +# and 2.3 (Float32) to 4.6 (BFloat16) times slower with 4096 outputs, which fill only 16 blocks. +# (Devices with a few compute units fill sooner: on a Ryzen 9950X's integrated GPU and a Jetson +# Xavier, 2048 outputs of 768 elements reduce 1.4 times as fast with one thread each.) +const BY_THREAD_ROW_BYTES = 64 * 1024 +const COLUMNS_MIN_REDUCE = 32 # Host-side canonicalization: split the dimensions into reduced and kept ("outer") segments, @@ -112,12 +121,9 @@ _segment_tuple(t::Tuple, n::Int) = ntuple(i -> t[i], n) # run of consecutive elements, of a length the # tuning's `segments_*` settings allow -> :segments (a group of threads per # output, several outputs per block) -# - at least COLUMNS_MIN_OUTPUTS contiguous outputs -# with a longer strided reduction, of at least -# the tuning's `columns_min_elements` in all -> :columns (blocks of up to 32 -# consecutive outputs; a second pass -# when an output's reduction is split) -# - square-like, contiguous output with strided input -> :tiled_strided +# - contiguous outputs, each reducing one strided +# segment -> :by_thread, :tiled_strided or :columns, +# see `_columns_shape` # - dst_size >= reduce_size -> :by_thread # - dst_size == 1, or dst_size < GS_DST_CUTOFF # (and dst_size < reduce_size) -> :multigroup (split one reduction @@ -138,6 +144,11 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceT (rstrides, rsizes, nr), (ostrides, _, no) = segs single_reduce = nr == 1 && rsizes[1] != 0 + if no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 + shape = _columns_shape(dst_size, reduce_size, block_size, t, A) + shape === nothing || return shape + end + # Each output reduces a run of consecutive elements: a group of lanes per output reads it # with consecutive loads, and a block holds several outputs, so that it stays busy on short # runs (where one block per output idles most of its threads, and one thread per output @@ -161,37 +172,16 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceT # by_block override 2: when by_thread would launch too few blocks for a square-like # output/reduction shape, split each output reduction across a full block. This - # fallback is mostly for layouts that do not satisfy the narrower tiled-strided - # pattern below; the common contiguous-output strided case uses tiled_strided. - # Avoid applying this to wide-output shapes, where by_thread's cross-output - # coalescing is better than a strided block reduction. + # fallback is for layouts `_columns_shape` does not take. Avoid applying this to + # wide-output shapes, where by_thread's cross-output coalescing is better than a + # strided block reduction. use_by_block_for_low_occupancy = dst_size == reduce_size && reduce_size >= block_size && cld(dst_size, block_size) < target_blocks && single_reduce use_by_block = use_by_block_for_coalescing || use_by_block_for_low_occupancy - use_tiled_strided = - dst_size == reduce_size && reduce_size >= block_size && - no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 && - block_size % TILED_STRIDED_ROWS_PER_BLOCK == 0 && - ispow2(block_size ÷ TILED_STRIDED_ROWS_PER_BLOCK) - - # Contiguous outputs with a longer strided reduction, e.g. the row sums of a wide matrix: - # :columns, with enough blocks per output to fill the device - if no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 && - COLUMNS_MIN_OUTPUTS <= dst_size < reduce_size && - widemul(dst_size, reduce_size) >= t.columns_min_elements && - block_size >= _columns_cols(dst_size) - cols = _columns_cols(dst_size) - lanes = min(block_size, COLUMNS_MAX_THREADS) ÷ cols - groups = clamp(cld(target_blocks, cld(dst_size, cols)), 1, - cld(reduce_size, lanes * MIN_ITEMS_PER_THREAD)) - return :columns, groups - end - if use_tiled_strided - return :tiled_strided, 0 - elseif dst_size >= reduce_size && !use_by_block + if dst_size >= reduce_size && !use_by_block return :by_thread, 0 elseif use_by_block return :by_block, 0 @@ -208,6 +198,38 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceT end end +# The launch shape of a strided reduction into contiguous outputs (e.g. `dims=2` of a matrix), or +# `nothing` for the shapes of `_mapreduce_nd_shape`: +# - one thread per output for rows of BY_THREAD_ROW_BYTES of outputs, or short reductions; +# - else, with at least as many outputs as elements in each, and too few for tiles of 32 outputs +# to launch the tuning's `target_blocks`, the tiled kernel: 8 outputs per block, in one pass; +# - else the columns kernel, with tiles of up to 32 outputs, and each output's reduction split +# across enough blocks to launch `target_blocks`. A split reduction takes a second pass, +# which only reductions of at least the tuning's `columns_min_elements` elements afford. +# Tiles of 32 outputs that need splitting take longer than the tiled kernel where its 8 outputs +# per block launch enough blocks, from 1.7 times as long on an RTX 5080 for a 1024×1024 matrix to +# 2.5 times on an Arc A770; with more outputs, the columns kernel's several loads in flight beat +# its single one (a 8192×8192 Float32 matrix takes 296 instead of 416 µs on the RTX 5080). +function _columns_shape(dst_size, reduce_size, block_size, t::ReduceTuning, + ::Type{A}) where {A} + if widemul(dst_size, _accsize(A)) >= BY_THREAD_ROW_BYTES || reduce_size < COLUMNS_MIN_REDUCE + return :by_thread, 0 + end + threads = min(block_size, COLUMNS_MAX_THREADS) + cols = _columns_cols(dst_size) + tiles = cld(dst_size, cols) + rows = TILED_STRIDED_ROWS_PER_BLOCK + if dst_size >= reduce_size && tiles < t.target_blocks && + block_size % rows == 0 && ispow2(block_size ÷ rows) && reduce_size >= block_size ÷ rows + return :tiled_strided, 0 + end + dst_size >= COLUMNS_MIN_OUTPUTS && threads >= cols || return nothing + lanes = threads ÷ cols + groups = clamp(cld(t.target_blocks, tiles), 1, cld(reduce_size, lanes * MIN_ITEMS_PER_THREAD)) + groups > 1 && widemul(dst_size, reduce_size) < t.columns_min_elements && return nothing + return :columns, groups +end + # The launch shape of a reduction in element order (`commutative=false`), from the shape # `_mapreduce_nd_shape` chose for it: one thread per output, which reduces its slice in order, # where that shape is one, or for runs shorter than a block; the columns kernel in order for diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index 92df385e..c76ba98e 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1192,6 +1192,21 @@ end v = view(p, 2:m + 1, 4:n + 3) @test Array(AK.sum(v; dims=2, alg)) == sum(Int.(Array(p)[2:m + 1, 4:n + 3]); dims=2) end + # (the shape the reduction of `x` along `dims` with accumulator type `A` takes) + shape_of(x, dims, A) = AK._mapreduce_nd_shape( + length(x) ÷ prod(size(x, d) for d in dims), prod(size(x, d) for d in dims), + AK._mapreduce_segments(size(x), AK._mapreduce_layout(x)[3], dims), + AK._resolve_reduce(alg, BACKEND, A, dims).block_size, AK.reduce_tuning(BACKEND, A), A)[1] + # More outputs than elements each reduces: the tiled kernel, unless tiles of 32 outputs + # launch enough blocks, and one thread per output where a row of outputs fills the device + for (m, n, A, shape) in ((4096, 1024, Int32, :tiled_strided), (8192, 300, Int32, :columns), + (8192, 300, Int64, :by_thread)) + h = rand(A(-9):A(9), m, n) + a = array_from_host(h) + TEST_KERNELS && @test shape_of(a, (2,), A) === shape + r = AK.reduce(+, a; dims=2, init=zero(A), alg) + @test eltype(r) === A && Array(r) == sum(h; dims=2) + end end From cb9b168a5d4c34f06cca4e5e96943b787c21e1a6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 24 Oct 2026 12:57:57 +0200 Subject: [PATCH 2/2] Reduce the middle dimensions of arrays with the columns kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The columns kernel took only contiguous outputs, so reducing the middle dimension of a 3-d array (`sum(x; dims=2)` of a `(features, sequence, batch)` array) fell back to one thread or one block per output: 42 µs for a 768×1024×8 BFloat16 array on an RTX 5080, a third of the bandwidth, and 275 µs on an RTX 6000 Ada. It now reduces runs of contiguous outputs at the offsets of the other kept dimensions, its tiles covering every run: that array takes 19.6 µs on the RTX 5080 and 21.5 µs on the RTX 6000 Ada, and summing a 64×4096×32 Float32 array along its second dimension 41 instead of 138 µs on the RTX 5080. --- src/reduce/mapreduce_nd.jl | 82 +++++++++++++++++++++++--------------- test/generic/reduce.jl | 12 +++++- 2 files changed, 60 insertions(+), 34 deletions(-) diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index e8fed420..569206d9 100644 --- a/src/reduce/mapreduce_nd.jl +++ b/src/reduce/mapreduce_nd.jl @@ -121,8 +121,8 @@ _segment_tuple(t::Tuple, n::Int) = ntuple(i -> t[i], n) # run of consecutive elements, of a length the # tuning's `segments_*` settings allow -> :segments (a group of threads per # output, several outputs per block) -# - contiguous outputs, each reducing one strided -# segment -> :by_thread, :tiled_strided or :columns, +# - runs of contiguous outputs, each reducing one +# strided segment -> :by_thread, :tiled_strided or :columns, # see `_columns_shape` # - dst_size >= reduce_size -> :by_thread # - dst_size == 1, or dst_size < GS_DST_CUTOFF @@ -141,11 +141,11 @@ _segment_tuple(t::Tuple, n::Int) = ntuple(i -> t[i], n) function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceTuning, ::Type{A}) where {A} target_blocks = t.target_blocks - (rstrides, rsizes, nr), (ostrides, _, no) = segs + (rstrides, rsizes, nr), (ostrides, osizes, no) = segs single_reduce = nr == 1 && rsizes[1] != 0 - if no == 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 - shape = _columns_shape(dst_size, reduce_size, block_size, t, A) + if no >= 1 && ostrides[1] == 1 && nr == 1 && rstrides[1] > 1 + shape = _columns_shape(dst_size, reduce_size, osizes[1], block_size, t, A) shape === nothing || return shape end @@ -198,32 +198,36 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceT end end -# The launch shape of a strided reduction into contiguous outputs (e.g. `dims=2` of a matrix), or -# `nothing` for the shapes of `_mapreduce_nd_shape`: +# The launch shape of a strided reduction into runs of `run` contiguous outputs (e.g. `dims=2` of a +# matrix, or of a 3-d array, whose third dimension holds the runs), or `nothing` for the shapes of +# `_mapreduce_nd_shape`: # - one thread per output for rows of BY_THREAD_ROW_BYTES of outputs, or short reductions; # - else, with at least as many outputs as elements in each, and too few for tiles of 32 outputs -# to launch the tuning's `target_blocks`, the tiled kernel: 8 outputs per block, in one pass; -# - else the columns kernel, with tiles of up to 32 outputs, and each output's reduction split -# across enough blocks to launch `target_blocks`. A split reduction takes a second pass, -# which only reductions of at least the tuning's `columns_min_elements` elements afford. +# to launch the tuning's `target_blocks`, the tiled kernel (of a single run): 8 outputs per +# block, in one pass; +# - else the columns kernel, with tiles of up to 32 outputs of a run, and each output's +# reduction split across enough blocks to launch `target_blocks`. A split reduction takes a +# second pass, which only reductions of at least the tuning's `columns_min_elements` elements +# afford. # Tiles of 32 outputs that need splitting take longer than the tiled kernel where its 8 outputs # per block launch enough blocks, from 1.7 times as long on an RTX 5080 for a 1024×1024 matrix to # 2.5 times on an Arc A770; with more outputs, the columns kernel's several loads in flight beat # its single one (a 8192×8192 Float32 matrix takes 296 instead of 416 µs on the RTX 5080). -function _columns_shape(dst_size, reduce_size, block_size, t::ReduceTuning, +function _columns_shape(dst_size, reduce_size, run, block_size, t::ReduceTuning, ::Type{A}) where {A} if widemul(dst_size, _accsize(A)) >= BY_THREAD_ROW_BYTES || reduce_size < COLUMNS_MIN_REDUCE return :by_thread, 0 end threads = min(block_size, COLUMNS_MAX_THREADS) - cols = _columns_cols(dst_size) - tiles = cld(dst_size, cols) + cols = _columns_cols(run) + runs = dst_size ÷ run + tiles = runs * cld(run, cols) rows = TILED_STRIDED_ROWS_PER_BLOCK - if dst_size >= reduce_size && tiles < t.target_blocks && + if runs == 1 && dst_size >= reduce_size && tiles < t.target_blocks && block_size % rows == 0 && ispow2(block_size ÷ rows) && reduce_size >= block_size ÷ rows return :tiled_strided, 0 end - dst_size >= COLUMNS_MIN_OUTPUTS && threads >= cols || return nothing + run >= COLUMNS_MIN_OUTPUTS && threads >= cols || return nothing lanes = threads ÷ cols groups = clamp(cld(t.target_blocks, tiles), 1, cld(reduce_size, lanes * MIN_ITEMS_PER_THREAD)) groups > 1 && widemul(dst_size, reduce_size) < t.columns_min_elements && return nothing @@ -450,18 +454,23 @@ function _mapreduce_nd_launch!( elseif shape === :columns || shape === :columns_ordered ordered = shape === :columns_ordered groups = reduce_groups - blocks = cld(dst_size, _columns_cols(dst_size)) * groups + # (runs of `outer_sizes[1]` contiguous outputs, at the other kept dimensions' offsets) + run = outer_sizes[1] + runs = Base.tail(outer_strides), Base.tail(outer_sizes) + cols = _columns_cols(run) + blocks = cld(run, cols) * (dst_size ÷ run) * groups threads = min(block_size, COLUMNS_MAX_THREADS) kernel! = _mapreduce_nd_columns!(backend, threads) - cols = Val(_columns_cols(dst_size)) if groups == 1 kernel!(buffer, dst, f, op, neutral, init, base_offset, reduce_strides[1], dst_size, - reduce_size, 1, cols, Val(ordered); ndrange=(threads * blocks,)) + run, dst_size ÷ run, runs..., reduce_size, 1, Val(cols), Val(ordered); + ndrange=(threads * blocks,)) else # (the partial results of each output's groups, combined by a second pass) partial = bufs.partials kernel!(buffer, partial, f, op, neutral, _NoFinish(), base_offset, - reduce_strides[1], dst_size, reduce_size, groups, cols, Val(ordered); + reduce_strides[1], dst_size, run, dst_size ÷ run, runs..., reduce_size, groups, + Val(cols), Val(ordered); ndrange=(threads * blocks,)) pass2_block_size = _pass2_block_size(groups) if ordered @@ -897,8 +906,9 @@ end end # GPU kernel: columns — a long strided reduction into contiguous outputs, e.g. the row sums of -# a wide column-major matrix (`dims=2`). Each block covers `cols` consecutive outputs with -# `N ÷ cols` lanes each, so that a warp's loads of one step are consecutive elements, and the +# a wide column-major matrix (`dims=2`). Each block covers `cols` consecutive outputs of a run of +# `run` (the outputs at an offset of the other kept dimensions, `run_strides` and `run_sizes`), +# with `N ÷ cols` lanes each, so that a warp's loads of one step are consecutive elements, and the # reduction of each output is split across `groups` blocks. With one group the block stores the # outputs; with more, `dst` holds `output_size × groups` partial results for a second pass (and # `init` is `_NoFinish()`). The lanes and groups take their elements in turn, or `ordered`, each a @@ -908,7 +918,7 @@ end @Const(src), dst, f, op, neutral, init, base_offset, reduce_stride, - output_size, reduce_size, groups, + output_size, run, runs, run_strides, run_sizes, reduce_size, groups, ::Val{cols}, ::Val{ordered}, ) where {cols, ordered} @uniform N = @groupsize()[1] @@ -918,23 +928,30 @@ end iblock = @index(Group, Linear) - 0x1 ithread = @index(Local, Linear) - 0x1 - tiles = cld(output_size, cols) + run_tiles = cld(run, cols) + tiles = run_tiles * runs - # (unsigned divisions by constants and by the tile count; see the tiled kernel) + # (unsigned divisions by constants and by the tile counts; see the tiled kernel. The block + # index is divided in 32 bits: the columns kernel takes fewer than 64 Ki outputs (see + # `_columns_shape`), in tiles of at least 8, split into groups only to approach + # `target_blocks` blocks. 64-bit divisions take it 7% longer on an RTX 6000 Ada) col = Int(unsigned(ithread) % unsigned(cols)) lane = Int(unsigned(ithread) ÷ unsigned(cols)) - tile = Int(unsigned(iblock) % unsigned(tiles)) - group = Int(unsigned(iblock) ÷ unsigned(tiles)) - iout = tile * cols + col + tile = Int((iblock % UInt32) % (tiles % UInt32)) + group = Int((iblock % UInt32) ÷ (tiles % UInt32)) + irun = isempty(run_sizes) ? 0 : Int((tile % UInt32) ÷ (run_tiles % UInt32)) + icol = (tile - irun * run_tiles) * cols + col + valid = icol < run + iout = irun * run + icol + sbase = base_offset + _outer_decode(irun, run_strides, run_sizes) + icol acc = neutral - if ordered && iout < output_size + if ordered && valid # (the group's range of the slice, and the lane's of that, as evenly divided as they can be) lo = fld(group * reduce_size, groups) len = fld((group + 1) * reduce_size, groups) - lo j = lo + fld(lane * len, lanes) stop = lo + fld((lane + 1) * len, lanes) - sbase = base_offset + iout while j + (COLUMNS_ITEMS - 1) < stop for k in 0:(COLUMNS_ITEMS - 1) acc = op_lanes(acc, f_lanes(src[sbase + (j + k) * reduce_stride + 0x1])) @@ -945,8 +962,7 @@ end acc = op_lanes(acc, f_lanes(src[sbase + j * reduce_stride + 0x1])) j += 1 end - elseif iout < output_size - sbase = base_offset + iout + elseif valid step = groups * lanes j = group * lanes + lane # (several loads in flight per thread: one at a time leaves the kernel at half the @@ -989,7 +1005,7 @@ end end end - if lane == 0x0 && iout < output_size + if lane == 0x0 && valid i = iout + group * output_size + 0x1 dst[i] = _finish(op, init, dst, i, sdata[col + 0x1]) end diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index c76ba98e..76f88f2a 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -1207,6 +1207,16 @@ end r = AK.reduce(+, a; dims=2, init=zero(A), alg) @test eltype(r) === A && Array(r) == sum(h; dims=2) end + # Runs of contiguous outputs at the offsets of other kept dimensions: the middle dimension of + # 3-d and 4-d arrays, and a view at an offset whose kept dimensions are not adjacent in memory + h = rand(Int32(-9):Int32(9), 50, 30_000, 7) + for (a, ha) in ((array_from_host(rand(Int32(-9):Int32(9), 96, 7000, 10)), nothing), + (array_from_host(rand(Int32(-9):Int32(9), 40, 30_000, 3, 2)), nothing), + (view(array_from_host(h), 3:42, :, 1:2:7), h[3:42, :, 1:2:7])) + ha = something(ha, Array(a)) + TEST_KERNELS && @test shape_of(a, (2,), Int32) === :columns + @test Array(AK.reduce(+, a; dims=2, init=Int32(0), alg)) == sum(ha; dims=2) + end end @@ -1631,7 +1641,7 @@ affine_rand(::Type{T}) where {T} = rand(T) | one(T) ((40, 5, 30), 3), ((40, 5, 30), (1, 3)), ((9, 4, 3000), (2, 3)), ((8, 300_000), 2), ((40, 70_000), 2), ((256, 256), 2), ((300, 300), 2), ((3000, 300), 1), ((300_000, 2), 1), ((70_000, 40), 1), ((64, 4000), 1), ((512, 40), 1), - ((1, 5000, 2), 2)) + ((1, 5000, 2), 2), ((40, 3000, 30), 2)) for (sz, dims) in cases if TEST_KERNELS segs = AK._mapreduce_segments(sz, AK._mapreduce_dense_strides(sz),