Repository navigation
Conversation
`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.
A partial result of a reduction or scan whose operator has no neutral
element is an `Optional{T}`, the value and a `Bool` flag. For a bits
type the flag is padded to `T`'s alignment, and storing an `Optional`
writes the fields only, so the padding of partials in device memory was
never written. The second pass loads partials whole and the host copies
the last one, so Compute Sanitizer's initcheck reported every such
reduction, including Base reductions with a custom operator that
GPUArrays sends to AK.
The flag is now the unsigned integer that fills those bytes
(`Optional{T, F}` with `F = flag_type(T)`), up to 64 bits, which is as
wide as integers go in OpenCL's SPIR-V. Sizes are unchanged, and storing
an `Optional` writes all of it, unless `T` has padding of its own or is
aligned to more than 8 bytes. Other types keep a `Bool`.
The flag that fills the padding after the value of `Optional` partials (up to 8 bytes) keeps Compute Sanitizer's initcheck quiet for most element types and makes scans faster on CUDA, but in local memory some compilers handle it worse than a `Bool`: PoCL accesses it with AVX-512 gathers and scatters, and Metal's scans along `dims` get slower. A new tuning field, `wide_optional_flag` of `ReduceTuning` and `ScanTuning`, selects the layout: the wide flag by default, a `Bool` on Metal and OpenCL. The choice is made at run time, behind a function barrier, from the canonical (wide) seed of the plan: kernels make every `Optional` like the seed, and the buffers of partials, which plans and workspaces size with the wide flag, are reinterpreted (both layouts have the same size). So plans, workspaces and result types don't depend on the tuning.
maleadt
force-pushed
the
tb/optional-flag-tuning
branch
from
October 28, 2026 14:40
f1b5600 to
9871af3
Compare
Only PoCL's CPU devices slow down with the wide flag. Intel's CPU runtime (2026.21) runs scans of Float32s along `dims=1` in 0.81-0.85x the time it takes with a `Bool` flag, and with a `Bool` flag it crashes compiling the `accumulate!` tile kernel (a segfault in X86 instruction selection). On an Arc A770 (NEO), scans take 0.73-1.0x as long with the wide flag, and NVIDIA's OpenCL runs both layouts alike.
This branch has not been deployed
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.
Builds on #194 and #192; only the last commit is new.
#192 widens the flag of
Optionalpartials to fill their padding, which keeps Compute Sanitizer's initcheck quiet. How that layout performs depends on the backend's compiler:Boolflag splits each access to a partial into a 32- and an 8-bit instruction, and the tile scans are bound by shared memory (~11-way bank conflicts), so one 64-bit access halves their wavefronts (scan kernel of a 10^7accumulate!: 214 → 120 µs on an RTX 5080). With the bank conflicts fixed (Pad the runs of the scan tiles to an odd stride in local memory #202) the gain shrinks to 1.15–1.2x (64 → 54 µs), so wide is still the better default there.{float, i8}loads and stores that LLVM keeps whole into separate 32-bit and byte SLM accesses and register spills (7680 against 1504 bytes per thread inscan_tile_kernel).UInt32flag takesslice_tiles_kernel!'s output pass from 67 to 70 registers and over an occupancy step (3 → 2 blocks of 256 per core), making scans alongdimsup to 1.2x slower.This PR adds a tuning field,
wide_optional_flag(inReduceTuningandScanTuning), that picks the layout per backend: the wide flag by default, aBoolon Metal and on PoCL's CPU devices. Host paths and other OpenCL drivers keep the wide flag.The choice is a run-time
Bool, so plans, workspaces and result types don't depend on it. Setup still builds the wide seed and sizes the buffers of partials with it. At the run entries (mapreduce_run,mapreducedim_run!,scan_run!),with_flag_layoutlooks up the tuning (only when the seed is anOptional) and calls the rest either as is, or with the seed converted to aBoolflag and those buffersreinterpreted (both layouts have the same size). Both branches return the same type, so results still infer. Kernels now build everyOptionallike the seed (OptionalMap,OptionalOp,lift) rather than fromflag_type.Time ratios against #194 alone,
Float32maxwithout a neutral element, 10^4 to 10^8 elements (minimum of 3–6 runs; the GPUs and PoCL were shared with other jobs):BoolBoolaccumulate!accumulate,dimsAgainst main, Metal's scans along
dimsare now 0.97–1.04, where #192 alone gave 1.06–1.19. The AMDGPU numbers confirm the wide default there. (oneAPI'saccumulate!of 10^4 elements varies with launch latency: 1.16 back to back, 0.89–0.90 with a 50 µs gap between calls.) These numbers predate #194's second commit, which only changes the combine of non-primitive values; a tuple accumulator still works with theBoollayout on PoCL.Other OpenCL drivers, measured against the
Boollayout: on an Arc A770 (NEO) scans take 0.73–1.0x as long with the wide flag, NVIDIA's OpenCL (RTX 5080, RTX 6000 Ada) runs both alike, and Intel's CPU runtime (2026.21) is 0.81–0.85x alongdims=1(0.98–1.3x alongdims=2of 8-row matrices) and crashes compiling theaccumulate!tile kernel with aBoolflag and #194's combine (a segfault in its X86 instruction selection). So only PoCL gets theBoolflag, by platform name.Open questions:
Boolwould make its partialsReinterpretArrays.