Skip to content

Widen the flag of Optional partial results to fill their padding - #192

Open
AzamatB wants to merge 2 commits into
JuliaGPU:mainfrom
AzamatB:ab/optional-padding
Open

AzamatB wants to merge 2 commits into
JuliaGPU:mainfrom
AzamatB:ab/optional-padding

Conversation

@AzamatB

@AzamatB AzamatB commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

A reduction or scan whose operator has no known neutral element (no neutral keyword and no GPUArraysCore.neutral_element method) keeps its partial results as Optional{T}: the value and a Bool flag, padded to T's alignment. LLVM stores the fields separately (st.global.b64 + st.global.b8), but the second pass loads partials whole (ld.global.nc.v2.b64) and the host copies the last one whole, reading bytes nothing wrote. Results are correct, but Compute Sanitizer's initcheck reports every such reduction, including Base reduce with a custom operator, which GPUArrays 12 sends to AK:

AK.reduce((a, b) -> max(a, b), CuArray(rand(100_000)); init=0.0)
# ========= Uninitialized __global__ memory read of size 16 bytes   (49 times)
# ========= Host API memory access error at host access to ... of size 16 bytes

This PR makes the flag the unsigned integer that fills those bytes: Optional{T, F} with F = flag_type(T), e.g. UInt64 for Float64. Sizes don't change, and storing an Optional writes all of its bytes. Bits unions and non-bits types keep a Bool. CUB's scan tile descriptors do the same. The flag is at most 64 bits, as OpenCL's SPIR-V has no wider integers, so 16-byte aligned types keep 8 bytes of padding; padding inside T isn't covered either. Storing the partial whole doesn't help, since LLVM treats padding as undefined.

A script of reductions and scans without a neutral element gives 309 initcheck reports on main and 2 with this PR (from Int128 and NTuple{4, VecElement{Float32}}). The test checks that partials have no padding for common element types (failing on main) and that the flag stays within 64 bits. reduce/accumulate tests pass on CUDA (RTX 5080, Orin, Xavier, Jetson Nano), AMDGPU, PoCL, oneAPI and Metal.

The layout changes codegen. Measured on top of #194, which keeps Optional partials out of stack slots (without it, the comparison is dominated by ptxas promoting those slots only for the wide flag), the time ratios of this PR for Float32 max, 10^4 to 10^8 elements, median of 3 runs:

CUDA (RTX 5080) Metal (M1) oneAPI (A770) PoCL
reductions 0.96–1.02 0.97–1.0 0.87–1.04 0.81–1.04
accumulate! 0.64–0.87 0.96–0.98 0.76–0.89 1.19–1.28
accumulate, dims=1 0.56–0.92 0.99–1.15 0.88–1.0 1.0–1.44

On PoCL, the work-group vectorizer accesses the wider flag in local-memory scan tiles with AVX-512 gathers and scatters, where it leaves the i8 flag scalar. The Metal scans along dims also keep the flag in threadgroup memory, now as i32 accesses instead of i8; I couldn't see why that's slower without the AGX assembly. Choosing the flag type per backend through the tunings was considered: it takes 15–20 lines through reduce_seed and the kernel-side constructors, and loses initcheck coverage on those backends, so I left it out.

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

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