Skip to content

Combine Optional partials without returning either operand - #194

Open
maleadt wants to merge 2 commits into
mainfrom
tb/optional-op-select
Open

maleadt wants to merge 2 commits into
mainfrom
tb/optional-op-select

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026 •

Copy link
Copy Markdown
Member

Reductions and scans whose operator has no known neutral element (no neutral keyword and no GPUArraysCore.neutral_element method, so also Base reduce with a custom operator via GPUArrays) keep their partial results as Optionals, combined by OptionalOp. That returned one of its operands or a new Optional, so the compiler merged the possible results through stack slots: two allocas and a phi ptr between them, which SROA can't promote. For example, in the loop of the one-thread-per-output reduction kernel:

%.sroa.08.1.in = phi ptr [ %"new::Optional40", %L130 ], [ %"new::Optional", %L102 ]
%.sroa.08.1 = load i64, ptr %.sroa.08.1.in, align 4

On CUDA that stays local memory: the SASS of that kernel keeps 4 STL and 1 LDL (local-memory stores and loads) for the running partial. This PR builds the result from the operands' fields instead (ifelse on the values, op called only when both are set), and the PTX has no local memory left. This only applies to primitive types: Julia selects between other values, such as the tuples of a mapreduce with (max(p[1], q[1]), max(p[2], q[2])), through stack slots too, which made such a pair reduction of 10^7 elements 3.3x slower in its kernels on an RTX 6000 Ada (14.1 µs on main, 46.6 µs). Those, and non-bits types (whose value can be undefined), keep the old code.

New/main time ratios for Float32 max without a neutral element, 10^4 to 10^8 elements (median of 3 runs):

CUDA (RTX 5080) Metal (M1) oneAPI (A770) PoCL
whole-array reduce 0.28–0.98 0.75–0.98 0.66–0.90 0.75–0.96
dims=2 reduce 0.23–0.97 1.0 0.81–0.94 0.42–0.81
accumulate! 0.76–0.98 0.73–0.96 0.87–0.99 0.84–1.08
accumulate, dims=1 0.95 1.0 (1.10 at 10^8) 0.88–0.94 1.08–1.22

Some cases get slower. PoCL scans along dims that take a block per slice get 1.03–1.3x slower: PoCL runs each work-item's serial scan of its run (thread_scan!) as scalar code, and LLVM compiles the new combine into a chain of two masked blends after op (select(sa, select(sb, op, a), b)), where main had branches that were always predicted. Fewer instructions run, but each step of the dependency chain takes longer. GPUs hide that latency; using the old combine in that kernel only would recover PoCL but give back the 4–12% gain on CUDA, so I left it.

On CUDA, accumulate along dims=2 of an (8, 10^7) array gets 1.12x slower on an RTX 5080 (not on an RTX 6000 Ada). That shape takes the split one-thread-per-slice scan, whose totals pass is bound by memory latency: on main, ptxas happened to unroll its loop and keep several loads in flight, which it doesn't do for the select-based combine. #199 makes that kernel load eight elements at a time, which makes this scan about 2x faster than main.

On oneAPI, small sizes are dominated by a 55–155 µs fixed per-call latency that comes in two modes depending on how soon the next call follows, so medians of back-to-back runs at 10^4 elements flip between them (one run gave the dims=2 reduce 1.35). The kernel itself is faster there (8.8–9.7 µs against 12.9–14.8 µs), as are the minimum wall times (71–72 µs against 74–78 µs).

Known issue: with this PR alone (a Bool flag), Intel's CPU OpenCL runtime (2026.21, not in CI) segfaults in its X86 instruction selection while compiling the accumulate! tile kernel; main and the wide flag of #192 compile fine. Merging this together with #192 (or #195, which keeps the wide flag there) avoids it. Tracked in JuliaGPU/OpenCL.jl#600 (not reduced yet).

Tested reduce/accumulate on CUDA, PoCL, Metal, oneAPI, and with OpenCL on an Arc A770 and NVIDIA GPUs.

`OptionalOp`, which combines the partial results of reductions and scans
whose operator has no neutral element, returned one of its operands or a
new `Optional`. The compiler merged those results through stack slots (a
phi of pointers to allocas, which SROA cannot promote), so on CUDA the
reduction loops kept local-memory loads and stores. Building one result
from the operands' fields avoids the stack slots.

Reductions without a neutral element get up to 3.5x faster on CUDA,
2.4x on PoCL, 1.5x on oneAPI and 1.3x on Metal.
For a value that isn't of a primitive type, such as the tuple of a
`mapreduce` with `(max(p[1], q[1]), max(p[2], q[2]))`, Julia lowers the
select between the operands' values through stack slots, so the
running partial went through local memory on every combine. On an
RTX 6000 Ada that made a `Float32` pair reduction of 10^7 elements 3.3x
slower in its kernels (14.1 us on main, 46.6 us with the select). Such
values keep the old combine, which returns either operand.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant