From f1b1fed6f3013cec99494b69abf814c85635687c Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 00:02:47 +0200
Subject: [PATCH 01/18] KernelInterface: more sub-group communication

Add the sub-group operations that e.g. Molly.jl's CUDA kernels use, so that they
can be written portably:

- shuffles `shfl` (from a given lane), `shfl_up` and `shfl_xor`, next to
  `shfl_down`. Backends implement them for primitive types; a fallback shuffles
  `isbits` structs and tuples field by field, and `supports_shuffle` checks
  their fields.
- votes `sub_group_any`, `sub_group_all` and `sub_group_ballot` (a `UInt64`
  mask, for sub-groups of at most 64 work-items), required with sub-group
  support.
- `get_max_sub_group_size` is now required to be a constant of the generated
  code.

Implement them for POCL; its sub-group width is folded into the IR before
optimization.

Assisted-by: Claude Code (Opus 5.5)
---
 docs/src/kernelinterface.md                |   8 +-
 lib/KernelInterface/src/KernelInterface.jl |   4 +-
 lib/KernelInterface/src/device.jl          | 166 ++++++++++++++++++++-
 lib/KernelInterface/src/host.jl            |  22 ++-
 lib/KernelInterface/test/interface.jl      | 126 ++++++++++++++++
 lib/KernelInterface/test/runtests.jl       |  18 ++-
 src/pocl/backend.jl                        |  27 +++-
 src/pocl/compiler/compilation.jl           |  15 ++
 src/pocl/device/subgroups.jl               |  25 ++++
 src/pocl/pocl.jl                           |   1 +
 test/runtests.jl                           |  16 ++
 11 files changed, 408 insertions(+), 20 deletions(-)
 create mode 100644 src/pocl/device/subgroups.jl

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index 4eba70529..c6f7cfb08 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
-| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries, [`sub_group_barrier`](@ref); if `supports_shuffle(backend, T)`: [`shfl_down`](@ref) for `T` | |
+| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for primitive `T` | |
 
 Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones),
 the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch),
@@ -192,7 +192,13 @@ localmemory
 ### Communication
 
 ```@docs
+shfl
 shfl_down
+shfl_up
+shfl_xor
+sub_group_any
+sub_group_all
+sub_group_ballot
 ```
 
 ### Printing
diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl
index 0f0fade1a..03311eaec 100644
--- a/lib/KernelInterface/src/KernelInterface.jl
+++ b/lib/KernelInterface/src/KernelInterface.jl
@@ -31,7 +31,9 @@ include("host.jl")
             :get_group_id, :get_num_groups,
             :get_sub_group_size, :get_max_sub_group_size, :get_num_sub_groups,
             :get_sub_group_id, :get_sub_group_local_id,
-            :localmemory, :shfl_down, :barrier, :sub_group_barrier, :_print,
+            :localmemory, :shfl, :shfl_down, :shfl_up, :shfl_xor,
+            :sub_group_any, :sub_group_all, :sub_group_ballot,
+            :barrier, :sub_group_barrier, :_print,
             # compilation and launch
             :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"),
             :launch_configuration, :max_work_group_size, :max_work_group_dims,
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index 5b0090914..f106691e0 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -155,12 +155,21 @@ See [`get_local_id`](@ref) for the supported types `T`.
 """
     get_max_sub_group_size([::Type{T}=Int])::T
 
-The sub-group width, [`sub_group_size(backend)`](@ref sub_group_size) on the host.
+The sub-group width (the warp or wavefront size), [`sub_group_size(backend)`](@ref
+sub_group_size) on the host.
+
+It is a compile-time constant of the generated code, so code that depends on it, e.g. a
+loop over the lanes or a shuffle butterfly, is specialized for it. It isn't known during
+type inference, though: to pick types or `Val` parameters from the width, e.g. the integer
+type of a mask with a bit per lane, pass [`sub_group_size(backend)`](@ref sub_group_size)
+from the host.
 
 See [`get_local_id`](@ref) for the supported types `T`.
 
 !!! note
-    Backend implementations that support sub-groups **must** implement:
+    Backend implementations that support sub-groups **must** implement this, returning the
+    width the kernel is compiled for as a constant (not by querying the device at run
+    time):
     ```
     @device_override get_max_sub_group_size(::Type{T})::T where {T}
     ```
@@ -254,6 +263,34 @@ localmemory(::Type{T}, ::Val) where {T} =
 
 ## communication
 
+# Shuffles exchange values between the work-items of a sub-group. Backends implement them for
+# the primitive types for which `supports_shuffle` returns `true`; the fallbacks below
+# shuffle other `isbits` types field by field.
+
+"""
+    shfl(val::T, lane::Integer)::T
+
+Return `val` of the work-item with [`get_sub_group_local_id`](@ref) equal to `lane` in the
+sub-group. When there is no such work-item, the result is an unspecified value (of type `T`).
+
+All work-items of the sub-group have to execute `shfl` together (not in a divergent branch),
+but they may read from different lanes.
+
+`shfl` exchanges values, not memory: it is not a memory fence.
+
+Types for which [`supports_shuffle`](@ref) returns `true` are supported, as well as `isbits`
+structs and tuples of such types, which are shuffled field by field.
+
+!!! note
+    Backend implementations **must** implement this for every primitive type `T` for which
+    [`supports_shuffle`](@ref) returns `true`, and only for those, so that other types reach
+    the fallback that shuffles structs field by field:
+    ```
+    @device_override shfl(val::T, lane::Integer) where {T <: Union{...}}
+    ```
+"""
+@inline shfl(val, lane::Integer) = shfl_fields(x -> shfl(x, lane), val)
+
 """
     shfl_down(val::T, offset::Integer)::T
 
@@ -264,16 +301,131 @@ no such work-item, the result is an unspecified value (of type `T`).
 All work-items of the sub-group have to execute `shfl_down` together (not in a divergent
 branch), with the same `offset`.
 
-`shfl_down` exchanges values, not memory: it is not a memory fence.
+`shfl_down` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for
+the supported types.
+
+!!! note
+    Backend implementations **must** implement this like [`shfl`](@ref):
+    ```
+    @device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}}
+    ```
+"""
+@inline shfl_down(val, offset::Integer) = shfl_fields(x -> shfl_down(x, offset), val)
+
+"""
+    shfl_up(val::T, offset::Integer)::T
+
+Return `val` of the work-item `offset` lanes earlier in the sub-group, i.e. with
+[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() - offset`. When there is
+no such work-item, the result is an unspecified value (of type `T`).
+
+All work-items of the sub-group have to execute `shfl_up` together (not in a divergent
+branch), with the same `offset`.
+
+`shfl_up` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for
+the supported types.
+
+!!! note
+    Backend implementations **must** implement this like [`shfl`](@ref):
+    ```
+    @device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}}
+    ```
+"""
+@inline shfl_up(val, offset::Integer) = shfl_fields(x -> shfl_up(x, offset), val)
+
+"""
+    shfl_xor(val::T, mask::Integer)::T
+
+Return `val` of the work-item whose 0-based lane id is the 0-based lane id of this work-item
+xor `mask`, i.e. with [`get_sub_group_local_id`](@ref) equal to
+`((get_sub_group_local_id() - 1) ⊻ mask) + 1`. When there is no such work-item, the result is
+an unspecified value (of type `T`).
+
+All work-items of the sub-group have to execute `shfl_xor` together (not in a divergent
+branch), with the same `mask`. A butterfly over the masks `width ÷ 2, …, 2, 1` (for the
+sub-group width [`get_max_sub_group_size`](@ref)) reduces a full sub-group such that every
+work-item gets the result.
+
+`shfl_xor` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for
+the supported types.
+
+!!! note
+    Backend implementations **must** implement this like [`shfl`](@ref):
+    ```
+    @device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}}
+    ```
+"""
+@inline shfl_xor(val, mask::Integer) = shfl_fields(x -> shfl_xor(x, mask), val)
+
+# Shuffle a value that the backend doesn't support directly field by field.
+@inline @generated function shfl_fields(f, val::T) where {T}
+    if !isbitstype(T) || isprimitivetype(T)
+        return :(
+            throw(
+                ArgumentError(
+                    $("Shuffling values of type $T is not supported by this backend, see `supports_shuffle`")
+                )
+            )
+        )
+    end
+    fields = (:(f(getfield(val, $i))) for i in 1:fieldcount(T))
+    return Expr(:new, T, fields...)
+end
+
+"""
+    sub_group_any(pred::Bool)::Bool
+
+Whether `pred` is `true` for any work-item of the sub-group. All work-items of the sub-group
+get the same result.
+
+All work-items of the sub-group have to execute `sub_group_any` together (not in a divergent
+branch).
+
+!!! note
+    Backend implementations that support sub-groups **must** implement:
+    ```
+    @device_override sub_group_any(pred::Bool)::Bool
+    ```
+"""
+function sub_group_any end
+
+"""
+    sub_group_all(pred::Bool)::Bool
+
+Whether `pred` is `true` for all work-items of the sub-group. All work-items of the
+sub-group get the same result.
+
+All work-items of the sub-group have to execute `sub_group_all` together (not in a divergent
+branch).
+
+!!! note
+    Backend implementations that support sub-groups **must** implement:
+    ```
+    @device_override sub_group_all(pred::Bool)::Bool
+    ```
+"""
+function sub_group_all end
+
+"""
+    sub_group_ballot(pred::Bool)::UInt64
+
+A mask of the work-items of the sub-group for which `pred` is `true`: bit `i - 1` (counting
+from the least significant bit) is set for the work-item with
+[`get_sub_group_local_id`](@ref) equal to `i`. All work-items of the sub-group get the same
+result. Use e.g. `count_ones` to count the work-items, or `trailing_zeros` to find the first
+one.
+
+All work-items of the sub-group have to execute `sub_group_ballot` together (not in a
+divergent branch). Only sub-groups of at most 64 work-items are supported.
 
 !!! note
-    Backend implementations **must** implement this for every `T` for which
-    [`supports_shuffle`](@ref) returns `true`:
+    Backend implementations that support sub-groups with a width of at most 64 **must**
+    implement:
     ```
-    @device_override shfl_down(val::T, offset::Integer) where T
+    @device_override sub_group_ballot(pred::Bool)::UInt64
     ```
 """
-function shfl_down end
+function sub_group_ballot end
 
 
 ## synchronization
diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl
index b769783cd..59661b84c 100644
--- a/lib/KernelInterface/src/host.jl
+++ b/lib/KernelInterface/src/host.jl
@@ -267,7 +267,9 @@ width [`sub_group_size`](@ref). See the manual for what KernelInterface guarante
 how work-groups are divided into sub-groups; a backend that can't ensure that reports
 `false`.
 
-Which types [`shfl_down`](@ref) supports is queried separately with [`supports_shuffle`](@ref).
+Sub-group support includes the votes [`sub_group_any`](@ref), [`sub_group_all`](@ref) and,
+for sub-groups of at most 64 work-items, [`sub_group_ballot`](@ref). Which types the
+shuffles support is queried separately with [`supports_shuffle`](@ref).
 
 !!! note
     Backend implementations **must** implement this function if they support sub-groups.
@@ -278,13 +280,21 @@ supports_subgroups(::Backend) = false
 """
     supports_shuffle(::Backend, ::Type{T})::Bool
 
-Whether kernels on the active device support [`shfl_down`](@ref) for values of type `T`.
+Whether kernels on the active device support the shuffles [`shfl`](@ref),
+[`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`.
+
+`isbits` structs and tuples are supported if all of their fields are.
 
 !!! note
-    Backend implementations **must** implement this function for the types they support.
-    The fallback returns `false`.
-"""
-supports_shuffle(::Backend, ::Type) = false
+    Backend implementations **must** implement this function for the primitive types they
+    support, with a signature that only matches those, e.g.
+    `supports_shuffle(::NewBackend, ::Type{<:Union{Int32, Float32}})`. The fallback returns
+    `false` for other primitive types, and checks the fields of other types.
+"""
+function supports_shuffle(backend::Backend, ::Type{T}) where {T}
+    (isbitstype(T) && !isprimitivetype(T)) || return false
+    return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T))
+end
 
 """
     allocate(::Backend, Type, dims...; unified=false)::AbstractArray
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 210b77123..ac91b803d 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -244,6 +244,65 @@ function shfl_down_lanes_kernel(out, ::Type{T}, offset) where {T}
     return
 end
 
+# Every lane reads the lane `shift` further, wrapping around: the rotation of a tile.
+function shfl_rotate_kernel(out, a::AbstractArray{T}, shift) where {T}
+    lane = KI.get_sub_group_local_id()
+    width = KI.get_sub_group_size()
+    val = @inbounds a[lane]
+    @inbounds out[lane] = KI.shfl(val, mod1(lane + shift, width))
+    return
+end
+
+function shfl_up_lanes_kernel(out, ::Type{T}, offset) where {T}
+    lane = KI.get_sub_group_local_id()
+    shuffled = KI.shfl_up(T(lane), offset)
+    @inbounds out[lane] = shuffled
+    return
+end
+
+# An all-reduce with a butterfly: every lane gets the sum of the sub-group. `N` is the
+# width, and the work-group is one sub-group.
+function shfl_xor_sum_kernel(out, a, ::Val{N}) where {N}
+    lane = KI.get_sub_group_local_id()
+    val = @inbounds a[lane]
+    mask = N >> 1
+    while mask > 0
+        val += KI.shfl_xor(val, mask)
+        mask >>= 1
+    end
+    @inbounds out[lane] = val
+    return
+end
+
+struct ShuffleStruct
+    a::Float32
+    b::Int64
+    c::NTuple{3, Int32}
+end
+
+function shfl_struct_kernel(out, a)
+    lane = KI.get_sub_group_local_id()
+    width = KI.get_sub_group_size()
+    val = @inbounds a[lane]
+    @inbounds out[lane] = KI.shfl(val, mod1(lane + 1, width))
+    return
+end
+
+function vote_kernel(out, pred)
+    lane = KI.get_sub_group_local_id()
+    p = @inbounds pred[lane]
+    any = KI.sub_group_any(p)
+    all = KI.sub_group_all(p)
+    ballot = KI.sub_group_ballot(p)
+    @inbounds begin
+        out[lane, 1] = any
+        out[lane, 2] = all
+        out[lane, 3] = ballot
+        out[lane, 4] = KI.get_max_sub_group_size()
+    end
+    return
+end
+
 # Every lane writes local and global memory, and reads another lane's write after a
 # sub-group barrier. `N` is the sub-group width; the work-group is one sub-group.
 function sub_group_barrier_kernel(scratch, out, ::Val{N}) where {N}
@@ -696,6 +755,71 @@ function interface_testsuite(backend::KI.Backend, AT)
             end
         end
 
+        @testset "shuffles" begin
+            candidates = (Int32, Int64, UInt32, Float32, Float64)
+            types = filter(T -> KI.supports_shuffle(backend, T), candidates)
+            @testset "$T" for T in types
+                a = T.(rand(1:100, sg_size))
+                @testset "shfl, shift $shift" for shift in unique((0, 1, sg_size - 1))
+                    out = AT(zeros(T, sg_size))
+                    KI.@launch backend workgroupsize = sg_size shfl_rotate_kernel(out, AT(a), shift)
+                    KI.synchronize(backend)
+                    @test Array(out) == circshift(a, -shift)
+                end
+
+                @testset "shfl_up, offset $offset" for offset in unique((1, 3, sg_size ÷ 2))
+                    1 <= offset < sg_size || continue
+                    out = AT(zeros(T, sg_size))
+                    KI.@launch backend workgroupsize = sg_size shfl_up_lanes_kernel(out, T, offset)
+                    KI.synchronize(backend)
+                    @test Array(out)[(offset + 1):end] == T.(1:(sg_size - offset))
+                end
+
+                if ispow2(sg_size)
+                    out = AT(zeros(T, sg_size))
+                    KI.@launch backend workgroupsize = sg_size shfl_xor_sum_kernel(out, AT(a), Val(sg_size))
+                    KI.synchronize(backend)
+                    @test all(==(sum(a)), Array(out))
+                end
+            end
+
+            @testset "structs" begin
+                T = ShuffleStruct
+                @test KI.supports_shuffle(backend, T) ==
+                    all(S -> KI.supports_shuffle(backend, S), (Float32, Int64, Int32))
+                @test !KI.supports_shuffle(backend, Ref{Int})
+                if KI.supports_shuffle(backend, T)
+                    a = [T(i, -i, (i, 2i, 3i)) for i in 1:sg_size]
+                    out = AT(fill(T(0, 0, (0, 0, 0)), sg_size))
+                    KI.@launch backend workgroupsize = sg_size shfl_struct_kernel(out, AT(a))
+                    KI.synchronize(backend)
+                    @test Array(out) == circshift(a, -1)
+                end
+            end
+        end
+
+        @testset "votes" begin
+            patterns = (
+                "none" => falses(sg_size),
+                "all" => trues(sg_size),
+                "some" => [i % 3 == 1 for i in 1:sg_size],
+                "last" => [i == sg_size for i in 1:sg_size],
+            )
+            @testset "$name" for (name, pred) in patterns
+                out = AT(zeros(UInt64, sg_size, 4))
+                KI.@launch backend workgroupsize = sg_size vote_kernel(out, AT(collect(pred)))
+                KI.synchronize(backend)
+                out = Array(out)
+                @test all(==(any(pred)), out[:, 1])
+                @test all(==(all(pred)), out[:, 2])
+                @test all(==(sg_size), out[:, 4])
+                if sg_size <= 64
+                    mask = reduce(|, (UInt64(1) << (i - 1) for i in 1:sg_size if pred[i]); init = UInt64(0))
+                    @test all(==(mask), out[:, 3])
+                end
+            end
+        end
+
         @testset "shfl_down" begin
             candidates = (
                 Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64,
@@ -751,6 +875,8 @@ function contract_testsuite(backend::KI.Backend, AT)
     @test hasmethod(KI.max_num_groups, Tuple{B})
     if KI.supports_subgroups(backend)
         @test hasmethod(KI.sub_group_size, Tuple{B})
+        # the device functions are overlays, so they can't be checked here; the sub-group
+        # testsuite runs them
     end
     return
 end
diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl
index a8e5b8def..3615cb8af 100644
--- a/lib/KernelInterface/test/runtests.jl
+++ b/lib/KernelInterface/test/runtests.jl
@@ -23,13 +23,16 @@ end
     @test !occursin("[sources]", toml)
 end
 
+struct ShuffleBackend <: KI.Backend end
+KI.supports_shuffle(::ShuffleBackend, ::Type{Int32}) = true
+
 # NOTE: this runs before the mock backend below defines methods on `argconvert`
 # and `kernel_function`.
 @testset "interface stubs" begin
     # These have no fallback on purpose: a backend that forgets to `@device_override`
     # them should get a MethodError rather than silently wrong behaviour.
     stubs = [
-        KI.shfl_down,
+        KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot,
         KI.max_work_group_size, KI.max_work_group_dims, KI.max_num_groups,
         KI.sub_group_size, KI.argconvert, KI.kernel_function, KI.launch,
         # Host-side stubs: required backend methods with no sensible fallback.
@@ -39,6 +42,19 @@ end
         @test isempty(methods(stub))
     end
 
+    # The shuffles only have the fallback that shuffles structs field by field, which
+    # doesn't handle the primitive types a backend has to implement.
+    for shfl in [KI.shfl, KI.shfl_down, KI.shfl_up, KI.shfl_xor]
+        @test_throws ArgumentError shfl(1.0f0, 1)
+        @test_throws ArgumentError shfl((1.0f0, 2), 1)
+        @test_throws ArgumentError shfl(Ref(1), 1)
+    end
+    @test !KI.supports_shuffle(ShuffleBackend(), Float32)
+    @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Float32, Int})
+    @test KI.supports_shuffle(ShuffleBackend(), Int32)
+    @test KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, NTuple{2, Int32}})
+    @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, Float32})
+
     # The primitive queries take an element type; only the zero-argument form has a
     # (forwarding) method, and it must reach the typed stub rather than recurse.
     primitives = [
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index 4d5df6726..c9927e632 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -174,9 +174,11 @@ function KI.multiprocessor_count(::POCLBackend)::Int
 end
 
 KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0
-function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T}
+# the types `sub_group_shuffle` supports; other types are shuffled field by field
+const ShuffleTypes = Union{SPIRVIntrinsics.gentypes...}
+
+function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T <: ShuffleTypes}
     KI.supports_subgroups(backend) || return false
-    T in SPIRVIntrinsics.gentypes || return false
     T === Float64 && return "cl_khr_fp64" in device().extensions
     T === Float16 && return "cl_khr_fp16" in device().extensions
     return true
@@ -247,8 +249,25 @@ end
     sub_group_barrier(POCL.LOCAL_MEM_FENCE | POCL.GLOBAL_MEM_FENCE)
 end
 
-@device_override function KI.shfl_down(val::T, offset::Integer) where {T}
-    sub_group_shuffle(val, get_sub_group_local_id() + offset)
+@device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} =
+    POCL.shuffle(val, lane)
+
+@device_override KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} =
+    POCL.shuffle(val, get_sub_group_local_id() + offset)
+
+@device_override KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} =
+    POCL.shuffle(val, get_sub_group_local_id() - offset)
+
+@device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} =
+    sub_group_shuffle_xor(val, mask % UInt32)
+
+@device_override KI.sub_group_any(pred::Bool) = POCL.sub_group_any(pred)
+
+@device_override KI.sub_group_all(pred::Bool) = POCL.sub_group_all(pred)
+
+@device_override function KI.sub_group_ballot(pred::Bool)
+    mask = POCL.sub_group_ballot(pred)
+    return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
 @device_override @inline function KI._print(args...)
diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl
index 62ff16027..464cecd2a 100644
--- a/src/pocl/compiler/compilation.jl
+++ b/src/pocl/compiler/compilation.jl
@@ -47,6 +47,7 @@ GPUCompiler.isintrinsic(job::OpenCLCompilerJob, fn::String) =
     job, fn
 ) ||
     in(fn, known_intrinsics) ||
+    in(fn, subgroup_intrinsics) ||
     contains(fn, "__spirv_")
 
 GPUCompiler.kernel_state_type(::OpenCLCompilerJob) = KernelState
@@ -64,6 +65,20 @@ function GPUCompiler.finish_module!(
     sg_size = job.config.params.sub_group_size
     if sg_size !== nothing
         entry.metadata["intel_reqd_sub_group_size"] = MDNode([ConstantInt(Int32(sg_size))])
+
+        # the sub-group width is fixed, so make `get_max_sub_group_size` a constant, as
+        # KernelInterface requires (this runs before optimization)
+        gvs = LLVM.globals(mod)
+        if haskey(gvs, "__spirv_BuiltInSubgroupMaxSize")
+            gv = gvs["__spirv_BuiltInSubgroupMaxSize"]
+            for use in collect(LLVM.uses(gv))
+                load = LLVM.user(use)
+                load isa LLVM.LoadInst || continue
+                LLVM.replace_uses!(load, ConstantInt(LLVM.value_type(load), sg_size))
+                LLVM.erase!(load)
+            end
+            isempty(LLVM.uses(gv)) && LLVM.erase!(gv)
+        end
     end
 
     # if this kernel uses our RNG, we should prime the shared state.
diff --git a/src/pocl/device/subgroups.jl b/src/pocl/device/subgroups.jl
new file mode 100644
index 000000000..07a2947bb
--- /dev/null
+++ b/src/pocl/device/subgroups.jl
@@ -0,0 +1,25 @@
+# Sub-group operations that SPIRVIntrinsics doesn't wrap (the way we need them).
+
+# `sub_group_shuffle`, with the lane passed modulo `UInt32`, so that an out-of-range lane
+# gives an unspecified value rather than an `InexactError`
+for T in SPIRVIntrinsics.gentypes
+    @eval @device_function shuffle(x::$T, lane::Integer) =
+        @builtin_ccall(
+        "__spirv_GroupNonUniformShuffle", $T, (UInt32, $T, UInt32),
+        UInt32(Scope.Subgroup), x, (lane - 1) % UInt32
+    )
+end
+
+# Votes, from `cl_khr_subgroups` and `cl_khr_subgroup_ballot`. The SPIR-V back-end lowers
+# these OpenCL built-ins, which have to be listed in `subgroup_intrinsics`.
+const subgroup_intrinsics = ["_Z13sub_group_anyi", "_Z13sub_group_alli", "_Z16sub_group_balloti"]
+
+@device_function sub_group_any(pred::Bool) =
+    ccall("extern _Z13sub_group_anyi", llvmcall, Int32, (Int32,), pred) != Int32(0)
+
+@device_function sub_group_all(pred::Bool) =
+    ccall("extern _Z13sub_group_alli", llvmcall, Int32, (Int32,), pred) != Int32(0)
+
+# bit `i` of the result is set for the lane with (0-based) id `i`
+@device_function sub_group_ballot(pred::Bool) =
+    ccall("extern _Z16sub_group_balloti", llvmcall, NTuple{4, VecElement{UInt32}}, (Int32,), pred)
diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl
index 17e47e296..c4b7f85df 100644
--- a/src/pocl/pocl.jl
+++ b/src/pocl/pocl.jl
@@ -109,6 +109,7 @@ include("device/array.jl")
 include("device/quirks.jl")
 include("device/runtime.jl")
 include("device/random.jl")
+include("device/subgroups.jl")
 
 include("compiler/compilation.jl")
 include("compiler/execution.jl")
diff --git a/test/runtests.jl b/test/runtests.jl
index 6fb138483..746747707 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -129,6 +129,22 @@ if "cl_khr_fp16" in POCL.device().extensions
     end
 end
 
+# KernelInterface requires the sub-group width to be a constant of the generated code
+function max_sub_group_size_kernel(out)
+    out[1] = KernelAbstractions.KernelInterface.get_max_sub_group_size()
+    return
+end
+@testset "POCL constant sub-group width" begin
+    out = zeros(Int, 1)
+    ir = sprint() do io
+        @device_code_llvm io = io debuginfo = :none @opencl max_sub_group_size_kernel(out)
+    end
+    width = POCL.device_limits().sub_group_size
+    @test out[1] == width
+    @test !occursin("BuiltInSubgroupMaxSize", ir)
+    @test occursin("store i64 $width", ir)
+end
+
 # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with
 # many arguments allocates unless every layer passes them on as a tuple
 @testset "POCL launch with many arguments" begin

From 6b3a360b5695eff27ded499a0a43a83a1c950066 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 08:44:23 +0200
Subject: [PATCH 02/18] KernelInterface: unroll nested fields in the struct
 shuffle fallback

On Julia 1.10, inference gives up on the recursive call of `shfl_fields` through
the shuffle of a nested field (e.g. a tuple in a struct), leaving a dynamic
invocation in the kernel. Generate the shuffles of all primitive fields
directly instead.

Assisted-by: Claude Code (Opus 5.5)
---
 lib/KernelInterface/src/device.jl | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index f106691e0..58d687345 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -357,7 +357,16 @@ the supported types.
 """
 @inline shfl_xor(val, mask::Integer) = shfl_fields(x -> shfl_xor(x, mask), val)
 
-# Shuffle a value that the backend doesn't support directly field by field.
+# The expression that shuffles `ex::S` field by field, calling `f` on the primitive fields
+function shfl_fields_expr(S, ex)
+    isprimitivetype(S) && return :(f($ex))
+    fields = (shfl_fields_expr(fieldtype(S, i), :(getfield($ex, $i))) for i in 1:fieldcount(S))
+    return Expr(:new, S, fields...)
+end
+
+# Shuffle a value that the backend doesn't support directly field by field. Nested fields are
+# unrolled here, rather than shuffled with a recursive call, which inference gives up on
+# (on Julia 1.10), so that `f` is only called on the primitive types.
 @inline @generated function shfl_fields(f, val::T) where {T}
     if !isbitstype(T) || isprimitivetype(T)
         return :(
@@ -368,8 +377,7 @@ the supported types.
             )
         )
     end
-    fields = (:(f(getfield(val, $i))) for i in 1:fieldcount(T))
-    return Expr(:new, T, fields...)
+    return shfl_fields_expr(T, :val)
 end
 
 """

From 8e3de8937e6d443d6745c3cf07ad650717830372 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 09:13:21 +0200
Subject: [PATCH 03/18] POCL: use the sub-group intrinsics from SPIRVIntrinsics

Replace the local workarounds with the votes and unchecked shuffle lanes from
JuliaGPU/OpenCL.jl#526, taken from its branch until it is released: through
`[sources]`, and explicitly where that doesn't apply (Julia 1.10 on CI, and the
Buildkite jobs, whose OpenCL job developed SPIRVIntrinsics from OpenCL.jl's
ka-0.10 branch).

Assisted-by: Claude Code (Opus 5.5)
---
 .buildkite/pipeline.yml          | 23 +++++++++++++++++------
 .github/workflows/ci.yml         |  5 ++++-
 Project.toml                     |  2 ++
 src/pocl/backend.jl              | 14 +++++++-------
 src/pocl/compiler/compilation.jl |  1 -
 src/pocl/device/subgroups.jl     | 25 -------------------------
 src/pocl/pocl.jl                 |  1 -
 7 files changed, 30 insertions(+), 41 deletions(-)
 delete mode 100644 src/pocl/device/subgroups.jl

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 9a7c65d96..189920191 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -24,7 +24,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing CUDA")
                 using Pkg
                 url="https://github.com/JuliaGPU/CUDA.jl"
@@ -64,7 +66,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing Metal")
                 using Pkg
                 Pkg.add([(; url="https://github.com/JuliaGPU/Metal.jl", rev="ka-0.10")])' || exit 3
@@ -100,7 +104,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing oneAPI")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -139,7 +145,9 @@ steps:
                 using Pkg
                 println("--- :julia: Instantiating project")
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing AMDGPU")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -178,11 +186,14 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing OpenCL")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="ka-0.10")
-                Pkg.develop(; name="SPIRVIntrinsics")' || exit 3
+                # TODO: develop the ka-0.10 copy again once JuliaGPU/OpenCL.jl#526 is merged
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index bdccba98d..47938dfd7 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -56,9 +56,12 @@ jobs:
       # applies to the test run only, not to `Pkg.develop`. Do this before building, which
       # already resolves the environment (and fails if this repo requires an unregistered
       # version of KernelInterface).
+      #
+      # The same applies to the SPIRVIntrinsics branch in `[sources]`.
+      # TODO: remove that once JuliaGPU/OpenCL.jl#526 is released
       - name: Dev KernelInterface
         shell: bash
-        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")'
+        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface"); Pkg.add(; url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")'
       # Nightly tracks the next Julia release; don't fail CI on it
       - uses: julia-actions/julia-buildpkg@v1
         continue-on-error: ${{ matrix.version == 'nightly' }}
diff --git a/Project.toml b/Project.toml
index b633c85bd..ac1763297 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,6 +31,8 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [sources]
 KernelInterface = {path = "lib/KernelInterface"}
+# TODO: remove once JuliaGPU/OpenCL.jl#526 is released
+SPIRVIntrinsics = {url = "https://github.com/JuliaGPU/OpenCL.jl", rev = "vc/subgroup-votes", subdir = "lib/intrinsics"}
 
 [extensions]
 LinearAlgebraExt = "LinearAlgebra"
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index c9927e632..c39857def 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -250,23 +250,23 @@ end
 end
 
 @device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} =
-    POCL.shuffle(val, lane)
+    sub_group_shuffle(val, lane)
 
 @device_override KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} =
-    POCL.shuffle(val, get_sub_group_local_id() + offset)
+    sub_group_shuffle(val, get_sub_group_local_id() + offset)
 
 @device_override KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} =
-    POCL.shuffle(val, get_sub_group_local_id() - offset)
+    sub_group_shuffle(val, get_sub_group_local_id() - offset)
 
 @device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} =
-    sub_group_shuffle_xor(val, mask % UInt32)
+    sub_group_shuffle_xor(val, mask)
 
-@device_override KI.sub_group_any(pred::Bool) = POCL.sub_group_any(pred)
+@device_override KI.sub_group_any(pred::Bool) = SPIRVIntrinsics.sub_group_any(pred)
 
-@device_override KI.sub_group_all(pred::Bool) = POCL.sub_group_all(pred)
+@device_override KI.sub_group_all(pred::Bool) = SPIRVIntrinsics.sub_group_all(pred)
 
 @device_override function KI.sub_group_ballot(pred::Bool)
-    mask = POCL.sub_group_ballot(pred)
+    mask = SPIRVIntrinsics.sub_group_ballot(pred)
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl
index 464cecd2a..90bb400fc 100644
--- a/src/pocl/compiler/compilation.jl
+++ b/src/pocl/compiler/compilation.jl
@@ -47,7 +47,6 @@ GPUCompiler.isintrinsic(job::OpenCLCompilerJob, fn::String) =
     job, fn
 ) ||
     in(fn, known_intrinsics) ||
-    in(fn, subgroup_intrinsics) ||
     contains(fn, "__spirv_")
 
 GPUCompiler.kernel_state_type(::OpenCLCompilerJob) = KernelState
diff --git a/src/pocl/device/subgroups.jl b/src/pocl/device/subgroups.jl
deleted file mode 100644
index 07a2947bb..000000000
--- a/src/pocl/device/subgroups.jl
+++ /dev/null
@@ -1,25 +0,0 @@
-# Sub-group operations that SPIRVIntrinsics doesn't wrap (the way we need them).
-
-# `sub_group_shuffle`, with the lane passed modulo `UInt32`, so that an out-of-range lane
-# gives an unspecified value rather than an `InexactError`
-for T in SPIRVIntrinsics.gentypes
-    @eval @device_function shuffle(x::$T, lane::Integer) =
-        @builtin_ccall(
-        "__spirv_GroupNonUniformShuffle", $T, (UInt32, $T, UInt32),
-        UInt32(Scope.Subgroup), x, (lane - 1) % UInt32
-    )
-end
-
-# Votes, from `cl_khr_subgroups` and `cl_khr_subgroup_ballot`. The SPIR-V back-end lowers
-# these OpenCL built-ins, which have to be listed in `subgroup_intrinsics`.
-const subgroup_intrinsics = ["_Z13sub_group_anyi", "_Z13sub_group_alli", "_Z16sub_group_balloti"]
-
-@device_function sub_group_any(pred::Bool) =
-    ccall("extern _Z13sub_group_anyi", llvmcall, Int32, (Int32,), pred) != Int32(0)
-
-@device_function sub_group_all(pred::Bool) =
-    ccall("extern _Z13sub_group_alli", llvmcall, Int32, (Int32,), pred) != Int32(0)
-
-# bit `i` of the result is set for the lane with (0-based) id `i`
-@device_function sub_group_ballot(pred::Bool) =
-    ccall("extern _Z16sub_group_balloti", llvmcall, NTuple{4, VecElement{UInt32}}, (Int32,), pred)
diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl
index c4b7f85df..17e47e296 100644
--- a/src/pocl/pocl.jl
+++ b/src/pocl/pocl.jl
@@ -109,7 +109,6 @@ include("device/array.jl")
 include("device/quirks.jl")
 include("device/runtime.jl")
 include("device/random.jl")
-include("device/subgroups.jl")
 
 include("compiler/compilation.jl")
 include("compiler/execution.jl")

From c93c3667961756bd6f7c532202e29b60dc903054 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 10:27:05 +0200
Subject: [PATCH 04/18] KernelInterface: shuffles of any bits type, with a
 width; match, reduce and scan

Fill the gaps that a survey of the packages using warp operations (KomaMRI,
KernelIntrinsics/KernelForge, AcceleratedKernels#93, ParallelStencil, ClimaCore,
...) showed:

- Primitive types that a backend doesn't support natively (e.g. `Bool`,
  `Char`, or 64-bit types on Metal) are shuffled as `UInt32` words, which
  backends now have to support. Structs keep being shuffled field by field.
- Shuffles within segments of `width` lanes (`shfl(val, lane, width)` etc.),
  with CUDA's semantics, built on `shfl`.
- `sub_group_match_any(val)`, the mask of the lanes with the same value, with a
  fallback built on `shfl` and `sub_group_ballot`.
- `sub_group_reduce(op, val)` and `sub_group_scan(op, val)` with fallbacks
  built on the shuffles, which backends can implement with native operations.
- Document how partial sub-groups behave.

The new tests are in a function of their own: as part of `interface_testsuite`,
compiling the host code crashed LLVM.

Assisted-by: Claude Code (Opus 5.5)
---
 docs/src/kernelinterface.md                |  12 +-
 lib/KernelInterface/src/KernelInterface.jl |   3 +-
 lib/KernelInterface/src/device.jl          | 225 +++++++++++++++++++--
 lib/KernelInterface/src/host.jl            |  10 +-
 lib/KernelInterface/test/interface.jl      | 186 +++++++++++++++++
 lib/KernelInterface/test/runtests.jl       |   7 +
 6 files changed, 417 insertions(+), 26 deletions(-)

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index c6f7cfb08..274a742cb 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
-| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for primitive `T` | |
+| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
 
 Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones),
 the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch),
@@ -163,6 +163,12 @@ report sub-group support have to ensure:
   `1:get_sub_group_size()`;
 - a 1-D work-group of at most `sub_group_size(backend)` work-items is a single sub-group.
 
+A sub-group is partial when it has fewer work-items than the sub-group width: the lanes
+`get_sub_group_size()+1:get_max_sub_group_size()` have no work-item. Shuffles from those lanes
+give unspecified values, and the votes, [`sub_group_match_any`](@ref),
+[`sub_group_reduce`](@ref) and [`sub_group_scan`](@ref) only take the work-items of the
+sub-group into account.
+
 In particular, [`get_num_sub_groups`](@ref) can be larger than
 `cld(prod(get_local_size()), get_max_sub_group_size())`. Storage for a value per sub-group
 has to be sized for up to one sub-group per work-item, and code combining those values has
@@ -196,9 +202,13 @@ shfl
 shfl_down
 shfl_up
 shfl_xor
+shfl(::Any, ::Integer, ::Integer)
 sub_group_any
 sub_group_all
 sub_group_ballot
+sub_group_match_any
+sub_group_reduce
+sub_group_scan
 ```
 
 ### Printing
diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl
index 03311eaec..61db3a339 100644
--- a/lib/KernelInterface/src/KernelInterface.jl
+++ b/lib/KernelInterface/src/KernelInterface.jl
@@ -32,7 +32,8 @@ include("host.jl")
             :get_sub_group_size, :get_max_sub_group_size, :get_num_sub_groups,
             :get_sub_group_id, :get_sub_group_local_id,
             :localmemory, :shfl, :shfl_down, :shfl_up, :shfl_xor,
-            :sub_group_any, :sub_group_all, :sub_group_ballot,
+            :sub_group_any, :sub_group_all, :sub_group_ballot, :sub_group_match_any,
+            :sub_group_reduce, :sub_group_scan,
             :barrier, :sub_group_barrier, :_print,
             # compilation and launch
             :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"),
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index 58d687345..59880daf4 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -132,6 +132,11 @@ end
 # doesn't change during the kernel's execution, and that a 1-D work-group of at most
 # `sub_group_size(backend)` work-items is a single sub-group. Backends that can't ensure
 # that don't report sub-group support. See the manual.
+#
+# In a partial sub-group, the lanes `get_sub_group_size()+1:get_max_sub_group_size()` have no
+# work-item: shuffles from them give unspecified values, and the votes, `sub_group_match_any`,
+# `sub_group_reduce` and `sub_group_scan` only take the work-items of the sub-group into
+# account.
 
 """
     get_sub_group_size([::Type{T}=Int])::T
@@ -264,8 +269,8 @@ localmemory(::Type{T}, ::Val) where {T} =
 ## communication
 
 # Shuffles exchange values between the work-items of a sub-group. Backends implement them for
-# the primitive types for which `supports_shuffle` returns `true`; the fallbacks below
-# shuffle other `isbits` types field by field.
+# the primitive types for which `supports_shuffle` returns `true`. The fallbacks below shuffle
+# other primitive types as `UInt32` words, and other `isbits` types field by field.
 
 """
     shfl(val::T, lane::Integer)::T
@@ -278,18 +283,21 @@ but they may read from different lanes.
 
 `shfl` exchanges values, not memory: it is not a memory fence.
 
-Types for which [`supports_shuffle`](@ref) returns `true` are supported, as well as `isbits`
-structs and tuples of such types, which are shuffled field by field.
+Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a
+backend supports natively, that includes other primitive types of 1, 2 or a multiple of 4
+bytes (e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, which are shuffled
+as `UInt32` words, and `isbits` structs and tuples of supported types, which are shuffled
+field by field.
 
 !!! note
-    Backend implementations **must** implement this for every primitive type `T` for which
-    [`supports_shuffle`](@ref) returns `true`, and only for those, so that other types reach
-    the fallback that shuffles structs field by field:
+    Backend implementations **must** implement this for the primitive types they support
+    natively, which have to include `UInt32`, and only for those, so that other types reach
+    the fallbacks:
     ```
     @device_override shfl(val::T, lane::Integer) where {T <: Union{...}}
     ```
 """
-@inline shfl(val, lane::Integer) = shfl_fields(x -> shfl(x, lane), val)
+@inline shfl(val, lane::Integer) = shfl_fallback(x -> shfl(x, lane), val)
 
 """
     shfl_down(val::T, offset::Integer)::T
@@ -310,7 +318,7 @@ the supported types.
     @device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}}
     ```
 """
-@inline shfl_down(val, offset::Integer) = shfl_fields(x -> shfl_down(x, offset), val)
+@inline shfl_down(val, offset::Integer) = shfl_fallback(x -> shfl_down(x, offset), val)
 
 """
     shfl_up(val::T, offset::Integer)::T
@@ -331,7 +339,7 @@ the supported types.
     @device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}}
     ```
 """
-@inline shfl_up(val, offset::Integer) = shfl_fields(x -> shfl_up(x, offset), val)
+@inline shfl_up(val, offset::Integer) = shfl_fallback(x -> shfl_up(x, offset), val)
 
 """
     shfl_xor(val::T, mask::Integer)::T
@@ -355,7 +363,45 @@ the supported types.
     @device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}}
     ```
 """
-@inline shfl_xor(val, mask::Integer) = shfl_fields(x -> shfl_xor(x, mask), val)
+@inline shfl_xor(val, mask::Integer) = shfl_fallback(x -> shfl_xor(x, mask), val)
+
+# Shuffle a value of a type that the backend doesn't support natively, with `f` shuffling a
+# value of a type it does support. The fallbacks are separate functions for primitive and for
+# other types, so that the fallback of a struct with a field the backend doesn't support
+# natively, e.g. an `Int64` on Metal, doesn't call itself, which inference gives up on (on
+# Julia 1.10).
+@inline function shfl_fallback(f, val::T) where {T}
+    return isprimitivetype(T) ? shfl_words(f, val) : shfl_fields(f, val)
+end
+
+shfl_unsupported(T) = throw(
+    ArgumentError(
+        "Shuffling values of type $T is not supported by this backend, see `supports_shuffle`"
+    )
+)
+
+# Whether a primitive type that a backend doesn't support natively can be shuffled as `UInt32`
+# words
+shuffle_as_words(T) = T !== UInt32 && sizeof(T) in (1, 2, 4, 8, 16)
+
+# The unsigned integer type of the size of a primitive type `T`
+const word_types = Dict(1 => UInt8, 2 => UInt16, 4 => UInt32, 8 => UInt64, 16 => UInt128)
+
+# Shuffle a primitive value as `UInt32` words: smaller values are zero-extended, larger ones
+# split into words.
+@inline @generated function shfl_words(f, val::T) where {T}
+    shuffle_as_words(T) || return :(shfl_unsupported($T))
+    U = word_types[sizeof(T)]
+    if sizeof(T) <= 4
+        return :(reinterpret($T, f(reinterpret($U, val) % UInt32) % $U))
+    end
+    n = sizeof(T) ÷ 4
+    words = (:((f((bits >> $(32 * (i - 1))) % UInt32) % $U) << $(32 * (i - 1))) for i in 1:n)
+    return quote
+        bits = reinterpret($U, val)
+        return reinterpret($T, |($(words...)))
+    end
+end
 
 # The expression that shuffles `ex::S` field by field, calling `f` on the primitive fields
 function shfl_fields_expr(S, ex)
@@ -368,18 +414,64 @@ end
 # unrolled here, rather than shuffled with a recursive call, which inference gives up on
 # (on Julia 1.10), so that `f` is only called on the primitive types.
 @inline @generated function shfl_fields(f, val::T) where {T}
-    if !isbitstype(T) || isprimitivetype(T)
-        return :(
-            throw(
-                ArgumentError(
-                    $("Shuffling values of type $T is not supported by this backend, see `supports_shuffle`")
-                )
-            )
-        )
-    end
+    isbitstype(T) || return :(shfl_unsupported($T))
     return shfl_fields_expr(T, :val)
 end
 
+# The shuffles within segments of `width` lanes, implemented with `shfl` from a lane.
+
+"""
+    shfl(val::T, lane::Integer, width::Integer)::T
+    shfl_down(val::T, offset::Integer, width::Integer)::T
+    shfl_up(val::T, offset::Integer, width::Integer)::T
+    shfl_xor(val::T, mask::Integer, width::Integer)::T
+
+Shuffles within segments of `width` consecutive lanes of the sub-group, as if each segment
+were a sub-group of its own: `lane` is the lane within the segment (between 1 and `width`, and
+taken modulo `width` otherwise), and `shfl_down`, `shfl_up` and `shfl_xor` read from lanes of
+the same segment. Where these would read from outside of the segment, they return `val` of the
+work-item itself (rather than an unspecified value, as without `width`), like CUDA's
+shuffles with a `width`. Reading from a lane of the segment that has no work-item (in a
+partial sub-group) gives an unspecified value.
+
+`width` has to be a power of two of at most the sub-group width
+[`get_max_sub_group_size`](@ref), and the same for all work-items of the sub-group.
+
+!!! note
+    Backends **may** implement these, e.g. if they have native shuffles with a width. The
+    fallbacks use [`shfl`](@ref) from a lane.
+"""
+@inline function shfl(val, lane::Integer, width::Integer)
+    l0 = get_sub_group_local_id(Int32) - Int32(1)
+    w = width % Int32
+    base = l0 & ~(w - Int32(1))
+    return shfl(val, base + ((lane % Int32 - Int32(1)) & (w - Int32(1))) + Int32(1))
+end
+
+@inline function shfl_down(val, offset::Integer, width::Integer)
+    l0 = get_sub_group_local_id(Int32) - Int32(1)
+    w = width % Int32
+    d = offset % Int32
+    src = ifelse((l0 & (w - Int32(1))) + d < w, l0 + d, l0)
+    return shfl(val, src + Int32(1))
+end
+
+@inline function shfl_up(val, offset::Integer, width::Integer)
+    l0 = get_sub_group_local_id(Int32) - Int32(1)
+    w = width % Int32
+    d = offset % Int32
+    src = ifelse((l0 & (w - Int32(1))) >= d, l0 - d, l0)
+    return shfl(val, src + Int32(1))
+end
+
+@inline function shfl_xor(val, mask::Integer, width::Integer)
+    l0 = get_sub_group_local_id(Int32) - Int32(1)
+    w = width % Int32
+    x = l0 ⊻ (mask % Int32)
+    src = ifelse((x & ~(w - Int32(1))) == (l0 & ~(w - Int32(1))), x, l0)
+    return shfl(val, src + Int32(1))
+end
+
 """
     sub_group_any(pred::Bool)::Bool
 
@@ -435,6 +527,99 @@ divergent branch). Only sub-groups of at most 64 work-items are supported.
 """
 function sub_group_ballot end
 
+"""
+    sub_group_match_any(val)::UInt64
+
+A mask of the work-items of the sub-group whose `val` is the same as the one of this
+work-item (compared bitwise, with `===`), in the format of [`sub_group_ballot`](@ref). Work-items
+with different values get different masks: e.g. `trailing_zeros(mask) + 1` is the lane of
+the first work-item with the same value, and `count_ones(mask)` the number of them.
+
+All work-items of the sub-group have to execute `sub_group_match_any` together (not in a
+divergent branch). Values of the types that the shuffles support are supported, see
+[`supports_shuffle`](@ref).
+
+!!! note
+    Backends **may** implement this, e.g. with CUDA's `match.any.sync`. The fallback finds
+    the groups of equal values one by one with [`shfl`](@ref) and [`sub_group_ballot`](@ref),
+    so it takes as many steps as there are distinct values.
+"""
+@inline function sub_group_match_any(val)
+    remaining = sub_group_ballot(true)
+    mask = zero(UInt64)
+    # uniform: every step removes the group of the lowest remaining lane from `remaining`
+    while remaining != zero(UInt64)
+        leader = trailing_zeros(remaining) + 1
+        same = val === shfl(val, leader)
+        group = sub_group_ballot(same)
+        mask = ifelse(same, group, mask)
+        remaining &= ~group
+    end
+    return mask
+end
+
+"""
+    sub_group_reduce(op, val::T)::T
+
+Reduce `val` over the work-items of the sub-group with the associative binary operator `op`,
+in the order of the lanes. All work-items of the sub-group get the result.
+
+All work-items of the sub-group have to execute `sub_group_reduce` together (not in a
+divergent branch). Values of the types that the shuffles support are supported, see
+[`supports_shuffle`](@ref).
+
+!!! note
+    Backends **may** implement this for operators and types with a native reduction (e.g.
+    `+` on `Float32`), dispatching on `typeof(op)`. The fallback combines ranges of doubling
+    length with [`shfl_down`](@ref), and broadcasts the result of the first lane with
+    [`shfl`](@ref).
+"""
+@inline function sub_group_reduce(op, val)
+    lane = get_sub_group_local_id(Int32)
+    sgsize = get_sub_group_size(Int32)
+    offset = Int32(1)
+    while offset < sgsize
+        other = shfl_down(val, offset)
+        # the result of shuffling from past the end of the sub-group is unspecified
+        if lane + offset <= sgsize
+            val = op(val, other)
+        end
+        offset <<= 1
+    end
+    return shfl(val, 1)
+end
+
+"""
+    sub_group_scan(op, val::T)::T
+
+The inclusive scan of `val` over the work-items of the sub-group with the associative binary
+operator `op`, in the order of the lanes: the work-item in lane `i` gets the reduction of the
+values of lanes `1` to `i`.
+
+All work-items of the sub-group have to execute `sub_group_scan` together (not in a divergent
+branch). Values of the types that the shuffles support are supported, see
+[`supports_shuffle`](@ref).
+
+!!! note
+    Backends **may** implement this for operators and types with a native scan (e.g. `+` on
+    `Float32`), dispatching on `typeof(op)`. The fallback is a Hillis-Steele scan with
+    [`shfl_up`](@ref).
+"""
+@inline function sub_group_scan(op, val)
+    lane = get_sub_group_local_id(Int32)
+    sgsize = get_sub_group_size(Int32)
+    offset = Int32(1)
+    while offset < sgsize
+        other = shfl_up(val, offset)
+        # the result of shuffling from before the start of the sub-group is unspecified
+        if lane > offset
+            val = op(other, val)
+        end
+        offset <<= 1
+    end
+    return val
+end
+
 
 ## synchronization
 
diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl
index 59661b84c..2952d1d3c 100644
--- a/lib/KernelInterface/src/host.jl
+++ b/lib/KernelInterface/src/host.jl
@@ -283,15 +283,17 @@ supports_subgroups(::Backend) = false
 Whether kernels on the active device support the shuffles [`shfl`](@ref),
 [`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`.
 
-`isbits` structs and tuples are supported if all of their fields are.
+Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and `isbits`
+structs and tuples if all of their fields are.
 
 !!! note
     Backend implementations **must** implement this function for the primitive types they
-    support, with a signature that only matches those, e.g.
-    `supports_shuffle(::NewBackend, ::Type{<:Union{Int32, Float32}})`. The fallback returns
-    `false` for other primitive types, and checks the fields of other types.
+    support natively, which have to include `UInt32`, with a signature that only matches
+    those, e.g. `supports_shuffle(::NewBackend, ::Type{<:Union{UInt32, Float32}})`. The
+    fallback handles other types.
 """
 function supports_shuffle(backend::Backend, ::Type{T}) where {T}
+    isprimitivetype(T) && return shuffle_as_words(T) && supports_shuffle(backend, UInt32)
     (isbitstype(T) && !isprimitivetype(T)) || return false
     return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T))
 end
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index ac91b803d..f648aee2a 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -303,6 +303,50 @@ function vote_kernel(out, pred)
     return
 end
 
+# Each work-item shuffles its value with `op` (`shfl`, `shfl_down`, `shfl_up`, `shfl_xor`) and
+# a `width`, recording its lane.
+function shfl_width_kernel(out, lanes, a, op, arg, width)
+    lane = KI.get_sub_group_local_id()
+    val = @inbounds a[lane]
+    @inbounds out[lane] = op(val, arg, width)
+    @inbounds lanes[lane] = lane
+    return
+end
+
+# `pred` determines which values are equal, `vals` gives the values
+function match_any_kernel(out, vals)
+    lane = KI.get_sub_group_local_id()
+    @inbounds out[lane] = KI.sub_group_match_any(vals[lane])
+    return
+end
+
+# a work-group of `length(a)` work-items, which is a single (possibly partial) sub-group
+function reduce_scan_kernel(red, scan, lanes, a, op)
+    i = KI.get_local_id().x
+    val = @inbounds a[i]
+    r = KI.sub_group_reduce(op, val)
+    s = KI.sub_group_scan(op, val)
+    @inbounds begin
+        red[i] = r
+        scan[i] = s
+        lanes[i] = KI.get_sub_group_local_id()
+    end
+    return
+end
+
+# primitive types that no backend supports natively, shuffled as `UInt32` words
+primitive type Bits64 64 end
+primitive type Bits16 16 end
+Bits64(x::Integer) = reinterpret(Bits64, x % UInt64)
+Bits16(x::Integer) = reinterpret(Bits16, x % UInt16)
+
+struct FallbackStruct
+    flag::Bool
+    c::Char
+    x::Bits64
+    limbs::NTuple{4, Float64}
+end
+
 # Every lane writes local and global memory, and reads another lane's write after a
 # sub-group barrier. `N` is the sub-group width; the work-group is one sub-group.
 function sub_group_barrier_kernel(scratch, out, ::Val{N}) where {N}
@@ -324,6 +368,146 @@ function captured_array_kernel(backend, AT, out)
     return kernel, WeakRef(a)
 end
 
+# an associative but not commutative operator: the composition of affine maps
+compose_affine(f, g) = (g[1] * f[1], g[1] * f[2] + g[2])
+# the (value, index) of the smallest value, the first one of equal values
+argmin_op(x, y) = ifelse(y[1] < x[1], y, x)
+
+# The sub-group operations that are built on the shuffles and votes. Separate functions, so
+# that `interface_testsuite` doesn't get too large to compile.
+
+# rotating values of type `T` by one lane, `f(i)` giving the value of lane `i`
+function shfl_type_testsuite(backend, AT, sg_size, ::Type{T}, f) where {T}
+    a = T[f(i) for i in 1:sg_size]
+    out = AT(fill(f(0), sg_size))
+    KI.@launch backend workgroupsize = sg_size shfl_rotate_kernel(out, AT(a), 1)
+    KI.synchronize(backend)
+    @test Array(out) == circshift(a, -1)
+    return
+end
+
+# `ref(lane, arg, width)` is the lane that `op(val, arg, width)` reads from
+function shfl_width_testsuite(backend, AT, sg_size, op, ref, width, arg)
+    a = Int32.(1:sg_size) .* Int32(10)
+    out = AT(zeros(Int32, sg_size))
+    lanes = AT(zeros(Int, sg_size))
+    KI.@launch backend workgroupsize = sg_size shfl_width_kernel(out, lanes, AT(a), op, arg, width)
+    KI.synchronize(backend)
+    out, lanes = Array(out), Array(lanes)
+    @test all(i -> out[i] == a[ref(lanes[i], arg, width)], 1:sg_size)
+    return
+end
+
+shfl_width_ref(l, arg, w) = (l - 1) ÷ w * w + mod1(arg, w)
+shfl_down_width_ref(l, arg, w) = (l - 1) % w + arg < w ? l + arg : l
+shfl_up_width_ref(l, arg, w) = (l - 1) % w >= arg ? l - arg : l
+shfl_xor_width_ref(l, arg, w) = (((l - 1) ⊻ arg) ÷ w == (l - 1) ÷ w) ? ((l - 1) ⊻ arg) + 1 : l
+
+function match_any_testsuite(backend, AT, sg_size, vals)
+    out = AT(zeros(UInt64, sg_size))
+    KI.@launch backend workgroupsize = sg_size match_any_kernel(out, AT(vals))
+    KI.synchronize(backend)
+    out = Array(out)
+    for i in 1:sg_size
+        expected = UInt64(0)
+        for j in 1:sg_size
+            vals[j] === vals[i] && (expected |= UInt64(1) << (j - 1))
+        end
+        @test out[i] == expected
+    end
+    return
+end
+
+# a work-group of `length(a)` work-items, i.e. a single sub-group, possibly partial
+function reduce_scan_testsuite(backend, AT, op, a)
+    n = length(a)
+    red = AT(similar(a))
+    scan = AT(similar(a))
+    lanes = AT(zeros(Int, n))
+    KI.@launch backend workgroupsize = n reduce_scan_kernel(red, scan, lanes, AT(a), op)
+    KI.synchronize(backend)
+    red, scan, lanes = Array(red), Array(scan), Array(lanes)
+    # in the order of the lanes
+    order = sortperm(lanes)
+    @test all(==(foldl(op, a[order])), red)
+    @test all(i -> scan[order[i]] == foldl(op, a[order[1:i]]), 1:n)
+    return
+end
+
+function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
+    @testset "shuffles of other types" begin
+        # primitive types that backends need not support natively are shuffled as words
+        if KI.supports_shuffle(backend, UInt32)
+            @test KI.supports_shuffle(backend, Bool)
+            @test KI.supports_shuffle(backend, Char)
+            @test KI.supports_shuffle(backend, Bits64)
+            @test KI.supports_shuffle(backend, Bits16)
+            @test KI.supports_shuffle(backend, FallbackStruct)
+        end
+        @test !KI.supports_shuffle(backend, Ref{Int})
+        @testset "Bool" begin
+            KI.supports_shuffle(backend, Bool) && shfl_type_testsuite(backend, AT, sg_size, Bool, isodd)
+        end
+        @testset "Char" begin
+            KI.supports_shuffle(backend, Char) &&
+                shfl_type_testsuite(backend, AT, sg_size, Char, i -> Char(0x0001F600 + i))
+        end
+        @testset "Bits64" begin
+            KI.supports_shuffle(backend, Bits64) &&
+                shfl_type_testsuite(backend, AT, sg_size, Bits64, i -> Bits64((UInt64(i) << 40) - i))
+        end
+        @testset "Bits16" begin
+            KI.supports_shuffle(backend, Bits16) &&
+                shfl_type_testsuite(backend, AT, sg_size, Bits16, i -> Bits16(0xa000 + i))
+        end
+        @testset "struct" begin
+            KI.supports_shuffle(backend, FallbackStruct) && shfl_type_testsuite(
+                backend, AT, sg_size, FallbackStruct,
+                i -> FallbackStruct(isodd(i), Char(64 + i), Bits64(-i), (i, -i, 1 / i, 2.0^i))
+            )
+        end
+    end
+
+    KI.supports_shuffle(backend, Int32) || return
+
+    @testset "shuffles with a width" begin
+        for w in (1, 2, 4, 8, 16, 32, 64)
+            (w <= sg_size && sg_size % w == 0) || continue
+            for arg in unique((1, 3, w - 1, w + 2))
+                @testset "width $w, $arg" begin
+                    shfl_width_testsuite(backend, AT, sg_size, KI.shfl, shfl_width_ref, w, arg)
+                    shfl_width_testsuite(backend, AT, sg_size, KI.shfl_down, shfl_down_width_ref, w, arg)
+                    shfl_width_testsuite(backend, AT, sg_size, KI.shfl_up, shfl_up_width_ref, w, arg)
+                    shfl_width_testsuite(backend, AT, sg_size, KI.shfl_xor, shfl_xor_width_ref, w, arg)
+                end
+            end
+        end
+    end
+
+    if sg_size <= 64
+        @testset "sub_group_match_any" begin
+            match_any_testsuite(backend, AT, sg_size, Int32[i % 3 for i in 1:sg_size])
+            # bitwise comparison, so NaN matches NaN, and -0.0 doesn't match 0.0
+            match_any_testsuite(
+                backend, AT, sg_size,
+                Float32[isodd(i) ? NaN32 : (i % 4 == 0 ? -0.0f0 : 0.0f0) for i in 1:sg_size]
+            )
+        end
+    end
+
+    for n in unique((sg_size, max(sg_size - 3, 1)))
+        @testset "sub_group_reduce and sub_group_scan, $n work-items" begin
+            reduce_scan_testsuite(backend, AT, +, Int32.(rand(1:100, n)))
+            reduce_scan_testsuite(
+                backend, AT, compose_affine,
+                [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n]
+            )
+            reduce_scan_testsuite(backend, AT, argmin_op, [(Float32(rand(1:20)), Int32(i)) for i in 1:n])
+        end
+    end
+    return
+end
+
 function interface_testsuite(backend::KI.Backend, AT)
     @testset "Launch parameters" begin
         # unequal group counts and sizes in every dimension, so that confusing them shows
@@ -798,6 +982,8 @@ function interface_testsuite(backend::KI.Backend, AT)
             end
         end
 
+        subgroup_communication_testsuite(backend, AT, sg_size)
+
         @testset "votes" begin
             patterns = (
                 "none" => falses(sg_size),
diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl
index 3615cb8af..762c0890c 100644
--- a/lib/KernelInterface/test/runtests.jl
+++ b/lib/KernelInterface/test/runtests.jl
@@ -25,6 +25,8 @@ end
 
 struct ShuffleBackend <: KI.Backend end
 KI.supports_shuffle(::ShuffleBackend, ::Type{Int32}) = true
+struct WordShuffleBackend <: KI.Backend end
+KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true
 
 # NOTE: this runs before the mock backend below defines methods on `argconvert`
 # and `kernel_function`.
@@ -49,6 +51,11 @@ KI.supports_shuffle(::ShuffleBackend, ::Type{Int32}) = true
         @test_throws ArgumentError shfl((1.0f0, 2), 1)
         @test_throws ArgumentError shfl(Ref(1), 1)
     end
+    # primitive types are shuffled as `UInt32` words, if the backend supports `UInt32`
+    @test KI.supports_shuffle(WordShuffleBackend(), Bool)
+    @test KI.supports_shuffle(WordShuffleBackend(), Float64)
+    @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Char, Int128})
+    @test !KI.supports_shuffle(WordShuffleBackend(), Ref{Int})
     @test !KI.supports_shuffle(ShuffleBackend(), Float32)
     @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Float32, Int})
     @test KI.supports_shuffle(ShuffleBackend(), Int32)

From f1dac4d9e976927e4c1232353ede1f69ece7d160 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 10:51:13 +0200
Subject: [PATCH 05/18] POCL: native sub-group reductions and scans

Implement `KI.sub_group_reduce` and `KI.sub_group_scan` with the collectives of
`cl_khr_subgroups` from SPIRVIntrinsics (JuliaGPU/OpenCL.jl#526): for `+` on
32- and 64-bit integers and floats, and `min`/`max` on integers. Floats keep the
fallback for `min` and `max`, as OpenCL treats NaN and the sign of zero
differently.

Test the operators and types that backends may implement natively, including a
NaN, and that POCL uses the native reduction.

Assisted-by: Claude Code (Opus 5.5)
---
 lib/KernelInterface/test/interface.jl | 12 ++++++++++--
 src/pocl/backend.jl                   | 19 +++++++++++++++++++
 test/runtests.jl                      | 17 +++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index f648aee2a..0c00ae2a4 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -429,8 +429,9 @@ function reduce_scan_testsuite(backend, AT, op, a)
     red, scan, lanes = Array(red), Array(scan), Array(lanes)
     # in the order of the lanes
     order = sortperm(lanes)
-    @test all(==(foldl(op, a[order])), red)
-    @test all(i -> scan[order[i]] == foldl(op, a[order[1:i]]), 1:n)
+    # `isequal`, as the results may be NaN
+    @test all(isequal(foldl(op, a[order])), red)
+    @test all(i -> isequal(scan[order[i]], foldl(op, a[order[1:i]])), 1:n)
     return
 end
 
@@ -498,6 +499,13 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     for n in unique((sg_size, max(sg_size - 3, 1)))
         @testset "sub_group_reduce and sub_group_scan, $n work-items" begin
             reduce_scan_testsuite(backend, AT, +, Int32.(rand(1:100, n)))
+            # operators and types that backends may implement natively
+            reduce_scan_testsuite(backend, AT, min, Int64.(rand(-100:100, n)))
+            reduce_scan_testsuite(backend, AT, max, UInt32.(rand(1:100, n)))
+            KI.supports_shuffle(backend, Float32) &&
+                reduce_scan_testsuite(backend, AT, +, Float32.(rand(1:100, n)))
+            KI.supports_shuffle(backend, Float32) &&
+                reduce_scan_testsuite(backend, AT, max, Float32[i == 2 ? NaN32 : rand(1:100) for i in 1:n])
             reduce_scan_testsuite(
                 backend, AT, compose_affine,
                 [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n]
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index c39857def..6b15f2d18 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -270,6 +270,25 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
+# native reductions and scans of `cl_khr_subgroups`. Only `+` for floats: OpenCL's `min` and
+# `max` treat NaN and the sign of zero differently from Julia's.
+const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
+const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
+
+@device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
+    SPIRVIntrinsics.sub_group_reduce_add(val)
+@device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_reduce_min(val)
+@device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_reduce_max(val)
+
+@device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
+@device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
+@device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
+
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 746747707..304d6c53a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -145,6 +145,23 @@ end
     @test occursin("store i64 $width", ir)
 end
 
+# the native sub-group collectives are used where they have Julia's semantics
+function sub_group_reduce_kernel(out, x)
+    i = KernelAbstractions.KernelInterface.get_global_id().x
+    out[i] = KernelAbstractions.KernelInterface.sub_group_reduce(+, x[i])
+    return
+end
+@testset "POCL native sub-group collectives" begin
+    for (T, native) in ((Int32, true), (Float32, true))
+        x, out = ones(T, 32), zeros(T, 32)
+        ir = sprint() do io
+            @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x)
+        end
+        @test all(==(32), out)
+        @test occursin("sub_group_reduce_add", ir) == native
+    end
+end
+
 # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with
 # many arguments allocates unless every layer passes them on as a tuple
 @testset "POCL launch with many arguments" begin

From 3a9eb4af7ee90b222f3de81e13bdb15eb9dbfda3 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 10:57:50 +0200
Subject: [PATCH 06/18] POCL: don't use the native sub-group collectives

PoCL's `cl_khr_subgroups` reductions and scans lose the values of work-items
that computed them in a divergent branch (PoCL 7.2; Intel's OpenCL runtime is
fine), so `@groupreduce` in a `@kernel`, whose padding work-items are masked,
returned garbage. Use KernelInterface's fallbacks again, and test reductions
and scans of values from a divergent branch.

Assisted-by: Claude Code (Opus 5.5)
---
 lib/KernelInterface/test/interface.jl | 29 +++++++++++++++++++++++++++
 src/pocl/backend.jl                   | 22 ++++----------------
 test/runtests.jl                      | 17 ----------------
 3 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 0c00ae2a4..2990aeb1f 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -340,6 +340,17 @@ primitive type Bits16 16 end
 Bits64(x::Integer) = reinterpret(Bits64, x % UInt64)
 Bits16(x::Integer) = reinterpret(Bits16, x % UInt16)
 
+# the values come from a divergent branch, like for the padding work-items of a `@kernel`
+function reduce_divergent_kernel(red, scan, a, m)
+    i = KI.get_local_id().x
+    val = i <= m ? (@inbounds a[i]) : zero(eltype(a))
+    r = KI.sub_group_reduce(+, val)
+    s = KI.sub_group_scan(+, val)
+    @inbounds red[i] = r
+    @inbounds scan[i] = s
+    return
+end
+
 struct FallbackStruct
     flag::Bool
     c::Char
@@ -435,6 +446,19 @@ function reduce_scan_testsuite(backend, AT, op, a)
     return
 end
 
+function reduce_divergent_testsuite(backend, AT, sg_size, ::Type{T}) where {T}
+    m = max(sg_size - 5, 1)
+    a = T.(rand(1:20, sg_size))
+    red = AT(zeros(T, sg_size))
+    scan = AT(zeros(T, sg_size))
+    KI.@launch backend workgroupsize = sg_size reduce_divergent_kernel(red, scan, AT(a), m)
+    KI.synchronize(backend)
+    # the sum doesn't depend on the order of the lanes
+    @test all(==(sum(a[1:m])), Array(red))
+    @test maximum(Array(scan)) == sum(a[1:m])
+    return
+end
+
 function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     @testset "shuffles of other types" begin
         # primitive types that backends need not support natively are shuffled as words
@@ -496,6 +520,11 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
         end
     end
 
+    @testset "sub_group_reduce and sub_group_scan of divergent values" begin
+        reduce_divergent_testsuite(backend, AT, sg_size, Int32)
+        KI.supports_shuffle(backend, Float32) && reduce_divergent_testsuite(backend, AT, sg_size, Float32)
+    end
+
     for n in unique((sg_size, max(sg_size - 3, 1)))
         @testset "sub_group_reduce and sub_group_scan, $n work-items" begin
             reduce_scan_testsuite(backend, AT, +, Int32.(rand(1:100, n)))
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index 6b15f2d18..691542607 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -270,24 +270,10 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
-# native reductions and scans of `cl_khr_subgroups`. Only `+` for floats: OpenCL's `min` and
-# `max` treat NaN and the sign of zero differently from Julia's.
-const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
-const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
-
-@device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
-    SPIRVIntrinsics.sub_group_reduce_add(val)
-@device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
-    SPIRVIntrinsics.sub_group_reduce_min(val)
-@device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
-    SPIRVIntrinsics.sub_group_reduce_max(val)
-
-@device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
-    SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
-@device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
-    SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
-@device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
-    SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
+# `KI.sub_group_reduce` and `KI.sub_group_scan` use KernelInterface's fallbacks: PoCL's native
+# `cl_khr_subgroups` collectives (`SPIRVIntrinsics.sub_group_reduce_add` etc.) lose the values
+# of work-items that computed them in a divergent branch, as the padding work-items of a
+# `@kernel` do (PoCL 7.2).
 
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
diff --git a/test/runtests.jl b/test/runtests.jl
index 304d6c53a..746747707 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -145,23 +145,6 @@ end
     @test occursin("store i64 $width", ir)
 end
 
-# the native sub-group collectives are used where they have Julia's semantics
-function sub_group_reduce_kernel(out, x)
-    i = KernelAbstractions.KernelInterface.get_global_id().x
-    out[i] = KernelAbstractions.KernelInterface.sub_group_reduce(+, x[i])
-    return
-end
-@testset "POCL native sub-group collectives" begin
-    for (T, native) in ((Int32, true), (Float32, true))
-        x, out = ones(T, 32), zeros(T, 32)
-        ir = sprint() do io
-            @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x)
-        end
-        @test all(==(32), out)
-        @test occursin("sub_group_reduce_add", ir) == native
-    end
-end
-
 # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with
 # many arguments allocates unless every layer passes them on as a tuple
 @testset "POCL launch with many arguments" begin

From 8030b4e6b4fe5e0dbb457b97cec1f5595814cf2b Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 12:17:57 +0200
Subject: [PATCH 07/18] KernelInterface: document that sub-group operations
 don't order memory

Like the shuffles, the votes, `sub_group_match_any`, `sub_group_reduce` and
`sub_group_scan` exchange values, not memory; communicating through memory
within a sub-group needs `sub_group_barrier`.

Assisted-by: Claude Code (Opus 5.5)
---
 docs/src/kernelinterface.md       |  6 ++++++
 lib/KernelInterface/src/device.jl | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index 274a742cb..cf1ca5aab 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -197,6 +197,12 @@ localmemory
 
 ### Communication
 
+The shuffles, votes and collectives below exchange values between the work-items of a
+sub-group, not memory: they don't order or make visible the work-items' accesses to local
+or global memory. To communicate through memory within a sub-group, e.g. a work-item reading
+what another one wrote to local memory, use [`sub_group_barrier`](@ref) between the write
+and the read.
+
 ```@docs
 shfl
 shfl_down
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index 59880daf4..c40e64cc3 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -481,6 +481,8 @@ get the same result.
 All work-items of the sub-group have to execute `sub_group_any` together (not in a divergent
 branch).
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backend implementations that support sub-groups **must** implement:
     ```
@@ -498,6 +500,8 @@ sub-group get the same result.
 All work-items of the sub-group have to execute `sub_group_all` together (not in a divergent
 branch).
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backend implementations that support sub-groups **must** implement:
     ```
@@ -518,6 +522,8 @@ one.
 All work-items of the sub-group have to execute `sub_group_ballot` together (not in a
 divergent branch). Only sub-groups of at most 64 work-items are supported.
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backend implementations that support sub-groups with a width of at most 64 **must**
     implement:
@@ -539,6 +545,8 @@ All work-items of the sub-group have to execute `sub_group_match_any` together (
 divergent branch). Values of the types that the shuffles support are supported, see
 [`supports_shuffle`](@ref).
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backends **may** implement this, e.g. with CUDA's `match.any.sync`. The fallback finds
     the groups of equal values one by one with [`shfl`](@ref) and [`sub_group_ballot`](@ref),
@@ -568,6 +576,8 @@ All work-items of the sub-group have to execute `sub_group_reduce` together (not
 divergent branch). Values of the types that the shuffles support are supported, see
 [`supports_shuffle`](@ref).
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backends **may** implement this for operators and types with a native reduction (e.g.
     `+` on `Float32`), dispatching on `typeof(op)`. The fallback combines ranges of doubling
@@ -600,6 +610,8 @@ All work-items of the sub-group have to execute `sub_group_scan` together (not i
 branch). Values of the types that the shuffles support are supported, see
 [`supports_shuffle`](@ref).
 
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
 !!! note
     Backends **may** implement this for operators and types with a native scan (e.g. `+` on
     `Float32`), dispatching on `typeof(op)`. The fallback is a Hillis-Steele scan with

From 03cfa30e8f5d362877be401301f505332f6e60fe Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 12:32:31 +0200
Subject: [PATCH 08/18] POCL: prepare the native sub-group collectives,
 disabled for now

The wrong results of the native `cl_khr_subgroups` collectives had two causes:
SPIRVIntrinsics declared them without `convergent`, so LLVM duplicated the
calls into divergent branches (fixed in JuliaGPU/OpenCL.jl#526), and PoCL 7.2
gives a peeled work-item its own copy of a collective's scratch memory after a
branch with an early exit, as bounds checks emit (fixed on PoCL's main branch,
backport to 7.2 in pocl/pocl#2373).

With both fixed, i.e. with `POCL_WORK_GROUP_METHOD=cbs` for now, the native
reductions and scans pass the tests. Keep them behind `NATIVE_COLLECTIVES`
until `pocl_standalone_jll` includes the PoCL fix.

Assisted-by: Claude Code (Opus 5.5)
---
 src/pocl/backend.jl | 32 ++++++++++++++++++++++++++++----
 test/runtests.jl    | 17 +++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index 691542607..6f2cad120 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -270,10 +270,34 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
-# `KI.sub_group_reduce` and `KI.sub_group_scan` use KernelInterface's fallbacks: PoCL's native
-# `cl_khr_subgroups` collectives (`SPIRVIntrinsics.sub_group_reduce_add` etc.) lose the values
-# of work-items that computed them in a divergent branch, as the padding work-items of a
-# `@kernel` do (PoCL 7.2).
+# Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and
+# floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
+# differently from Julia's).
+#
+# Disabled until `pocl_standalone_jll` includes pocl/pocl#2373: in PoCL 7.2, a collective after
+# a branch with an early exit (as bounds checks emit) gets wrong values, because WorkitemLoops
+# gives the peeled first work-item its own copy of the collective's scratch memory. Until
+# then, `KI.sub_group_reduce` and `KI.sub_group_scan` use KernelInterface's fallbacks.
+const NATIVE_COLLECTIVES = false
+
+@static if NATIVE_COLLECTIVES
+    const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
+    const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
+
+    @device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
+        SPIRVIntrinsics.sub_group_reduce_add(val)
+    @device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_reduce_min(val)
+    @device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_reduce_max(val)
+
+    @device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
+    @device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
+    @device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
+end
 
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
diff --git a/test/runtests.jl b/test/runtests.jl
index 746747707..2f59678b6 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -145,6 +145,23 @@ end
     @test occursin("store i64 $width", ir)
 end
 
+# the native sub-group collectives are used where they have Julia's semantics
+function sub_group_reduce_kernel(out, x)
+    i = KernelAbstractions.KernelInterface.get_global_id().x
+    out[i] = KernelAbstractions.KernelInterface.sub_group_reduce(+, x[i])
+    return
+end
+@testset "POCL native sub-group collectives" begin
+    for (T, native) in ((Int32, true), (Float32, true))
+        x, out = ones(T, 32), zeros(T, 32)
+        ir = sprint() do io
+            @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x)
+        end
+        @test all(==(32), out)
+        @test occursin("sub_group_reduce_add", ir) == (native && POCL.POCLKernels.NATIVE_COLLECTIVES)
+    end
+end
+
 # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with
 # many arguments allocates unless every layer passes them on as a tuple
 @testset "POCL launch with many arguments" begin

From b3c71351b7588bd95a712cbc291681809dffe48a Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 12:43:08 +0200
Subject: [PATCH 09/18] POCL: point at the Yggdrasil backport of the PoCL fix

Assisted-by: Claude Code (Opus 5.5)
---
 src/pocl/backend.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index 6f2cad120..cd46d36af 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -274,7 +274,8 @@ end
 # floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
 # differently from Julia's).
 #
-# Disabled until `pocl_standalone_jll` includes pocl/pocl#2373: in PoCL 7.2, a collective after
+# Disabled until `pocl_standalone_jll` includes pocl/pocl#2239 (JuliaPackaging/Yggdrasil#15001):
+# in PoCL 7.2, a collective after
 # a branch with an early exit (as bounds checks emit) gets wrong values, because WorkitemLoops
 # gives the peeled first work-item its own copy of the collective's scratch memory. Until
 # then, `KI.sub_group_reduce` and `KI.sub_group_scan` use KernelInterface's fallbacks.

From dbae32e827dc9c2111268e33855fb21f2c62851e Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 12:58:28 +0200
Subject: [PATCH 10/18] KernelInterface: votes within segments of a width

`sub_group_any`, `sub_group_all`, `sub_group_ballot` and `sub_group_match_any`
with a `width`, like the shuffles: the votes of segments of `width` lanes, with
masks that have a bit per lane of the segment. Fallbacks use the ballot and
match of the whole sub-group. With the shuffles with a width, this allows e.g.
tiles of 32 lanes on sub-groups of 64.

Assisted-by: Claude Code (Opus 5.5)
---
 docs/src/kernelinterface.md           |  3 +-
 lib/KernelInterface/src/device.jl     | 51 +++++++++++++++++++++
 lib/KernelInterface/test/interface.jl | 66 +++++++++++++++++++++++++++
 lib/KernelInterface/test/runtests.jl  |  8 +++-
 4 files changed, 126 insertions(+), 2 deletions(-)

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index cf1ca5aab..f282d05cc 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
-| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
+| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
 
 Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones),
 the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch),
@@ -213,6 +213,7 @@ sub_group_any
 sub_group_all
 sub_group_ballot
 sub_group_match_any
+sub_group_any(::Bool, ::Integer)
 sub_group_reduce
 sub_group_scan
 ```
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index c40e64cc3..291278800 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -566,6 +566,57 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri
     return mask
 end
 
+# The votes within segments of `width` lanes, implemented with `sub_group_ballot`.
+
+# The 0-based lane of the first work-item of this work-item's segment of `width` lanes
+@inline segment_base(width) = (get_sub_group_local_id(Int32) - Int32(1)) & ~(width % Int32 - Int32(1))
+
+# The bits of the segment of `width` lanes starting at `base` of a mask, moved to the bottom
+@inline function segment_bits(mask::UInt64, base, width)
+    w = width % UInt32
+    bits = w >= UInt32(64) ? typemax(UInt64) : (UInt64(1) << w) - UInt64(1)
+    return (mask >> (base % UInt32)) & bits
+end
+
+"""
+    sub_group_any(pred::Bool, width::Integer)::Bool
+    sub_group_all(pred::Bool, width::Integer)::Bool
+    sub_group_ballot(pred::Bool, width::Integer)::UInt64
+    sub_group_match_any(val, width::Integer)::UInt64
+
+Votes within segments of `width` consecutive lanes of the sub-group, as if each segment were a
+sub-group of its own (like the shuffles with a `width`, e.g. [`shfl`](@ref
+shfl(::Any, ::Integer, ::Integer))): the work-items of a segment get the result for the
+work-items of their segment. The masks of `sub_group_ballot` and `sub_group_match_any` have a
+bit per lane of the segment: bit `i - 1` for its `i`-th lane, i.e. the lane
+`get_sub_group_local_id()` with `(get_sub_group_local_id() - 1) % width == i - 1`. In a
+partial sub-group, only the work-items of the segment take part.
+
+`width` has to be a power of two of at most the sub-group width
+[`get_max_sub_group_size`](@ref) (and of at most 64), and the same for all work-items of the
+sub-group, which all have to execute the vote together (not in a divergent branch).
+
+It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref).
+
+!!! note
+    Backends **may** implement these. The fallbacks use [`sub_group_ballot`](@ref) and
+    [`sub_group_match_any`](@ref) of the whole sub-group.
+"""
+@inline sub_group_ballot(pred::Bool, width::Integer) =
+    segment_bits(sub_group_ballot(pred), segment_base(width), width)
+
+@inline sub_group_any(pred::Bool, width::Integer) = sub_group_ballot(pred, width) != zero(UInt64)
+
+@inline function sub_group_all(pred::Bool, width::Integer)
+    # compare with the work-items of the segment, which may be partial
+    base = segment_base(width)
+    return segment_bits(sub_group_ballot(pred), base, width) ==
+        segment_bits(sub_group_ballot(true), base, width)
+end
+
+@inline sub_group_match_any(val, width::Integer) =
+    segment_bits(sub_group_match_any(val), segment_base(width), width)
+
 """
     sub_group_reduce(op, val::T)::T
 
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 2990aeb1f..86b99f356 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -351,6 +351,25 @@ function reduce_divergent_kernel(red, scan, a, m)
     return
 end
 
+# the votes within segments of `width` lanes; a work-group of `length(pred)` work-items, i.e. a
+# single (possibly partial) sub-group
+function segmented_vote_kernel(out, lanes, pred, vals, width)
+    i = KI.get_local_id().x
+    p = @inbounds pred[i]
+    any = KI.sub_group_any(p, width)
+    all = KI.sub_group_all(p, width)
+    ballot = KI.sub_group_ballot(p, width)
+    match = KI.sub_group_match_any(@inbounds(vals[i]), width)
+    @inbounds begin
+        out[i, 1] = any
+        out[i, 2] = all
+        out[i, 3] = ballot
+        out[i, 4] = match
+        lanes[i] = KI.get_sub_group_local_id()
+    end
+    return
+end
+
 struct FallbackStruct
     flag::Bool
     c::Char
@@ -459,6 +478,42 @@ function reduce_divergent_testsuite(backend, AT, sg_size, ::Type{T}) where {T}
     return
 end
 
+function segmented_vote_testsuite(backend, AT, n, width)
+    pred = [rand(Bool) for _ in 1:n]
+    # make some segments all true and some all false
+    for i in 1:n
+        seg = (i - 1) ÷ width
+        seg % 3 == 1 && (pred[i] = true)
+        seg % 3 == 2 && (pred[i] = false)
+    end
+    vals = Int32.(rand(1:3, n))
+    out = AT(zeros(UInt64, n, 4))
+    lanes = AT(zeros(Int, n))
+    KI.@launch backend workgroupsize = n segmented_vote_kernel(out, lanes, AT(pred), AT(vals), width)
+    KI.synchronize(backend)
+    out, lanes = Array(out), Array(lanes)
+    by_lane = zeros(Int, maximum(lanes))
+    for i in 1:n
+        by_lane[lanes[i]] = i
+    end
+    for i in 1:n
+        base = (lanes[i] - 1) ÷ width * width
+        # the work-items of the segment, by their position in the segment
+        seg = [(l - base, by_lane[l]) for l in (base + 1):min(base + width, length(by_lane)) if by_lane[l] != 0]
+        ballot = UInt64(0)
+        match = UInt64(0)
+        for (k, j) in seg
+            pred[j] && (ballot |= UInt64(1) << (k - 1))
+            vals[j] == vals[i] && (match |= UInt64(1) << (k - 1))
+        end
+        @test out[i, 1] == any(pred[j] for (_, j) in seg)
+        @test out[i, 2] == all(pred[j] for (_, j) in seg)
+        @test out[i, 3] == ballot
+        @test out[i, 4] == match
+    end
+    return
+end
+
 function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     @testset "shuffles of other types" begin
         # primitive types that backends need not support natively are shuffled as words
@@ -520,6 +575,17 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
         end
     end
 
+    if sg_size <= 64
+        @testset "votes with a width" begin
+            for n in unique((sg_size, max(sg_size - 3, 1))), w in (1, 2, 4, 8, 16, 32, 64)
+                (w <= sg_size && sg_size % w == 0) || continue
+                @testset "$n work-items, width $w" begin
+                    segmented_vote_testsuite(backend, AT, n, w)
+                end
+            end
+        end
+    end
+
     @testset "sub_group_reduce and sub_group_scan of divergent values" begin
         reduce_divergent_testsuite(backend, AT, sg_size, Int32)
         KI.supports_shuffle(backend, Float32) && reduce_divergent_testsuite(backend, AT, sg_size, Float32)
diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl
index 762c0890c..3d1325504 100644
--- a/lib/KernelInterface/test/runtests.jl
+++ b/lib/KernelInterface/test/runtests.jl
@@ -34,7 +34,6 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true
     # These have no fallback on purpose: a backend that forgets to `@device_override`
     # them should get a MethodError rather than silently wrong behaviour.
     stubs = [
-        KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot,
         KI.max_work_group_size, KI.max_work_group_dims, KI.max_num_groups,
         KI.sub_group_size, KI.argconvert, KI.kernel_function, KI.launch,
         # Host-side stubs: required backend methods with no sensible fallback.
@@ -44,6 +43,13 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true
         @test isempty(methods(stub))
     end
 
+    # The votes only have the fallbacks within segments of a `width`, built on the votes
+    # of the whole sub-group that backends implement.
+    for vote in [KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot]
+        @test !hasmethod(vote, Tuple{Bool})
+        @test hasmethod(vote, Tuple{Bool, Int})
+    end
+
     # The shuffles only have the fallback that shuffles structs field by field, which
     # doesn't handle the primitive types a backend has to implement.
     for shfl in [KI.shfl, KI.shfl_down, KI.shfl_up, KI.shfl_xor]

From 4b37566fdfd2df3a914fe6072e237f45abd88bc6 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 13:45:12 +0200
Subject: [PATCH 11/18] KernelInterface: guarantee the sub-group layout, define
 shuffles past the width

- If a work-group is 1-D or its x extent is a multiple of the sub-group width,
  sub-groups are formed from consecutive work-items, x fastest. This holds on
  CUDA, AMD, PoCL, rusticl and Intel's CPU OpenCL runtime (which forms
  sub-groups per row otherwise), and lets kernels with e.g. (32, 8)
  work-groups rely on the layout.
- `shfl_down` and `shfl_up` return the work-item's own value where the source
  lane is past the sub-group width, like CUDA's shuffles. That's free on CUDA
  and Metal, and a select on AMD and SPIR-V (POCL), and makes them consistent
  with the shuffles with a `width`. `shfl_xor` requires a mask below the width.

Test both.

Assisted-by: Claude Code (Opus 5.5)
---
 docs/src/kernelinterface.md           | 21 ++++++++----
 lib/KernelInterface/src/device.jl     | 30 ++++++++++-------
 lib/KernelInterface/test/interface.jl | 48 +++++++++++++++++++++++++++
 src/pocl/backend.jl                   | 15 ++++++---
 4 files changed, 92 insertions(+), 22 deletions(-)

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index f282d05cc..9950833f2 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -151,17 +151,26 @@ get_global_size
 Sub-groups are optional ([`supports_subgroups`](@ref)). A work-group is divided into
 sub-groups of at most [`sub_group_size(backend)`](@ref sub_group_size) work-items. Which
 work-items form a sub-group, how many sub-groups there are, and which of them are partial
-is unspecified, and differs between devices and work-group shapes. For example, CUDA forms
-warps from consecutive linear work-item indices, while Intel's CPU OpenCL runtime forms
-sub-groups per row of a multi-dimensional work-group, so that a 33×2 work-group consists of
-four sub-groups of 32 and 1 work-items. What KernelInterface guarantees, and backends that
-report sub-group support have to ensure:
+can differ between devices and work-group shapes. For example, CUDA forms warps from
+consecutive linear work-item indices, while Intel's CPU OpenCL runtime forms sub-groups per
+row of a multi-dimensional work-group, so that a 33×2 work-group consists of four sub-groups
+of 32 and 1 work-items. What KernelInterface guarantees, and backends that report sub-group
+support have to ensure:
 
 - every work-item has a unique `(get_sub_group_id(), get_sub_group_local_id())` pair in its
   work-group, which doesn't change during the kernel;
 - the sub-group ids are `1:get_num_sub_groups()`, and the lanes of a sub-group are
   `1:get_sub_group_size()`;
-- a 1-D work-group of at most `sub_group_size(backend)` work-items is a single sub-group.
+- if the work-group is 1-D, or its x extent `get_local_size().x` is a multiple of the
+  sub-group width `W` ([`get_max_sub_group_size`](@ref)), sub-groups are formed from
+  consecutive work-items, x fastest: the work-item with the linear index
+  `lin = x + (y - 1) * size.x + (z - 1) * size.x * size.y` (for `(; x, y, z) =
+  get_local_id()` and `size = get_local_size()`) is in sub-group `(lin - 1) ÷ W + 1`, lane
+  `(lin - 1) % W + 1`. Only the last sub-group can be partial. In particular, a 1-D
+  work-group of at most `W` work-items is a single sub-group.
+
+Other shapes can form sub-groups differently, e.g. per row of the work-group: for portable
+code, make the x extent of multi-dimensional work-groups a multiple of the sub-group width.
 
 A sub-group is partial when it has fewer work-items than the sub-group width: the lanes
 `get_sub_group_size()+1:get_max_sub_group_size()` have no work-item. Shuffles from those lanes
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index 291278800..e0992c7c8 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -129,9 +129,10 @@ end
 # sub-groups of at most `sub_group_size(backend)` work-items. Which work-items form a
 # sub-group, how many sub-groups there are and which are partial is unspecified, except
 # that `(get_sub_group_id(), get_sub_group_local_id())` is unique within a work-group and
-# doesn't change during the kernel's execution, and that a 1-D work-group of at most
-# `sub_group_size(backend)` work-items is a single sub-group. Backends that can't ensure
-# that don't report sub-group support. See the manual.
+# doesn't change during the kernel's execution, and that sub-groups are formed from
+# consecutive work-items, x fastest, if the work-group is 1-D or its x extent is a multiple of
+# the sub-group width (see the manual). Backends that can't ensure that don't report sub-group
+# support.
 #
 # In a partial sub-group, the lanes `get_sub_group_size()+1:get_max_sub_group_size()` have no
 # work-item: shuffles from them give unspecified values, and the votes, `sub_group_match_any`,
@@ -303,8 +304,11 @@ field by field.
     shfl_down(val::T, offset::Integer)::T
 
 Return `val` of the work-item `offset` lanes further in the sub-group, i.e. with
-[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`. When there is
-no such work-item, the result is an unspecified value (of type `T`).
+[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`, for an `offset`
+of at least 0. When that lane is past the sub-group width, i.e.
+`get_sub_group_local_id() + offset > get_max_sub_group_size()`, the result is `val` of the
+work-item itself, like CUDA's `shfl_down_sync`. When the lane is within the width but has no
+work-item, in a partial sub-group, the result is an unspecified value (of type `T`).
 
 All work-items of the sub-group have to execute `shfl_down` together (not in a divergent
 branch), with the same `offset`.
@@ -324,8 +328,9 @@ the supported types.
     shfl_up(val::T, offset::Integer)::T
 
 Return `val` of the work-item `offset` lanes earlier in the sub-group, i.e. with
-[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() - offset`. When there is
-no such work-item, the result is an unspecified value (of type `T`).
+[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() - offset`, for an `offset`
+of at least 0. When there is no such lane, i.e. `get_sub_group_local_id() <= offset`, the
+result is `val` of the work-item itself, like CUDA's `shfl_up_sync`.
 
 All work-items of the sub-group have to execute `shfl_up` together (not in a divergent
 branch), with the same `offset`.
@@ -346,8 +351,9 @@ the supported types.
 
 Return `val` of the work-item whose 0-based lane id is the 0-based lane id of this work-item
 xor `mask`, i.e. with [`get_sub_group_local_id`](@ref) equal to
-`((get_sub_group_local_id() - 1) ⊻ mask) + 1`. When there is no such work-item, the result is
-an unspecified value (of type `T`).
+`((get_sub_group_local_id() - 1) ⊻ mask) + 1`, for a `mask` between 0 and
+`get_max_sub_group_size() - 1`. When that lane has no work-item, in a partial sub-group, the
+result is an unspecified value (of type `T`).
 
 All work-items of the sub-group have to execute `shfl_xor` together (not in a divergent
 branch), with the same `mask`. A butterfly over the masks `width ÷ 2, …, 2, 1` (for the
@@ -430,9 +436,9 @@ Shuffles within segments of `width` consecutive lanes of the sub-group, as if ea
 were a sub-group of its own: `lane` is the lane within the segment (between 1 and `width`, and
 taken modulo `width` otherwise), and `shfl_down`, `shfl_up` and `shfl_xor` read from lanes of
 the same segment. Where these would read from outside of the segment, they return `val` of the
-work-item itself (rather than an unspecified value, as without `width`), like CUDA's
-shuffles with a `width`. Reading from a lane of the segment that has no work-item (in a
-partial sub-group) gives an unspecified value.
+work-item itself, like CUDA's shuffles with a `width` (and like `shfl_down` and `shfl_up`
+without a `width` past the sub-group width). Reading from a lane of the segment that has no
+work-item (in a partial sub-group) gives an unspecified value.
 
 `width` has to be a power of two of at most the sub-group width
 [`get_max_sub_group_size`](@ref), and the same for all work-items of the sub-group.
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 86b99f356..08b05e9b4 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -370,6 +370,19 @@ function segmented_vote_kernel(out, lanes, pred, vals, width)
     return
 end
 
+# the sub-group and lane of every work-item, by its linear index (x fastest)
+function sub_group_layout_kernel(out)
+    l = KI.get_local_id()
+    sz = KI.get_local_size()
+    lin = l.x + (l.y - 1) * sz.x + (l.z - 1) * sz.x * sz.y
+    @inbounds begin
+        out[1, lin] = KI.get_sub_group_id()
+        out[2, lin] = KI.get_sub_group_local_id()
+        out[3, lin] = KI.get_sub_group_size()
+    end
+    return
+end
+
 struct FallbackStruct
     flag::Bool
     c::Char
@@ -514,6 +527,30 @@ function segmented_vote_testsuite(backend, AT, n, width)
     return
 end
 
+# 1-D work-groups and ones whose x extent is a multiple of the sub-group width form
+# sub-groups from consecutive work-items, x fastest
+function sub_group_layout_testsuite(backend, AT, sg_size, fits)
+    shapes = ((sg_size + 5,), (3 * sg_size,), (sg_size, 4), (2 * sg_size, 2), (sg_size, 2, 2))
+    for dims in shapes
+        n = prod(dims)
+        out = AT(zeros(Int, 3, n))
+        kernel = KI.@launch backend launch = false sub_group_layout_kernel(out)
+        if !fits(kernel, dims)
+            @test_skip "work-groups of $dims work-items"
+            continue
+        end
+        kernel(out; workgroupsize = dims)
+        KI.synchronize(backend)
+        out = Array(out)
+        @testset "$dims" begin
+            @test out[1, :] == [(lin - 1) ÷ sg_size + 1 for lin in 1:n]
+            @test out[2, :] == [(lin - 1) % sg_size + 1 for lin in 1:n]
+            @test out[3, :] == [min(sg_size, n - (lin - 1) ÷ sg_size * sg_size) for lin in 1:n]
+        end
+    end
+    return
+end
+
 function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     @testset "shuffles of other types" begin
         # primitive types that backends need not support natively are shuffled as words
@@ -940,6 +977,10 @@ function interface_testsuite(backend::KI.Backend, AT)
         # whether `kernel` can be launched with work-groups of size `dims`
         fits(kernel, dims) = all(dims .<= max_dims[1:length(dims)]) && prod(dims) <= KI.max_work_group_size(kernel)
 
+        @testset "Sub-group layout" begin
+            sub_group_layout_testsuite(backend, AT, sg_size, fits)
+        end
+
         @testset "Sub-group return types" begin
             @test sg_size isa Int && sg_size >= 1
 
@@ -1060,6 +1101,8 @@ function interface_testsuite(backend::KI.Backend, AT)
                     KI.@launch backend workgroupsize = sg_size shfl_up_lanes_kernel(out, T, offset)
                     KI.synchronize(backend)
                     @test Array(out)[(offset + 1):end] == T.(1:(sg_size - offset))
+                    # lanes without a lane `offset` earlier get their own value
+                    @test Array(out)[1:offset] == T.(1:offset)
                 end
 
                 if ispow2(sg_size)
@@ -1143,6 +1186,11 @@ function interface_testsuite(backend::KI.Backend, AT)
                     end
                     @test !isempty(in_range)
                     @test out[in_range, 3] == out[in_range, 1] .+ offset
+                    # a 1-D work-group forms full sub-groups (but the last one), and lanes
+                    # past the sub-group width get their own value
+                    past = findall(i -> out[i, 1] + offset > sg_size, 1:N)
+                    @test !isempty(past)
+                    @test out[past, 3] == out[past, 1]
                 end
             end
         end
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index cd46d36af..77bf84195 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -252,11 +252,18 @@ end
 @device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} =
     sub_group_shuffle(val, lane)
 
-@device_override KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} =
-    sub_group_shuffle(val, get_sub_group_local_id() + offset)
+# past the sub-group width, `shfl_down` and `shfl_up` return the work-item's own value, which
+# `sub_group_shuffle` (like SPIR-V's `OpGroupNonUniformShuffleDown`) leaves undefined
+@device_override function KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes}
+    lane = get_sub_group_local_id()
+    src = lane + offset
+    return sub_group_shuffle(val, ifelse(src <= get_max_sub_group_size(), src, lane))
+end
 
-@device_override KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} =
-    sub_group_shuffle(val, get_sub_group_local_id() - offset)
+@device_override function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes}
+    lane = get_sub_group_local_id()
+    return sub_group_shuffle(val, ifelse(lane > offset, lane - offset, lane))
+end
 
 @device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} =
     sub_group_shuffle_xor(val, mask)

From 1386de4ded953218690aff6155d6d9bddea83c71 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 13:57:58 +0200
Subject: [PATCH 12/18] KernelInterface: faster fallbacks of sub_group_reduce
 and sub_group_scan

Porting KomaMRI.jl showed the fallback of `sub_group_reduce` to cost ~40% of a
reduction-heavy kernel on CUDA: its loop ran to the run-time sub-group size, so it
didn't unroll, and every lane was masked.

`sub_group_reduce` is now an ordered butterfly with `shfl_xor` over the constant
sub-group width (combining the lower block first, so it still only needs
associativity), which gives every work-item of a full sub-group the result
without a broadcast. Partial sub-groups skip the blocks without work-items and
broadcast the result of the first lane. All sub-groups run the same shuffles,
which PoCL needs; document that requirement of the POCL backend. `sub_group_scan`
loops to the constant width as well.

Test reductions and scans in a work-group of several sub-groups, the last one
partial.

Assisted-by: Claude Code (Opus 5.5)
---
 lib/KernelInterface/src/device.jl     | 55 ++++++++++++++++++++-------
 lib/KernelInterface/test/interface.jl | 37 ++++++++++++++++++
 src/KernelAbstractions.jl             |  8 ++++
 3 files changed, 86 insertions(+), 14 deletions(-)

diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index e0992c7c8..405d3caa2 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -638,22 +638,48 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri
 !!! note
     Backends **may** implement this for operators and types with a native reduction (e.g.
     `+` on `Float32`), dispatching on `typeof(op)`. The fallback combines ranges of doubling
-    length with [`shfl_down`](@ref), and broadcasts the result of the first lane with
-    [`shfl`](@ref).
+    length with a butterfly of [`shfl_xor`](@ref), and in a partial sub-group broadcasts the
+    result of the first lane with [`shfl`](@ref).
 """
 @inline function sub_group_reduce(op, val)
-    lane = get_sub_group_local_id(Int32)
+    width = get_max_sub_group_size(Int32)
     sgsize = get_sub_group_size(Int32)
-    offset = Int32(1)
-    while offset < sgsize
-        other = shfl_down(val, offset)
-        # the result of shuffling from past the end of the sub-group is unspecified
-        if lane + offset <= sgsize
-            val = op(val, other)
+    lane0 = get_sub_group_local_id(Int32) - Int32(1)
+    if ispow2(width)
+        # A butterfly with `shfl_xor`, which unrolls for the constant width. Each step combines
+        # the block of a work-item with the neighboring block, the lower one first, so that
+        # only associativity is needed. In a full sub-group, every work-item ends up with the
+        # reduction; in a partial one, blocks without work-items are skipped, which keeps the
+        # result of the first lane correct, and it is broadcast. All sub-groups run the same
+        # shuffles, since some backends (PoCL) need that across the sub-groups of a work-group.
+        mask = Int32(1)
+        while mask < width
+            other = shfl_xor(val, mask)
+            if (lane0 ⊻ mask) < sgsize
+                if lane0 & mask == Int32(0)
+                    val = op(val, other)
+                else
+                    val = op(other, val)
+                end
+            end
+            mask <<= 1
         end
-        offset <<= 1
+        first = shfl(val, 1)
+        return ifelse(sgsize == width, val, first)
+    else
+        # combine ranges of doubling length with `shfl_down`, skipping the lanes without a
+        # work-item, and broadcast the result of the first lane
+        lane = lane0 + Int32(1)
+        offset = Int32(1)
+        while offset < width
+            other = shfl_down(val, offset)
+            if lane + offset <= sgsize
+                val = op(val, other)
+            end
+            offset <<= 1
+        end
+        return shfl(val, 1)
     end
-    return shfl(val, 1)
 end
 
 """
@@ -676,11 +702,12 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri
 """
 @inline function sub_group_scan(op, val)
     lane = get_sub_group_local_id(Int32)
-    sgsize = get_sub_group_size(Int32)
+    # loop to the constant width, so that the loop unrolls: the lanes `shfl_up` reads from
+    # always have a work-item, also in a partial sub-group
+    width = get_max_sub_group_size(Int32)
     offset = Int32(1)
-    while offset < sgsize
+    while offset < width
         other = shfl_up(val, offset)
-        # the result of shuffling from before the start of the sub-group is unspecified
         if lane > offset
             val = op(other, val)
         end
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 08b05e9b4..824acfa5f 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -383,6 +383,20 @@ function sub_group_layout_kernel(out)
     return
 end
 
+# reductions and scans in a 1-D work-group of several sub-groups, the last one partial
+function reduce_scan_multi_kernel(red, scan, sgs, a)
+    i = KI.get_local_id().x
+    val = @inbounds a[i]
+    r = KI.sub_group_reduce(+, val)
+    s = KI.sub_group_scan(+, val)
+    @inbounds begin
+        red[i] = r
+        scan[i] = s
+        sgs[i] = KI.get_sub_group_id()
+    end
+    return
+end
+
 struct FallbackStruct
     flag::Bool
     c::Char
@@ -551,6 +565,24 @@ function sub_group_layout_testsuite(backend, AT, sg_size, fits)
     return
 end
 
+function reduce_scan_multi_testsuite(backend, AT, sg_size, ::Type{T}) where {T}
+    n = 2 * sg_size + 5
+    a = T.(rand(1:20, n))
+    red, scan, sgs = AT(zeros(T, n)), AT(zeros(T, n)), AT(zeros(Int, n))
+    kernel = KI.@launch backend launch = false reduce_scan_multi_kernel(red, scan, sgs, AT(a))
+    n <= KI.max_work_group_size(kernel) || return
+    kernel(red, scan, sgs, AT(a); workgroupsize = n)
+    KI.synchronize(backend)
+    red, scan, sgs = Array(red), Array(scan), Array(sgs)
+    for i in 1:n
+        group = findall(==(sgs[i]), sgs)
+        @test red[i] == sum(a[group])
+    end
+    # 1-D work-groups form sub-groups from consecutive work-items
+    @test scan == [sum(a[((i - 1) ÷ sg_size * sg_size + 1):i]) for i in 1:n]
+    return
+end
+
 function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     @testset "shuffles of other types" begin
         # primitive types that backends need not support natively are shuffled as words
@@ -623,6 +655,11 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
         end
     end
 
+    @testset "sub_group_reduce and sub_group_scan, several sub-groups" begin
+        reduce_scan_multi_testsuite(backend, AT, sg_size, Int32)
+        KI.supports_shuffle(backend, Float32) && reduce_scan_multi_testsuite(backend, AT, sg_size, Float32)
+    end
+
     @testset "sub_group_reduce and sub_group_scan of divergent values" begin
         reduce_divergent_testsuite(backend, AT, sg_size, Int32)
         KI.supports_shuffle(backend, Float32) && reduce_divergent_testsuite(backend, AT, sg_size, Float32)
diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl
index 276962365..89c8bdb5a 100644
--- a/src/KernelAbstractions.jl
+++ b/src/KernelAbstractions.jl
@@ -734,6 +734,14 @@ export POCLBackend
 
 CPU backend that compiles kernels to OpenCL via [POCL](https://portablecl.org/) and executes
 them on the host. This is the concrete type behind the [`CPU`](@ref) alias.
+
+!!! warning "Sub-group operations"
+    PoCL implements sub-group operations (shuffles, votes, and the collectives built on
+    them) with barriers across the whole work-group. Beyond KernelInterface's requirement
+    that all work-items of a sub-group execute them together, all sub-groups of a work-group
+    have to execute the same sub-group operations, in the same order: e.g. not in a loop whose
+    trip count depends on the sub-group, or in a branch that only some sub-groups take.
+    Otherwise, results are wrong.
 """
 POCLBackend
 

From 8026fefad139312d471a5a79d2220ed58b98f7e1 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 14:19:04 +0200
Subject: [PATCH 13/18] POCL: reduce and scan with a loop that PoCL 7.2
 compiles correctly

The new fallbacks of `sub_group_reduce` and `sub_group_scan` loop over the
constant sub-group width. PoCL 7.2 miscompiles the unrolled shuffles after a
branch with an early exit, as bounds-checked `@kernel`s have (the bug fixed by
pocl/pocl#2239), so `@groupreduce` with sub-groups gave wrong results. Override
them for POCL with a loop bounded by the smaller of the width and the work-group
size, which doesn't unroll and is the same for all sub-groups of a work-group.

Assisted-by: Claude Code (Opus 5.5)
---
 src/pocl/backend.jl | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index 77bf84195..b2fc209c5 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -277,6 +277,49 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
+# The fallbacks of `KI.sub_group_reduce` and `KI.sub_group_scan` loop over the constant
+# sub-group width, and PoCL 7.2 miscompiles the unrolled shuffles after a branch with an early
+# exit (as bounds checks emit), like the native collectives below (fixed by pocl/pocl#2239,
+# JuliaPackaging/Yggdrasil#15001). Until `pocl_standalone_jll` includes the fix, loop to a
+# bound that isn't a constant, but is the same for all sub-groups of the work-group, which
+# PoCL needs (see `POCLBackend`): the smaller of the width and the work-group size. Reduce by
+# combining ranges of doubling length with `shfl_down`, skipping the lanes without a
+# work-item, and broadcast the result of the first lane.
+@inline function uniform_sub_group_bound()
+    sz = KI.get_local_size(Int32)
+    return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z)
+end
+
+@device_override @inline function KI.sub_group_reduce(op, val)
+    lane = KI.get_sub_group_local_id(Int32)
+    sgsize = KI.get_sub_group_size(Int32)
+    offset = Int32(1)
+    bound = uniform_sub_group_bound()
+    while offset < bound
+        other = KI.shfl_down(val, offset)
+        if lane + offset <= sgsize
+            val = op(val, other)
+        end
+        offset <<= 1
+    end
+    return KI.shfl(val, 1)
+end
+
+# likewise for `KI.sub_group_scan`
+@device_override @inline function KI.sub_group_scan(op, val)
+    lane = KI.get_sub_group_local_id(Int32)
+    offset = Int32(1)
+    bound = uniform_sub_group_bound()
+    while offset < bound
+        other = KI.shfl_up(val, offset)
+        if lane > offset
+            val = op(other, val)
+        end
+        offset <<= 1
+    end
+    return val
+end
+
 # Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and
 # floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
 # differently from Julia's).

From bd36294e3223117d20ac7a3ca64a68206af7776f Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 14:30:38 +0200
Subject: [PATCH 14/18] KernelInterface: test the sub-group layout of a
 work-group past 256 work-items

Assisted-by: Claude Code (Opus 5.5)
---
 lib/KernelInterface/test/interface.jl | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index 824acfa5f..c6ca72316 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -544,7 +544,12 @@ end
 # 1-D work-groups and ones whose x extent is a multiple of the sub-group width form
 # sub-groups from consecutive work-items, x fastest
 function sub_group_layout_testsuite(backend, AT, sg_size, fits)
-    shapes = ((sg_size + 5,), (3 * sg_size,), (sg_size, 4), (2 * sg_size, 2), (sg_size, 2, 2))
+    # including one with a partial last sub-group past 256 work-items, which catches 8-bit
+    # arithmetic in the index computations
+    shapes = (
+        (sg_size + 5,), (3 * sg_size,), (9 * sg_size + 5,),
+        (sg_size, 4), (2 * sg_size, 2), (sg_size, 2, 2),
+    )
     for dims in shapes
         n = prod(dims)
         out = AT(zeros(Int, 3, n))

From 9cb6933beef56093617c6baa4d6f95961362d5eb Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 16:43:12 +0200
Subject: [PATCH 15/18] POCL: use the native sub-group collectives with
 pocl_standalone_jll 7.2.1+1

pocl_standalone_jll 7.2.1+1 includes the WorkitemLoops fix (pocl/pocl#2239,
JuliaPackaging/Yggdrasil#15001), so implement `KI.sub_group_reduce` and
`KI.sub_group_scan` with the native `cl_khr_subgroups` collectives where they
have Julia's semantics, and only keep the workaround for the fallbacks with
older builds. The version is checked at precompile time, since the compat bound
can't distinguish builds.

Note: PoCL's kernel cache doesn't distinguish 7.2.1+0 and +1, so binaries
miscompiled by the former can be reused by the latter until the cache is
cleared (pocl/pocl#2374, JuliaPackaging/Yggdrasil#15002).

Assisted-by: Claude Code (Opus 5.5)
---
 src/pocl/backend.jl | 86 ++++++++++++++++++++++-----------------------
 1 file changed, 43 insertions(+), 43 deletions(-)

diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index b2fc209c5..dd430ea9f 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -277,59 +277,59 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
-# The fallbacks of `KI.sub_group_reduce` and `KI.sub_group_scan` loop over the constant
-# sub-group width, and PoCL 7.2 miscompiles the unrolled shuffles after a branch with an early
-# exit (as bounds checks emit), like the native collectives below (fixed by pocl/pocl#2239,
-# JuliaPackaging/Yggdrasil#15001). Until `pocl_standalone_jll` includes the fix, loop to a
-# bound that isn't a constant, but is the same for all sub-groups of the work-group, which
-# PoCL needs (see `POCLBackend`): the smaller of the width and the work-group size. Reduce by
-# combining ranges of doubling length with `shfl_down`, skipping the lanes without a
-# work-item, and broadcast the result of the first lane.
-@inline function uniform_sub_group_bound()
-    sz = KI.get_local_size(Int32)
-    return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z)
-end
+# PoCL 7.2 miscompiles sub-group operations after a branch with an early exit (as bounds checks
+# emit), because WorkitemLoops gives the peeled first work-item its own copy of their scratch
+# memory: the native collectives below, and the unrolled shuffles of KernelInterface's
+# fallbacks of `KI.sub_group_reduce` and `KI.sub_group_scan`. `pocl_standalone_jll` includes
+# the fix (pocl/pocl#2239) since 7.2.1+1 (JuliaPackaging/Yggdrasil#15001).
+const POCL_REPLICA_FIX = pkgversion(cl.pocl_standalone_jll) >= v"7.2.1+1"
+
+@static if !POCL_REPLICA_FIX
+    # Without the fix, loop to a bound that isn't a constant, but is the same for all
+    # sub-groups of the work-group, which PoCL needs (see `POCLBackend`): the smaller of the
+    # width and the work-group size. Reduce by combining ranges of doubling length with
+    # `shfl_down`, skipping the lanes without a work-item, and broadcast the result of the
+    # first lane.
+    @inline function uniform_sub_group_bound()
+        sz = KI.get_local_size(Int32)
+        return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z)
+    end
 
-@device_override @inline function KI.sub_group_reduce(op, val)
-    lane = KI.get_sub_group_local_id(Int32)
-    sgsize = KI.get_sub_group_size(Int32)
-    offset = Int32(1)
-    bound = uniform_sub_group_bound()
-    while offset < bound
-        other = KI.shfl_down(val, offset)
-        if lane + offset <= sgsize
-            val = op(val, other)
+    @device_override @inline function KI.sub_group_reduce(op, val)
+        lane = KI.get_sub_group_local_id(Int32)
+        sgsize = KI.get_sub_group_size(Int32)
+        offset = Int32(1)
+        bound = uniform_sub_group_bound()
+        while offset < bound
+            other = KI.shfl_down(val, offset)
+            if lane + offset <= sgsize
+                val = op(val, other)
+            end
+            offset <<= 1
         end
-        offset <<= 1
+        return KI.shfl(val, 1)
     end
-    return KI.shfl(val, 1)
-end
 
-# likewise for `KI.sub_group_scan`
-@device_override @inline function KI.sub_group_scan(op, val)
-    lane = KI.get_sub_group_local_id(Int32)
-    offset = Int32(1)
-    bound = uniform_sub_group_bound()
-    while offset < bound
-        other = KI.shfl_up(val, offset)
-        if lane > offset
-            val = op(other, val)
+    # likewise for `KI.sub_group_scan`
+    @device_override @inline function KI.sub_group_scan(op, val)
+        lane = KI.get_sub_group_local_id(Int32)
+        offset = Int32(1)
+        bound = uniform_sub_group_bound()
+        while offset < bound
+            other = KI.shfl_up(val, offset)
+            if lane > offset
+                val = op(other, val)
+            end
+            offset <<= 1
         end
-        offset <<= 1
+        return val
     end
-    return val
 end
 
 # Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and
 # floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
-# differently from Julia's).
-#
-# Disabled until `pocl_standalone_jll` includes pocl/pocl#2239 (JuliaPackaging/Yggdrasil#15001):
-# in PoCL 7.2, a collective after
-# a branch with an early exit (as bounds checks emit) gets wrong values, because WorkitemLoops
-# gives the peeled first work-item its own copy of the collective's scratch memory. Until
-# then, `KI.sub_group_reduce` and `KI.sub_group_scan` use KernelInterface's fallbacks.
-const NATIVE_COLLECTIVES = false
+# differently from Julia's), with the fix for PoCL's peeling.
+const NATIVE_COLLECTIVES = POCL_REPLICA_FIX
 
 @static if NATIVE_COLLECTIVES
     const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}

From 759e579d6ea7058837682b18bf3c0b8491a37fba Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 6 Oct 2026 14:08:43 +0200
Subject: [PATCH 16/18] Address review: SPIRVIntrinsics 1.3, match_any on PoCL,
 small fixes

- Require the released SPIRVIntrinsics 1.3 instead of the deleted
  `vc/subgroup-votes` branch of OpenCL.jl, in `[sources]` and in CI.
  The OpenCL.jl job installs 1.3 rather than the ka-0.10 copy (1.1.4).
- The fallback of `KI.sub_group_match_any` takes a step per distinct
  value, which differs between the sub-groups of a work-group. PoCL needs
  them to take the same steps, so override it with a loop to a bound that
  is the same for the whole work-group. Test it with several sub-groups.
- Merge `NATIVE_COLLECTIVES` into `POCL_REPLICA_FIX`.
- `supports_shuffle` of a type without fields follows `UInt32` instead of
  being `true` on every backend.
- Fix the sizes of primitive types in the `shfl` docstring and the
  `supports_shuffle` entry of the backend table.

Assisted-by: Claude Code
---
 .buildkite/pipeline.yml               | 25 +++-----
 .github/workflows/ci.yml              |  5 +-
 Project.toml                          |  4 +-
 docs/src/kernelinterface.md           |  2 +-
 lib/KernelInterface/src/device.jl     |  4 +-
 lib/KernelInterface/src/host.jl       |  5 +-
 lib/KernelInterface/test/interface.jl | 36 +++++++++++
 src/pocl/backend.jl                   | 86 ++++++++++++++++-----------
 test/runtests.jl                      |  4 +-
 9 files changed, 107 insertions(+), 64 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 189920191..9edf96667 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -24,9 +24,7 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
-                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing CUDA")
                 using Pkg
                 url="https://github.com/JuliaGPU/CUDA.jl"
@@ -66,9 +64,7 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
-                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing Metal")
                 using Pkg
                 Pkg.add([(; url="https://github.com/JuliaGPU/Metal.jl", rev="ka-0.10")])' || exit 3
@@ -104,9 +100,7 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
-                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing oneAPI")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -145,9 +139,7 @@ steps:
                 using Pkg
                 println("--- :julia: Instantiating project")
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
-                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing AMDGPU")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -186,14 +178,13 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
-                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
       julia -e 'println("--- :julia: Developing OpenCL")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="ka-0.10")
-                # TODO: develop the ka-0.10 copy again once JuliaGPU/OpenCL.jl#526 is merged
-                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
+                # not the ka-0.10 copy of SPIRVIntrinsics: KernelAbstractions needs the
+                # sub-group votes and collectives of SPIRVIntrinsics 1.3
+                Pkg.add(name="SPIRVIntrinsics", version="1.3")' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 47938dfd7..bdccba98d 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -56,12 +56,9 @@ jobs:
       # applies to the test run only, not to `Pkg.develop`. Do this before building, which
       # already resolves the environment (and fails if this repo requires an unregistered
       # version of KernelInterface).
-      #
-      # The same applies to the SPIRVIntrinsics branch in `[sources]`.
-      # TODO: remove that once JuliaGPU/OpenCL.jl#526 is released
       - name: Dev KernelInterface
         shell: bash
-        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface"); Pkg.add(; url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")'
+        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")'
       # Nightly tracks the next Julia release; don't fail CI on it
       - uses: julia-actions/julia-buildpkg@v1
         continue-on-error: ${{ matrix.version == 'nightly' }}
diff --git a/Project.toml b/Project.toml
index ac1763297..9fe440915 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,8 +31,6 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [sources]
 KernelInterface = {path = "lib/KernelInterface"}
-# TODO: remove once JuliaGPU/OpenCL.jl#526 is released
-SPIRVIntrinsics = {url = "https://github.com/JuliaGPU/OpenCL.jl", rev = "vc/subgroup-votes", subdir = "lib/intrinsics"}
 
 [extensions]
 LinearAlgebraExt = "LinearAlgebra"
@@ -53,7 +51,7 @@ Printf = "<0.0.1, 1"
 Random = "1"
 Random123 = "1.7.1"
 RandomNumbers = "1.6.0"
-SPIRVIntrinsics = "1.1.3"
+SPIRVIntrinsics = "1.3"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_Tools_jll = "2024.4, 2025.1"
 SparseArrays = "<0.0.1, 1.6"
diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index 9950833f2..93a1431c3 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -74,7 +74,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Execution** | [`synchronize`](@ref) (cooperative) | [`record_event`](@ref)/[`wait_event`](@ref) (synchronize), [`priority!`](@ref) (no-op) |
 | **Devices** | with more than one device: [`ndevices`](@ref), [`device`](@ref), [`device!`](@ref), `device(backend, A)` | all four (a single device) |
 | **Queries** | [`max_work_group_size`](@ref) (for the backend and for a kernel), [`max_work_group_dims`](@ref), [`max_num_groups`](@ref) | [`launch_configuration`](@ref) (the limit), [`multiprocessor_count`](@ref) (0), [`functional`](@ref) (`missing`), [`versioninfo`](@ref) |
-| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) |
+| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
 | **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index 405d3caa2..df7c2d89e 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -285,8 +285,8 @@ but they may read from different lanes.
 `shfl` exchanges values, not memory: it is not a memory fence.
 
 Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a
-backend supports natively, that includes other primitive types of 1, 2 or a multiple of 4
-bytes (e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, which are shuffled
+backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes
+(e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, which are shuffled
 as `UInt32` words, and `isbits` structs and tuples of supported types, which are shuffled
 field by field.
 
diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl
index 2952d1d3c..7fcf80143 100644
--- a/lib/KernelInterface/src/host.jl
+++ b/lib/KernelInterface/src/host.jl
@@ -284,7 +284,8 @@ Whether kernels on the active device support the shuffles [`shfl`](@ref),
 [`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`.
 
 Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and `isbits`
-structs and tuples if all of their fields are.
+structs and tuples if all of their fields are (and, for those without fields, if `UInt32`
+is).
 
 !!! note
     Backend implementations **must** implement this function for the primitive types they
@@ -295,6 +296,8 @@ structs and tuples if all of their fields are.
 function supports_shuffle(backend::Backend, ::Type{T}) where {T}
     isprimitivetype(T) && return shuffle_as_words(T) && supports_shuffle(backend, UInt32)
     (isbitstype(T) && !isprimitivetype(T)) || return false
+    # not `true` for a type without fields on a backend without shuffles
+    fieldcount(T) == 0 && return supports_shuffle(backend, UInt32)
     return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T))
 end
 
diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl
index c6ca72316..a3942d81a 100644
--- a/lib/KernelInterface/test/interface.jl
+++ b/lib/KernelInterface/test/interface.jl
@@ -397,6 +397,13 @@ function reduce_scan_multi_kernel(red, scan, sgs, a)
     return
 end
 
+# `sub_group_match_any` in a 1-D work-group of several sub-groups, the last one partial
+function match_any_multi_kernel(out, vals)
+    i = KI.get_local_id().x
+    @inbounds out[i] = KI.sub_group_match_any(vals[i])
+    return
+end
+
 struct FallbackStruct
     flag::Bool
     c::Char
@@ -588,6 +595,29 @@ function reduce_scan_multi_testsuite(backend, AT, sg_size, ::Type{T}) where {T}
     return
 end
 
+# sub-groups with different numbers of distinct values, which the fallback takes different
+# numbers of steps for
+function match_any_multi_testsuite(backend, AT, sg_size)
+    n = 2 * sg_size + 5
+    vals = Int32[i <= sg_size ? i : (i <= 2 * sg_size ? 7 : i % 3) for i in 1:n]
+    out = AT(zeros(UInt64, n))
+    kernel = KI.@launch backend launch = false match_any_multi_kernel(out, AT(vals))
+    n <= KI.max_work_group_size(kernel) || return
+    kernel(out, AT(vals); workgroupsize = n)
+    KI.synchronize(backend)
+    out = Array(out)
+    # 1-D work-groups form sub-groups from consecutive work-items
+    for i in 1:n
+        base = (i - 1) ÷ sg_size * sg_size
+        expected = UInt64(0)
+        for j in (base + 1):min(base + sg_size, n)
+            vals[j] === vals[i] && (expected |= UInt64(1) << (j - 1 - base))
+        end
+        @test out[i] == expected
+    end
+    return
+end
+
 function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
     @testset "shuffles of other types" begin
         # primitive types that backends need not support natively are shuffled as words
@@ -660,6 +690,12 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size)
         end
     end
 
+    if sg_size <= 64
+        @testset "sub_group_match_any, several sub-groups" begin
+            match_any_multi_testsuite(backend, AT, sg_size)
+        end
+    end
+
     @testset "sub_group_reduce and sub_group_scan, several sub-groups" begin
         reduce_scan_multi_testsuite(backend, AT, sg_size, Int32)
         KI.supports_shuffle(backend, Float32) && reduce_scan_multi_testsuite(backend, AT, sg_size, Float32)
diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index dd430ea9f..f1bd20c07 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -277,6 +277,35 @@ end
     return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32)
 end
 
+# The smaller of the width and the work-group size: a bound for loops over the lanes that isn't
+# a constant, but is the same for all sub-groups of the work-group, which PoCL needs (see
+# `POCLBackend`).
+@inline function uniform_sub_group_bound()
+    sz = KI.get_local_size(Int32)
+    return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z)
+end
+
+# KernelInterface's fallback of `KI.sub_group_match_any` takes a step per distinct value, which
+# differs between the sub-groups of a work-group. Take the same number of steps in all of them,
+# as PoCL needs.
+@device_override @inline function KI.sub_group_match_any(val)
+    remaining = KI.sub_group_ballot(true)
+    mask = zero(UInt64)
+    step = Int32(0)
+    bound = uniform_sub_group_bound()
+    while step < bound
+        done = remaining == zero(UInt64)
+        leader = ifelse(done, Int32(1), trailing_zeros(remaining) % Int32 + Int32(1))
+        other = KI.shfl(val, leader)
+        same = !done & (val === other)
+        group = KI.sub_group_ballot(same)
+        mask = ifelse(same, group, mask)
+        remaining &= ~group
+        step += Int32(1)
+    end
+    return mask
+end
+
 # PoCL 7.2 miscompiles sub-group operations after a branch with an early exit (as bounds checks
 # emit), because WorkitemLoops gives the peeled first work-item its own copy of their scratch
 # memory: the native collectives below, and the unrolled shuffles of KernelInterface's
@@ -284,17 +313,30 @@ end
 # the fix (pocl/pocl#2239) since 7.2.1+1 (JuliaPackaging/Yggdrasil#15001).
 const POCL_REPLICA_FIX = pkgversion(cl.pocl_standalone_jll) >= v"7.2.1+1"
 
-@static if !POCL_REPLICA_FIX
-    # Without the fix, loop to a bound that isn't a constant, but is the same for all
-    # sub-groups of the work-group, which PoCL needs (see `POCLBackend`): the smaller of the
-    # width and the work-group size. Reduce by combining ranges of doubling length with
-    # `shfl_down`, skipping the lanes without a work-item, and broadcast the result of the
-    # first lane.
-    @inline function uniform_sub_group_bound()
-        sz = KI.get_local_size(Int32)
-        return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z)
-    end
+@static if POCL_REPLICA_FIX
+    # Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers
+    # and floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign
+    # of zero differently from Julia's).
+    const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
+    const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
+
+    @device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
+        SPIRVIntrinsics.sub_group_reduce_add(val)
+    @device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_reduce_min(val)
+    @device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_reduce_max(val)
 
+    @device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
+    @device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
+    @device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
+        SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
+else
+    # Without the fix, loop to `uniform_sub_group_bound()` rather than to the constant width.
+    # Reduce by combining ranges of doubling length with `shfl_down`, skipping the lanes
+    # without a work-item, and broadcast the result of the first lane.
     @device_override @inline function KI.sub_group_reduce(op, val)
         lane = KI.get_sub_group_local_id(Int32)
         sgsize = KI.get_sub_group_size(Int32)
@@ -326,30 +368,6 @@ const POCL_REPLICA_FIX = pkgversion(cl.pocl_standalone_jll) >= v"7.2.1+1"
     end
 end
 
-# Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and
-# floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
-# differently from Julia's), with the fix for PoCL's peeling.
-const NATIVE_COLLECTIVES = POCL_REPLICA_FIX
-
-@static if NATIVE_COLLECTIVES
-    const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
-    const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
-
-    @device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
-        SPIRVIntrinsics.sub_group_reduce_add(val)
-    @device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_reduce_min(val)
-    @device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_reduce_max(val)
-
-    @device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
-    @device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
-    @device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
-end
-
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 2f59678b6..1aa1e0dc7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -152,13 +152,13 @@ function sub_group_reduce_kernel(out, x)
     return
 end
 @testset "POCL native sub-group collectives" begin
-    for (T, native) in ((Int32, true), (Float32, true))
+    for T in (Int32, Float32)
         x, out = ones(T, 32), zeros(T, 32)
         ir = sprint() do io
             @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x)
         end
         @test all(==(32), out)
-        @test occursin("sub_group_reduce_add", ir) == (native && POCL.POCLKernels.NATIVE_COLLECTIVES)
+        @test occursin("sub_group_reduce_add", ir) == POCL.POCLKernels.POCL_REPLICA_FIX
     end
 end
 

From 874350640b0cdacd32bc4cfe6d65064fa35db4a3 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 6 Oct 2026 14:16:47 +0200
Subject: [PATCH 17/18] POCL: drop the sub-group fallbacks for PoCL without the
 replica fix

pocl_standalone_jll 7.2.2, which main requires now, includes the fix
for PoCL 7.2's peeling of the first work-item (pocl/pocl#2239), so the
native reductions and scans are always used, and the loops to a uniform
bound that replaced them without the fix can't be reached anymore.

Assisted-by: Claude Code
---
 src/pocl/backend.jl | 81 +++++++++++----------------------------------
 test/runtests.jl    |  2 +-
 2 files changed, 21 insertions(+), 62 deletions(-)

diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl
index f1bd20c07..6741c4ebd 100644
--- a/src/pocl/backend.jl
+++ b/src/pocl/backend.jl
@@ -306,67 +306,26 @@ end
     return mask
 end
 
-# PoCL 7.2 miscompiles sub-group operations after a branch with an early exit (as bounds checks
-# emit), because WorkitemLoops gives the peeled first work-item its own copy of their scratch
-# memory: the native collectives below, and the unrolled shuffles of KernelInterface's
-# fallbacks of `KI.sub_group_reduce` and `KI.sub_group_scan`. `pocl_standalone_jll` includes
-# the fix (pocl/pocl#2239) since 7.2.1+1 (JuliaPackaging/Yggdrasil#15001).
-const POCL_REPLICA_FIX = pkgversion(cl.pocl_standalone_jll) >= v"7.2.1+1"
-
-@static if POCL_REPLICA_FIX
-    # Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers
-    # and floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign
-    # of zero differently from Julia's).
-    const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
-    const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
-
-    @device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
-        SPIRVIntrinsics.sub_group_reduce_add(val)
-    @device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_reduce_min(val)
-    @device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_reduce_max(val)
-
-    @device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
-    @device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
-    @device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
-        SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
-else
-    # Without the fix, loop to `uniform_sub_group_bound()` rather than to the constant width.
-    # Reduce by combining ranges of doubling length with `shfl_down`, skipping the lanes
-    # without a work-item, and broadcast the result of the first lane.
-    @device_override @inline function KI.sub_group_reduce(op, val)
-        lane = KI.get_sub_group_local_id(Int32)
-        sgsize = KI.get_sub_group_size(Int32)
-        offset = Int32(1)
-        bound = uniform_sub_group_bound()
-        while offset < bound
-            other = KI.shfl_down(val, offset)
-            if lane + offset <= sgsize
-                val = op(val, other)
-            end
-            offset <<= 1
-        end
-        return KI.shfl(val, 1)
-    end
-
-    # likewise for `KI.sub_group_scan`
-    @device_override @inline function KI.sub_group_scan(op, val)
-        lane = KI.get_sub_group_local_id(Int32)
-        offset = Int32(1)
-        bound = uniform_sub_group_bound()
-        while offset < bound
-            other = KI.shfl_up(val, offset)
-            if lane > offset
-                val = op(other, val)
-            end
-            offset <<= 1
-        end
-        return val
-    end
-end
+# Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and
+# floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero
+# differently from Julia's). They need the fix of PoCL 7.2's peeling of the first work-item
+# (pocl/pocl#2239), which `pocl_standalone_jll` includes since 7.2.1+1.
+const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64}
+const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64}
+
+@device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) =
+    SPIRVIntrinsics.sub_group_reduce_add(val)
+@device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_reduce_min(val)
+@device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_reduce_max(val)
+
+@device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_add(val)
+@device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_min(val)
+@device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) =
+    SPIRVIntrinsics.sub_group_scan_inclusive_max(val)
 
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
diff --git a/test/runtests.jl b/test/runtests.jl
index 1aa1e0dc7..fc71a101a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -158,7 +158,7 @@ end
             @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x)
         end
         @test all(==(32), out)
-        @test occursin("sub_group_reduce_add", ir) == POCL.POCLKernels.POCL_REPLICA_FIX
+        @test occursin("sub_group_reduce_add", ir)
     end
 end
 

From 64ac66815480e0ed8ccef659280e58c85e1f85f6 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 6 Oct 2026 14:16:47 +0200
Subject: [PATCH 18/18] KernelInterface: shuffle other 8-byte primitive types
 as a UInt64

Primitive types that a backend doesn't shuffle natively were always
split into UInt32 words, so a `Ptr` or another 64-bit primitive type
took two shuffles even on backends with native 64-bit shuffles. Shuffle
them as a `UInt64` instead, which a backend without native 64-bit
shuffles splits into two words as before, and 16-byte values as two
`UInt64` words. Check the generated code on PoCL.

Assisted-by: Claude Code
---
 docs/src/kernelinterface.md       |  2 +-
 lib/KernelInterface/src/device.jl | 24 +++++++++++++++---------
 test/codegen_checks.jl            | 24 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md
index 93a1431c3..90000f687 100644
--- a/docs/src/kernelinterface.md
+++ b/docs/src/kernelinterface.md
@@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
-| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
+| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (field by field) |
 
 Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones),
 the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch),
diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl
index df7c2d89e..0e1d75099 100644
--- a/lib/KernelInterface/src/device.jl
+++ b/lib/KernelInterface/src/device.jl
@@ -271,7 +271,7 @@ localmemory(::Type{T}, ::Val) where {T} =
 
 # Shuffles exchange values between the work-items of a sub-group. Backends implement them for
 # the primitive types for which `supports_shuffle` returns `true`. The fallbacks below shuffle
-# other primitive types as `UInt32` words, and other `isbits` types field by field.
+# other primitive types as unsigned words, and other `isbits` types field by field.
 
 """
     shfl(val::T, lane::Integer)::T
@@ -286,9 +286,10 @@ but they may read from different lanes.
 
 Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a
 backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes
-(e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, which are shuffled
-as `UInt32` words, and `isbits` structs and tuples of supported types, which are shuffled
-field by field.
+(e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, and `isbits` structs and
+tuples of supported types, which are shuffled field by field. Primitive types of up to 4
+bytes are shuffled as a `UInt32`, and larger ones as `UInt64` words, each of which is shuffled
+natively if the backend supports `UInt64`, and as two `UInt32` words otherwise.
 
 !!! note
     Backend implementations **must** implement this for the primitive types they support
@@ -386,23 +387,28 @@ shfl_unsupported(T) = throw(
     )
 )
 
-# Whether a primitive type that a backend doesn't support natively can be shuffled as `UInt32`
+# Whether a primitive type that a backend doesn't support natively can be shuffled as unsigned
 # words
 shuffle_as_words(T) = T !== UInt32 && sizeof(T) in (1, 2, 4, 8, 16)
 
 # The unsigned integer type of the size of a primitive type `T`
 const word_types = Dict(1 => UInt8, 2 => UInt16, 4 => UInt32, 8 => UInt64, 16 => UInt128)
 
-# Shuffle a primitive value as `UInt32` words: smaller values are zero-extended, larger ones
-# split into words.
+# Shuffle a primitive value as unsigned words: values of up to 4 bytes are zero-extended to a
+# `UInt32`. Other 8-byte values are shuffled as a `UInt64`, which the backend may support
+# natively, and a `UInt64` it doesn't is split into two `UInt32` words. 16-byte values are
+# split into two `UInt64` words.
 @inline @generated function shfl_words(f, val::T) where {T}
     shuffle_as_words(T) || return :(shfl_unsupported($T))
     U = word_types[sizeof(T)]
     if sizeof(T) <= 4
         return :(reinterpret($T, f(reinterpret($U, val) % UInt32) % $U))
+    elseif sizeof(T) == 8 && T !== UInt64
+        return :(reinterpret($T, f(reinterpret(UInt64, val))))
     end
-    n = sizeof(T) ÷ 4
-    words = (:((f((bits >> $(32 * (i - 1))) % UInt32) % $U) << $(32 * (i - 1))) for i in 1:n)
+    W = sizeof(T) == 8 ? UInt32 : UInt64
+    nbits = 8 * sizeof(W)
+    words = (:((f((bits >> $(nbits * (i - 1))) % $W) % $U) << $(nbits * (i - 1))) for i in 1:2)
     return quote
         bits = reinterpret($U, val)
         return reinterpret($T, |($(words...)))
diff --git a/test/codegen_checks.jl b/test/codegen_checks.jl
index 6db27c38c..f970d20a9 100644
--- a/test/codegen_checks.jl
+++ b/test/codegen_checks.jl
@@ -16,6 +16,7 @@
 using FileCheck
 using KernelAbstractions
 using KernelAbstractions: @atomic
+import KernelAbstractions.KernelInterface as KI
 using Test
 
 import KernelAbstractions.POCL: @device_code_llvm
@@ -57,6 +58,14 @@ end
     @inbounds A[I] = I
 end
 
+# a 64-bit primitive type that no backend shuffles natively
+primitive type CodegenBits64 64 end
+
+@kernel function codegen_shfl_bits64(A)
+    I = @index(Global, Linear)
+    @inbounds A[I] = reinterpret(UInt64, KI.shfl(reinterpret(CodegenBits64, A[I]), 1))
+end
+
 # `@inbounds` is only honoured under `--check-bounds=auto`; several checks below assert
 # that it removes code, so refuse to run under anything else rather than fail obscurely.
 if Base.JLOptions().check_bounds != 0
@@ -181,6 +190,21 @@ end
         end
     end
 
+    # A 64-bit primitive type that the backend doesn't shuffle natively uses its native 64-bit
+    # shuffle (`__spirv_GroupNonUniformShuffle` mangled for an `i64` value is `...jmj`, for an
+    # `i32` value `...jjj`), rather than being split into two 32-bit words.
+    @testset "shuffles of other primitive types" begin
+        A = KernelAbstractions.zeros(backend, UInt64, 32)
+        @test @filecheck implicit_check_not = "GroupNonUniformShufflejjj" begin
+            @check "define spir_kernel void @{{.*}}gpu_codegen_shfl_bits64"
+            @check "GroupNonUniformShufflejmj"
+            @check_not "GroupNonUniformShuffle"
+            @check "ret void"
+            @device_code_llvm debuginfo = :none codegen_shfl_bits64(backend, 32)(A, ndrange = 32)
+            KernelAbstractions.synchronize(backend)
+        end
+    end
+
     # `@print` lowers to a single variadic printf call, not to one call per argument.
     @testset "print" begin
         @test @filecheck begin