Repository navigation
Choose the flag layout of Optional partials per backend - #195
Merged
Merged
Conversation
maleadt
force-pushed
the
tb/optional-flag-tuning
branch
from
October 28, 2026 14:40
f1b5600 to
9871af3
Compare
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 OpenCL and on Metal from Julia 1.12. Before 1.12, a `Bool` flag on Metal crashes Apple's compiler (Metal.jl#1125). 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.
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.
maleadt
force-pushed
the
tb/optional-flag-tuning
branch
from
November 5, 2026 09:28
8317fe6 to
e349535
Compare
maleadt
marked this pull request as ready for review
November 5, 2026 09:28
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.
Stacked on #194.
#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. On Metal, only from Julia 1.12: before that, kernels can merge partials through stack slots, and SROA then loads the padding after aBoolflag as a byte-alignedi24, which crashes Apple's compiler (JuliaGPU/Metal.jl#1125). 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.