diff --git a/src/reduce/mapreduce_nd.jl b/src/reduce/mapreduce_nd.jl index 59b8cf47..569206d9 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 +# - 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 @@ -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 @@ -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,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 @@ -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 @@ -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 @@ -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] @@ -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])) @@ -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 @@ -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 diff --git a/test/generic/reduce.jl b/test/generic/reduce.jl index 92df385e..76f88f2a 100644 --- a/test/generic/reduce.jl +++ b/test/generic/reduce.jl @@ -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 @@ -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),