Repository navigation
Support bits-union element types in reductions and scans - #209
Merged
Merged
Conversation
Metal.jl 1.12, oneAPI.jl 2.10, OpenCL.jl 1.0 and AMDGPU.jl 2.10 store arrays and local memory of bits-union element types, and CUDA.jl 6.5 reads them through `Const`, so the tests no longer leave out Metal, oneAPI and OpenCL. The test environment requires those releases instead, and KernelAbstractions 0.10.5 for its PoCL backend.
A sort in one task passed the plan's scratch to `Base.sort!`, whose radix sort of the values
without `missing` reinterprets the scratch as unsigned integers. That fails for a scratch of a
union element type, such as `Union{Missing, Float32}`, from a few thousand elements on. Base
allocates its own scratch for such elements now.
`promote_type` of two different unions is abstract (`promote_type(Union{Float32, Int32},
Union{Float32, Int64})` is `Real`), so the running values of `accumulate(+, v; init=0)` of a
`Union{Int32, Float32}` vector were `Any`. Types are now joined into a bits union where promotion
gives an abstract type, as Base's inference does, and the operator's results are inferred per
union member where inference does not split them: through a call for `Base.reduce_first`, whose
fallback otherwise matches the whole union, and member by member where there are more
combinations than inference splits.
`accumulate` also allocates Base's result type, the type of the stored values, rather than that
of the running values, which include one-element partial results: `Union{Float32, Int}` rather
than `Union{Float32, Int32, Int}` in that example.
maleadt
force-pushed
the
tb/isbits-unions
branch
from
November 5, 2026 20:08
2b310c4 to
756e85c
Compare
A reduction or scan whose accumulator type is a bits union, such as `sum` or `cumsum` of a
`Union{Missing, Float32}` array, seeds its partial results with an `Optional` of that union. It
is stored inline, in arrays and local memory, but is not a bits type, which kernel arguments must
be, so such accumulators were rejected. The seed is now passed as an `OptionalArg`, with a slot
for each union member and the index of the one in use, so that its type does not depend on the
value. Kernels turn it back into the `Optional`.
The kernels also combine two union values through calls with one union argument each, which
inference splits into their members, and scans convert their running values to the
destination's element type with an `InexactError` for union members it cannot hold (those of
one-element partial results) rather than a dynamic `MethodError`, which kernels cannot throw.
maleadt
force-pushed
the
tb/isbits-unions
branch
from
November 5, 2026 20:08
756e85c to
ea5809e
Compare
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.
Arrays of isbits unions such as
Union{Missing, Float32}orUnion{Int32, Float32}already worked with the sorts,findall,searchsorted*,reverse,map!and the predicates. They also worked with reductions whose accumulator is a bits type. They did not work with reductions or scans whose accumulator is a union:sum,maximum,cumsumoraccumulateof such an array threwArgumentError: ... is not a bits type, which the kernels need. This PR makes those work too, with Base's values, and with Base's result types foraccumulate.The cause was more than the check. The accumulator type of such a reduction is a union, so it has no neutral element, and the partial results start from an
Optionalof the union. That struct is stored inline (in arrays and local memory) but is not a bits type, and kernel arguments must be bits types. Changes, one commit each:Base.sort!. Base's radix sort of the non-missingvalues reinterprets that scratch, which fails for a union scratch from about 3000 elements (sort!(v; scratch=similar(v))reproduces it in plain Base). For such elements Base now allocates its own scratch.promote_typeof two different unions is abstract, so e.g.accumulate(+, v; init=0)of aUnion{Int32, Float32}vector ran inAny. Types are now joined into a bits union where promotion would be abstract. The operator's result types are inferred per union member where inference does not split the union.accumulateallocates Base's result type, which is the type of the stored values, e.g.Union{Float32, Int}forcumsum.Optionalis passed as anOptionalArg(anAdaptrule), and kernels turn it back into theOptional. Two union values are combined through calls with one union argument each, so that inference splits them. A scan's final conversion to the destination raises anInexactErrorfor union members the destination cannot hold, instead of compiling a dynamicMethodError.Two deliberate differences from Base, both documented in "Differences from Base":
dimskeep AK's contract: the result has the accumulator type. Base usesBase.promote_unionof the element type there, e.g.Float32forUnion{Int32, Float32}.Any, AK keeps the bits union, e.g. forcumsum(Union{Missing, Int32, Float32}[...]).Known issue: on Julia 1.10, sorting
Union{Int32, Float32}arrays gives wrong results on PoCL (the--cpu-kaand OpenCL jobs). PoCL's CBS work-group method doesn't give each work-item its own copy of a private stack slot that reaches a barrier only through a pointer phi, which Julia 1.10 generates for union-typed locals. AdaptiveCpp fixed the same code in AdaptiveCpp/AdaptiveCpp#1763. The earlier SPIR-V validation failure was an LLVM bug, fixed by llvm/llvm-project#239343 andSPIRV_LLVM_Backend_jll23.1.4+1.Tested with Julia 1.13.2: CPU (threads), CUDA (RTX 5080), OpenCL on PoCL,
--cpu-ka(except the issue above), and Metal (M-series).