Skip to content

Support bits-union element types in reductions and scans - #209

Merged
maleadt merged 4 commits into
mainfrom
tb/isbits-unions
Nov 6, 2026
Merged

maleadt merged 4 commits into
mainfrom
tb/isbits-unions

Conversation

@maleadt

@maleadt maleadt commented Nov 5, 2026 •

Copy link
Copy Markdown
Member

Arrays of isbits unions such as Union{Missing, Float32} or Union{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, cumsum or accumulate of such an array threw ArgumentError: ... 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 for accumulate.

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 Optional of 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:

  • Tests on every backend. The test gate that left out Metal, oneAPI and OpenCL is gone: their current releases store bits-union arrays. The test environment requires those releases (Metal 1.12, oneAPI 2.10, OpenCL 1, AMDGPU 2.10, CUDACore 6.5, KernelAbstractions 0.10.5).
  • Sorting on Julia threads. A sort in one task passed AK's scratch to Base.sort!. Base's radix sort of the non-missing values 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.
  • Accumulator types. promote_type of two different unions is abstract, so e.g. accumulate(+, v; init=0) of a Union{Int32, Float32} vector ran in Any. 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. accumulate allocates Base's result type, which is the type of the stored values, e.g. Union{Float32, Int} for cumsum.
  • Kernels. A seed of a union Optional is passed as an OptionalArg (an Adapt rule), and kernels turn it back into the Optional. 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 an InexactError for union members the destination cannot hold, instead of compiling a dynamic MethodError.

Two deliberate differences from Base, both documented in "Differences from Base":

  • Out-of-place reductions along dims keep AK's contract: the result has the accumulator type. Base uses Base.promote_union of the element type there, e.g. Float32 for Union{Int32, Float32}.
  • Where Base's inference widens a union of three or more members to Any, AK keeps the bits union, e.g. for cumsum(Union{Missing, Int32, Float32}[...]).

Known issue: on Julia 1.10, sorting Union{Int32, Float32} arrays gives wrong results on PoCL (the --cpu-ka and 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 and SPIRV_LLVM_Backend_jll 23.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).

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.
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
maleadt merged commit d990e0d into main Nov 6, 2026
45 of 49 checks passed
@maleadt
maleadt deleted the tb/isbits-unions branch November 6, 2026 07:40
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