Repository navigation
Choose the launch of reductions along dims once, as shape structs - #181
Merged
Merged
Conversation
The plan sized the scratch by choosing a launch shape, and the run chose it again, along with the source's layout and segments. The ordered shape had to keep the commutative shape's blocks per output so that both took the same scratch, and `plan(mapreduce, ...; commutative)` ignored `commutative`. The setup now chooses the layout, segments and final shape once (in element order where `commutative=false`) and keeps them in a type-stable NamedTuple, from which both the scratch and the run derive. The run only takes the buffer from the source it reduces, which before Julia 1.12 is a materialized `Broadcasted` object.
Shapes were a `Symbol` with an `Int` that meant blocks per output, lanes per output or nothing, and one launcher ran an if-chain over them. They are now structs named along CUB's DeviceSegmentedReduce and PyTorch's ReduceConfig (`ByThread`, `ByLanes`, `ByBlock`, `ByBlocks`, `Columns`, `Tiled`, `Ordered`, `Whole`), each carrying its own parameters, and the kernel shapes have a `launch!` method each, which takes the indices as one NamedTuple. `:by_block` and `:by_block_each` differed only in their block count, which `ByBlock` now holds; `:columns` and `:columns_ordered` are one `Columns` shape. The setup holds the shape behind an abstract field so that it infers. To keep small reductions as fast as before, the part of `mapreduce`'s setup that follows `reduced_dims`, whose result is not inferred, sits behind a function barrier, and the whole-array kernels are still called directly.
The commutative second pass of `ByBlocks` and `Columns` had a kernel of its own, which is the by-block kernel over the partial results (a `dst_size × groups` matrix reduced along its second dimension with `PartialInputs()`), as the ordered second pass already used the ordered kernel. Both passes now go through `combine_groups!`.
The kernel that gave each output of a source without a dense buffer a thread, iterating over the Cartesian indices of its slice, was a copy of the threaded path's loop body. The threaded path now takes the backend and runs that body through `foreachindex_on` on either.
The kernels decoded output offsets with `outer_offset`, which divided by every segment's size including the last, and element offsets with `reduce_offset`, which skips the last division and shifts by sizes that are powers of two. Both decode indices below the product of the sizes, so `segment_offset` serves both with `reduce_offset`'s body.
`segments_lane_bytes` and `segments_short_runs` tune the `ByLanes` shape, and `columns_min_elements` when `Columns` splits its outputs across blocks; "segments" now only means merged dimensions. They become `lanes_max_bytes`, `lanes_short_runs` and `columns_split_min_elems`, spelled like `threads_min_elems` beside them.
The tiled kernel is the columns kernel with one group and one load per thread, and the ByBlocks kernel divides by its output count as a compile-time constant. Folding the first into the columns kernel and dividing by a run-time count in the second were both measured slower, so both stay, with the numbers next to them.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reductions along
dimspick one of several launch shapes (a thread, a group of threads or a block per output, several blocks per output, tiles of strided outputs, an in-order variant) from the sizes and strides of the source and the output. This decision was represented as a(Symbol, Int)pair whose integer meant blocks, lanes or nothing depending on the symbol, it was made twice per call (once to size the scratch memory, once to launch, with the in-order variant forced to reuse the other's group count so both agreed), and the launcher was a longifchain on the symbol. A few kernels also duplicated others.Now the setup chooses the launch once and the scratch and the launch both derive from it. The shapes are small structs that carry their own parameters (
ByThread,ByLanes(lanes),ByBlock(blocks),ByBlocks(groups),Columns(outputs, groups, ordered),Tiled,Ordered(groups)), named along CUB's segmented reduction (a thread, a group of threads or a block per segment) and PyTorch'sReduceConfig(tiles across strided outputs, several blocks per output), each with onelaunch!method.plan(...; commutative)now honourscommutative, so workspaces of in-order reductions get the scratch their shape needs.Kernel deduplication: the second pass of split reductions is the by-block kernel over the partial results; lazy sources reduce one thread per output through the same loop body as the threaded path; one offset decoder. Two merges were measured and rejected, with the numbers in comments: folding the tiled kernel into the columns kernel (17-63% slower for 1024² to 4096²
dims=2on an RTX 5080), and replacing the multi-block kernel's compile-time output count with a run-time one (4-9% slower for 3-16 outputs).ReduceTuning's fields are renamed after the shapes they tune (lanes_max_bytes,lanes_short_runs,columns_split_min_elems); the tuning structs are not public API.Performance, as the geometric mean of new over old times across the benchmark cases (min of several interleaved rounds): RTX 5080 0.985 (144 cases), AMD iGPU 1.002, PoCL 0.986, M1 0.999. Small, host-bound calls are slightly faster (e.g.
mapreducedim!of 64×64 with a workspace 4.5 → 4.2 µs on the 5080). Individual cases that came out slower in the full runs were equal when rerun on their own.Tested on CPU threads, PoCL (also with AK's kernels on the host backend), CUDA (Julia 1.12 and 1.10, which materializes
Broadcastedsources), AMDGPU and Metal.