Skip to content

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

Draft
maleadt wants to merge 5 commits into
mainfrom
tb/optional-flag-tuning
Draft

maleadt wants to merge 5 commits into
mainfrom
tb/optional-flag-tuning

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026 •

Copy link
Copy Markdown
Member

Builds on #194 and #192; only the last commit is new.

#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. 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.

maleadt and others added 4 commits October 28, 2026 11:34
`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
maleadt force-pushed the tb/optional-flag-tuning branch from f1b5600 to 9871af3 Compare October 28, 2026 14:40
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

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.

2 participants