Skip to content

Faster strided reductions into contiguous outputs (dims=2 of matrices, middle dimensions of arrays) #176

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maleadt merged 2 commits into from
Oct 25, 2026
Merged

Faster strided reductions into contiguous outputs (dims=2 of matrices, middle dimensions of arrays) #176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0436c65
Reduce strided slices with the tiled kernel where one thread per outp…
maleadt Oct 24, 2026
cb9b168
Reduce the middle dimensions of arrays with the columns kernel
maleadt Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 87 additions & 49 deletions src/reduce/mapreduce_nd.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
# - 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
# (and dst_size < reduce_size) -> :multigroup (split one reduction
Expand All @@ -135,9 +141,14 @@ _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, osizes[1], 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
Expand All @@ -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
Expand All @@ -208,6 +198,42 @@ function _mapreduce_nd_shape(dst_size, reduce_size, segs, block_size, t::ReduceT
end
end

# 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 (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, 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(run)
runs = dst_size ÷ run
tiles = runs * cld(run, cols)
rows = TILED_STRIDED_ROWS_PER_BLOCK
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
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
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
Expand Down Expand Up @@ -428,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
Expand Down Expand Up @@ -875,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
Expand All @@ -886,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]
Expand All @@ -896,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]))
Expand All @@ -923,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
Expand Down Expand Up @@ -967,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
Expand Down
27 changes: 26 additions & 1 deletion test/generic/reduce.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1192,6 +1192,31 @@ 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
# 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


Expand Down Expand Up @@ -1616,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),
Expand Down
Loading