From fb353af9138a2676f7fe3820ebc19ba305212bf5 Mon Sep 17 00:00:00 2001 From: Azamat Berdyshev Date: Wed, 28 Oct 2026 00:23:41 +0000 Subject: [PATCH] Widen the flag of `Optional` partial results to fill their padding 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`. --- src/accumulator.jl | 28 +++++++++++++++++++++------- test/algorithms.jl | 13 +++++++++++-- test/generic/accumulate.jl | 2 +- test/generic/tunings.jl | 3 ++- 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/accumulator.jl b/src/accumulator.jl index 447fbf19..5539d1b2 100644 --- a/src/accumulator.jl +++ b/src/accumulator.jl @@ -84,17 +84,31 @@ resolve_acctype(op, ::Type, ::Type, acctype) = # gives the same result as seeding with a neutral element. # # An unset `Optional` of a bits type (every one kernels see) holds zero bits; one of another type -# leaves the value undefined, which is how it is recognised. -struct Optional{T} +# leaves the value undefined, which is how it is recognised. The flag `F` is `flag_type(T)`. +struct Optional{T, F} value::T - isset::Bool - Optional{T}() where {T} = isinline(T) ? new{T}(zero_bits(T), false) : new{T}() - Optional{T}(value) where {T} = new{T}(value, true) - Optional{T}(value, isset::Bool) where {T} = new{T}(value, isset) + isset::F + Optional{T, F}() where {T, F} = new{T, F}() + Optional{T, F}(value, isset::Bool) where {T, F} = new{T, F}(value, isset) +end +Optional{T}() where {T} = + isinline(T) ? Optional{T, flag_type(T)}(zero_bits(T), false) : Optional{T, flag_type(T)}() +Optional{T}(value) where {T} = Optional{T, flag_type(T)}(value, true) +Optional{T}(value, isset::Bool) where {T} = Optional{T, flag_type(T)}(value, isset) + +# The flag of an `Optional{T}`. A `Bool` would leave the bytes up to `T`'s alignment as padding, +# which stores of an `Optional` don't write but kernels load (partials are loaded whole) and the +# host copies, so Compute Sanitizer's initcheck reports reading them. For a bits type, the flag is +# the unsigned integer that fills those bytes, up to 8 of them: OpenCL's SPIR-V has no wider +# integers, so an `Optional` of a 16-byte aligned type keeps some padding (as does a `T` with +# padding of its own). +@generated function flag_type(::Type{T}) where {T} + isbitstype(T) || return Bool + return (UInt8, UInt16, UInt32, UInt64)[trailing_zeros(min(Base.datatype_alignment(T), 8)) + 1] end isinline(::Type{T}) where {T} = isbitstype(T) || Base.isbitsunion(T) -@inline isset(l::Optional{T}) where {T} = isinline(T) ? l.isset : isdefined(l, :value) +@inline isset(l::Optional{T}) where {T} = isinline(T) ? !iszero(l.isset) : isdefined(l, :value) # A value of the bits type or bits union `T` with all-zero bits function zero_bits(::Type{T}) where {T} diff --git a/test/algorithms.jl b/test/algorithms.jl index b441c93d..2b64a11e 100644 --- a/test/algorithms.jl +++ b/test/algorithms.jl @@ -533,6 +533,15 @@ end @test_throws ArgumentError resolve_reduce(AK.Auto(); T=Union{Missing, Int32}) @test_throws ArgumentError resolve_reduce(AK.BlockReduce(); T=String) @test resolve_reduce(AK.Auto(); T=String, backend=HOST) isa AK.CPUThreads.Partitioned + + # Partial results of an operator without a neutral element are `Optional`s, which kernels load + # whole and the host copies: for these types, they have no padding that stores leave unwritten + for T in (Bool, Int8, Int32, Float32, Float64, ComplexF32, NTuple{3, Int32}, + Tuple{Float64, Float64}) + @test !Base.datatype_haspadding(typeof(AK.reduce_seed((a, b) -> a, T, nothing))) + end + # (but their flag is no wider than 64 bits, which OpenCL's SPIR-V can't go beyond) + @test AK.flag_type(NTuple{4, VecElement{Float32}}) === UInt64 end @@ -594,13 +603,13 @@ end @test_throws ArgumentError resolve_scan(AK.ScanPrefixes(); dims=1) @test_throws ArgumentError resolve_scan(AK.SliceScan()) - # Tiles of partial results: lanes (a Bool and a value) are wider than elements, so fewer fit. + # Tiles of partial results: lanes (a value and a flag) are wider than elements, so fewer fit. # NTuple{3, Int32} is 4-byte aligned everywhere, so its lanes take 16 bytes on every platform # (Int64 lanes take 12 on 32-bit x86) with_tuning(AK.ScanTuning(block_size=256, local_mem_bytes=32 * 1024, max_items=8)) do @test resolve_scan(AK.Auto(); T=NTuple{3, Int32}) === AK.ScanPrefixes(256, 8) @test AK.resolve_scan(AK.Auto(), FAKE, NTuple{3, Int32}, nothing, - AK.Optional{NTuple{3, Int32}}) === AK.ScanPrefixes(256, 7) + typeof(AK.Optional{NTuple{3, Int32}}())) === AK.ScanPrefixes(256, 7) end # Kernels need a bits-type element diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index 117578ea..461bcd6e 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -213,7 +213,7 @@ end # in turn. Tiles of exactly one and two blocks' worth, and one element more, for both block # sizes. for neutral in (I2, nothing), inclusive in (true, false), block_size in (64, 256) - S = neutral === nothing ? AK.Optional{Mat2} : Mat2 + S = neutral === nothing ? typeof(AK.Optional{Mat2}()) : Mat2 tile = !TEST_KERNELS ? 2 * block_size : block_size * AK.resolve_scan(AK.SliceScan(; block_size), BACKEND, Mat2, 1, S).items_per_thread diff --git a/test/generic/tunings.jl b/test/generic/tunings.jl index f5014a01..45bb8684 100644 --- a/test/generic/tunings.jl +++ b/test/generic/tunings.jl @@ -34,7 +34,8 @@ end end for dims in (nothing, 1) @test AK.resolve_scan(AK.Auto(), BACKEND, T, dims) isa AK.Algorithm - @test AK.resolve_scan(AK.Auto(), BACKEND, T, dims, AK.Optional{T}) isa AK.Algorithm + @test AK.resolve_scan(AK.Auto(), BACKEND, T, dims, typeof(AK.Optional{T}())) isa + AK.Algorithm end if kernels @test AK.resolve_scan(AK.ScanPrefixes(), BACKEND, T, nothing) isa AK.ScanPrefixes