Skip to content

Choose the flag layout of Optional partials per backend - #195

Merged
maleadt merged 2 commits into
tb/optional-op-selectfrom
tb/optional-flag-tuning
Nov 5, 2026
Merged

maleadt merged 2 commits into
tb/optional-op-selectfrom
tb/optional-flag-tuning

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026 •

Copy link
Copy Markdown
Member

Stacked on #194.

#192 widens the flag of Optional partials to fill their padding, which keeps Compute Sanitizer's initcheck quiet. How that layout performs depends on the backend's compiler:

  • CUDA: a Bool flag 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^7 accumulate!: 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.
  • oneAPI: IGC turns the padded {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 in scan_tile_kernel).
  • Metal: registers are allocated in 16-bit halves, so a UInt32 flag takes slice_tiles_kernel!'s output pass from 67 to 70 registers and over an occupancy step (3 → 2 blocks of 256 per core), making scans along dims up to 1.2x slower.
  • PoCL: the work-group vectorizer accesses the wide flag in local-memory tiles with AVX-512 gathers and scatters, making scans 1.0–1.4x slower.

This PR adds a tuning field, wide_optional_flag (in ReduceTuning and ScanTuning), that picks the layout per backend: the wide flag by default, a Bool on 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 a Bool flag as a byte-aligned i24, 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_layout looks up the tuning (only when the seed is an Optional) and calls the rest either as is, or with the seed converted to a Bool flag and those buffers reinterpreted (both layouts have the same size). Both branches return the same type, so results still infer. Kernels now build every Optional like the seed (OptionalMap, OptionalOp, lift) rather than from flag_type.

Time ratios against #194 alone, Float32 max without a neutral element, 10^4 to 10^8 elements (minimum of 3–6 runs; the GPUs and PoCL were shared with other jobs):

CUDA (RTX 5080) AMDGPU (gfx1036) oneAPI (A770) Metal (M1) PoCL
layout wide wide wide Bool Bool
reductions 0.99–1.13 1.0–1.2 (1.47 once) 0.66–1.03 0.95–1.08 0.84–1.06
accumulate! 0.62–0.88 0.65–0.82 0.76–0.86 0.98–1.08 0.96–1.02
accumulate, dims 0.55–0.91 0.63–0.82 0.82–0.97 0.96–1.07 0.90–1.07

Against main, Metal's scans along dims are now 0.97–1.04, where #192 alone gave 1.06–1.19. The AMDGPU numbers confirm the wide default there. (oneAPI's accumulate! 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 the Bool layout on PoCL.

Other OpenCL drivers, measured against the Bool layout: 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 along dims=1 (0.98–1.3x along dims=2 of 8-row matrices) and crashes compiling the accumulate! tile kernel with a Bool flag and #194's combine (a segfault in its X86 instruction selection). So only PoCL gets the Bool flag, by platform name.

Open questions:

  • Tests: every backend in CI runs one layout. Testing both on one backend, and a workspace reused across a change of layout, needs a way to override a real backend's tuning in the tests; sol checked both by hand.
  • Each reduction or scan without a neutral element compiles both host branches. The kernels compile only for the layout used.
  • One field in two tunings, or a separate hook?
  • With KernelAbstractions 0.10, the host backend's kernels also run on PoCL, but switching the host to Bool would make its partials ReinterpretArrays.

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
maleadt force-pushed the tb/optional-flag-tuning branch from 8317fe6 to e349535 Compare November 5, 2026 09:28
@maleadt
maleadt changed the base branch from main to tb/optional-op-select November 5, 2026 09:28
@maleadt
maleadt marked this pull request as ready for review November 5, 2026 09:28
@maleadt
maleadt added this pull request to stack #203 November 5, 2026 09:28
@maleadt
maleadt merged commit ba29949 into main Nov 5, 2026
16 of 17 checks passed
@maleadt
maleadt deleted the tb/optional-flag-tuning branch November 5, 2026 13:29
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