From 9332535d1935a6bfec9a9a5eefb86f7d062f2a3d Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:03:23 +0200
Subject: [PATCH 01/13] Give every static local memory call site its own
 allocation

`alloc_special` declared local memory as an external global named after its
`id`, so call sites using the same `id` shared one allocation. Define it with
internal linkage and an `undef` initializer instead, so that each call site
gets its own, as KernelInterface's `localmemory` requires: it has no `id`.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 src/device/gcn/memory_static.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/device/gcn/memory_static.jl b/src/device/gcn/memory_static.jl
index 4962630f6..63d137497 100644
--- a/src/device/gcn/memory_static.jl
+++ b/src/device/gcn/memory_static.jl
@@ -20,8 +20,10 @@
     gv = GlobalVariable(current_module(builder), gv_typ, name, as)
     if len > 0
         if as == AS.Local
-            gv.linkage = LLVM.Linkage.External
-            # NOTE: Backend doesn't support initializer for local AS
+            # every call site gets its own allocation. local memory can't be
+            # initialized, so use an `undef` initializer, which the backend accepts.
+            gv.linkage = LLVM.Linkage.Internal
+            gv.initializer = UndefValue(gv_typ)
         elseif as == AS.Private
             gv.linkage = LLVM.Linkage.Internal
             gv.initializer = null(gv_typ)

From 72d6b71bb65a7450c139e4952e69e66ec224804c Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:03:51 +0200
Subject: [PATCH 02/13] Add `HIP.max_workgroup_dims`

The maximum number of work-items along each dimension of a workgroup, from the
device's `maxThreadsDim`. It is cached in `HIPDevice`, since KernelInterface
queries it on every launch whose workgroup size it chooses.
---
 docs/src/api/devices.md |  1 +
 src/hip/device.jl       | 11 ++++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/docs/src/api/devices.md b/docs/src/api/devices.md
index cb9d28135..ab2b33f86 100644
--- a/docs/src/api/devices.md
+++ b/docs/src/api/devices.md
@@ -56,6 +56,7 @@ AMDGPU.device_id!
 ```@docs
 AMDGPU.HIP.name
 AMDGPU.HIP.wavefrontsize
+AMDGPU.HIP.max_workgroup_dims
 AMDGPU.HIP.gcn_arch
 AMDGPU.HIP.device_id
 AMDGPU.HIP.properties
diff --git a/src/hip/device.jl b/src/hip/device.jl
index accd4d069..6cfb62e42 100644
--- a/src/hip/device.jl
+++ b/src/hip/device.jl
@@ -3,6 +3,7 @@ struct HIPDevice
     device_id::Cint
     gcn_arch::String
     wavefrontsize::Cint
+    max_workgroup_dims::NTuple{3, Int}
 end
 
 const DEFAULT_DEVICE = Ref{Union{Nothing, HIPDevice}}(nothing)
@@ -18,7 +19,8 @@ function HIPDevice(device_id::Integer)
     gcn_arch = unsafe_string(pointer([props.gcnArchName...]))
 
     wavefrontsize = props.warpSize
-    HIPDevice(device_ref[], device_id, gcn_arch, wavefrontsize)
+    max_workgroup_dims = Int.(props.maxThreadsDim)
+    HIPDevice(device_ref[], device_id, gcn_arch, wavefrontsize, max_workgroup_dims)
 end
 
 """
@@ -36,6 +38,13 @@ Get size of the wavefront. AMD GPUs support either 32 or 64.
 """
 wavefrontsize(d::HIPDevice)::Cint = d.wavefrontsize
 
+"""
+    max_workgroup_dims(d::HIPDevice)::NTuple{3, Int}
+
+Get the maximum number of work-items along each dimension of a workgroup.
+"""
+max_workgroup_dims(d::HIPDevice)::NTuple{3, Int} = d.max_workgroup_dims
+
 """
     gcn_arch(d::HIPDevice)::String
 

From 06b723b7671a2bfecee022263ee6c5856a492bfa Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Tue, 6 Oct 2026 13:33:31 +0200
Subject: [PATCH 03/13] Implement KernelInterface, and support
 KernelAbstractions 0.10

KernelAbstractions 0.10 builds on KernelInterface, which defines what a back
end provides: memory and device management, compiling and launching kernels,
and the device-side intrinsics. KernelAbstractions then launches `@kernel`
kernels itself on any KernelInterface back end, which replaces AMDGPU's copy of
that launch path: partitioning the ndrange, building the kernel's context, and
tuning the workgroup size.

`ROCBackend` now implements KernelInterface, which AMDGPU depends on instead of
on KernelAbstractions. What KernelAbstractions still needs from a back end
moves to an extension: the `MArray` behind `@private`, the Adapt rule for
`@Const`, and the `maxthreads` hint for a static workgroup size.

`KI.launch` receives the kernel arguments as a tuple, and passes it on to the
`HIPKernel` without splatting it. It rejects `groupsize` and `gridsize`, which
would override the launch geometry that KernelInterface validated.
`KI.copyto!` accepts dense arrays and contiguous views of them.

`KI.kernel_function` receives the callable unconverted. It compiles its
converted form, and the kernel keeps the original alive, since the converted
form only holds pointers to the arrays a closure captures. `KI.launch` converts
it again for the launch's stream, like the arguments, so that those arrays are
made available to the stream too.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 Project.toml                       |   7 +-
 ext/AMDGPUKernelAbstractionsExt.jl |  42 +++++
 src/AMDGPU.jl                      |   2 +
 src/ROCKernels.jl                  | 269 ++++++++++++++---------------
 src/utils.jl                       |   3 +-
 test/Project.toml                  |   1 +
 test/kernelabstractions_tests.jl   |  56 +++++-
 test/kernelinterface_tests.jl      |  65 +++++++
 8 files changed, 303 insertions(+), 142 deletions(-)
 create mode 100644 ext/AMDGPUKernelAbstractionsExt.jl
 create mode 100644 test/kernelinterface_tests.jl

diff --git a/Project.toml b/Project.toml
index 41229fed4..0b1278866 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,7 +17,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -38,12 +38,14 @@ UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
+AMDGPUKernelAbstractionsExt = "KernelAbstractions"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
@@ -60,7 +62,8 @@ ExprTools = "0.1"
 GPUArrays = "12.1"
 GPUCompiler = "2.13"
 GPUToolbox = "3.3.1"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 LLVMDowngrader_jll = "0.12"
 Libdl = "1"
diff --git a/ext/AMDGPUKernelAbstractionsExt.jl b/ext/AMDGPUKernelAbstractionsExt.jl
new file mode 100644
index 000000000..177628f28
--- /dev/null
+++ b/ext/AMDGPUKernelAbstractionsExt.jl
@@ -0,0 +1,42 @@
+module AMDGPUKernelAbstractionsExt
+
+import AMDGPU
+import AMDGPU.Device: @device_override
+using AMDGPU: GPUArrays, ROCBackend
+
+import Adapt
+import KernelAbstractions as KA
+import LLVM
+
+using StaticArraysCore: MArray
+
+Adapt.adapt_storage(::KA.CPU, a::Union{AMDGPU.ROCArray, GPUArrays.AbstractGPUSparseArray}) =
+    Adapt.adapt(Array, a)
+
+## kernel launch
+
+# KernelAbstractions launches kernels through the KernelInterface back-end; tell the
+# compiler about statically sized workgroups
+function KA.compiler_options(obj::KA.Kernel{ROCBackend})
+    if KA.workgroupsize(obj) <: KA.StaticSize
+        return (; maxthreads = prod(KA.get(KA.workgroupsize(obj))))
+    else
+        return (;)
+    end
+end
+
+## scratch memory
+
+@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
+    MArray{Tuple{Dims...}, T}(undef)
+end
+
+## other
+
+# `@Const` arrays are read through the constant address space
+function Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) where T
+    ptr = LLVM.Interop.addrspacecast(Core.LLVMPtr{T,AMDGPU.Device.AS.Constant}, a.ptr)
+    AMDGPU.ROCDeviceArray(a.dims, ptr)
+end
+
+end
diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 799d958eb..2b97bc11e 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -11,6 +11,7 @@ using LLVM
 using Preferences
 using Printf
 
+import KernelInterface
 import UnsafeAtomics
 import Atomix
 import Atomix: @atomic, @atomicswap, @atomicreplace
@@ -132,6 +133,7 @@ include("dnn/MIOpen.jl")
 
 include("random.jl")
 
+# KernelInterface
 include("ROCKernels.jl")
 import .ROCKernels: ROCBackend
 export ROCBackend
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 2513d1c2f..a1f69186a 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -4,205 +4,198 @@ export ROCBackend
 
 import AMDGPU
 import AMDGPU.Device: @device_override
-using AMDGPU: GPUArrays, rocSPARSE
+using AMDGPU: GPUArrays, rocSPARSE, HIP, Device, rocconvert, hipfunction
 
 import Adapt
-import KernelAbstractions as KA
-import LLVM
-
-using StaticArraysCore: MArray
+import KernelInterface as KI
 
 """
-    ROCBackend <: KernelAbstractions.GPU
+    ROCBackend <: KernelInterface.Backend
 
-KernelAbstractions backend that executes kernels on an AMD GPU via AMDGPU.jl.
+KernelInterface backend that executes kernels on an AMD GPU via AMDGPU.jl.
 Pass `ROCBackend()` to a KernelAbstractions kernel to run it on the GPU, or
-obtain it from an array with `KernelAbstractions.get_backend(::ROCArray)`.
+obtain it from an array with `KernelInterface.get_backend(::ROCArray)`.
+
+Printing from a kernel (`KernelAbstractions.@print`) is not supported: it does nothing.
 """
-struct ROCBackend <: KA.GPU end
+struct ROCBackend <: KI.Backend end
 
-KA.functional(::ROCBackend) = AMDGPU.functional()
-KA.ndevices(::ROCBackend) = AMDGPU.HIP.ndevices()
-KA.device(::ROCBackend) = AMDGPU.device_id()
-function KA.device!(kab::ROCBackend, id::Int)
-    (0 < id <= KA.ndevices(kab)) || throw(ArgumentError("Device id $id out of bounds."))
+KI.functional(::ROCBackend) = AMDGPU.functional()
+KI.ndevices(::ROCBackend) = AMDGPU.HIP.ndevices()
+KI.device(::ROCBackend) = AMDGPU.device_id()
+function KI.device!(kab::ROCBackend, id::Int)
+    (0 < id <= KI.ndevices(kab)) || throw(ArgumentError("Device id $id out of bounds."))
     AMDGPU.device_id!(id)
     return
 end
+KI.device(::ROCBackend, A::AMDGPU.ROCArray) = AMDGPU.device_id(AMDGPU.device(A))
 
 Adapt.adapt_storage(::ROCBackend, a::AbstractArray) = Adapt.adapt(AMDGPU.ROCArray, a)
 Adapt.adapt_storage(::ROCBackend, a::Union{AMDGPU.ROCArray, GPUArrays.AbstractGPUSparseArray}) = a
-Adapt.adapt_storage(::KA.CPU, a::Union{AMDGPU.ROCArray, GPUArrays.AbstractGPUSparseArray}) =
-    Adapt.adapt(Array, a)
-function Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) where T
-    ptr = LLVM.Interop.addrspacecast(Core.LLVMPtr{T,AMDGPU.Device.AS.Constant}, a.ptr)
-    AMDGPU.ROCDeviceArray(a.dims, ptr)
-end
 
-KA.get_backend(::AMDGPU.ROCArray) = ROCBackend()
-KA.get_backend(::AMDGPU.rocSPARSE.ROCSparseVector) = ROCBackend()
-KA.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSC) = ROCBackend()
-KA.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSR) = ROCBackend()
+KI.get_backend(::AMDGPU.ROCArray) = ROCBackend()
+KI.get_backend(::AMDGPU.rocSPARSE.ROCSparseVector) = ROCBackend()
+KI.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSC) = ROCBackend()
+KI.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSR) = ROCBackend()
 
-KA.argconvert(::KA.Kernel{ROCBackend}, arg) = AMDGPU.rocconvert(arg)
-KA.synchronize(::ROCBackend) = AMDGPU.synchronize()
+KI.synchronize(::ROCBackend) = AMDGPU.synchronize()
 
-KA.unsafe_free!(x::AMDGPU.ROCArray) = AMDGPU.unsafe_free!(x)
-KA.allocate(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ROCArray{T}(undef, dims)
-KA.zeros(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.zeros(T, dims)
-KA.ones(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ones(T, dims)
+KI.supports_float64(::ROCBackend) = true
+KI.supports_atomics(::ROCBackend) = true
 
-function KA.priority!(::ROCBackend, priority::Symbol)
+function KI.priority!(::ROCBackend, priority::Symbol)
     priority ∉ (:high, :normal, :low) && error(
         "Priority `$priority` must be one of `:high`, `:normal`, `:low`.")
     AMDGPU.priority!(priority)
+    return
 end
 
-function KA.copyto!(::ROCBackend, A, B)
-    GC.@preserve A B begin
-        copyto!(A, 1, B, 1, length(A))
+## memory operations
+
+KI.unsafe_free!(x::AMDGPU.ROCArray) = AMDGPU.unsafe_free!(x)
+KI.allocate(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ROCArray{T}(undef, dims)
+
+# dense arrays, and contiguous views of them
+const ContiguousArray{T} = Union{DenseArray{T}, Base.FastContiguousSubArray{T, <:Any, <:DenseArray}}
+on_device(A::ContiguousArray) = parent(A) isa AMDGPU.ROCArray
+
+function KI.copyto!(::ROCBackend, A::ContiguousArray{T}, B::ContiguousArray{T}) where T
+    length(A) == length(B) ||
+        throw(ArgumentError("Arrays must have the same length, got $(length(A)) and $(length(B))"))
+    if isbitstype(T) && (on_device(A) || on_device(B))
+        # queued on the task's stream, after the work queued before it
+        GC.@preserve A B begin
+            AMDGPU.Mem.memcpy!(pointer(A), pointer(B), length(A) * AMDGPU.aligned_sizeof(T);
+                               stream=AMDGPU.stream())
+        end
+    else
+        # host-to-host copies, and bits unions, whose type tags are stored separately.
+        # queued work may still access host arrays (e.g. a copy from the device), so wait
+        # for it first.
+        on_device(A) || on_device(B) || AMDGPU.synchronize()
+        copyto!(A, B)
     end
-    return
+    return A
 end
+KI.copyto!(::ROCBackend, A, B) =
+    throw(ArgumentError("KernelInterface.copyto! only supports contiguous arrays of the same element type, got $(typeof(A)) and $(typeof(B))"))
 
-function KA.pagelock!(::ROCBackend, x::Array)
+function KI.pagelock!(::ROCBackend, x::Array)
     AMDGPU.Mem.pin(pointer(x), sizeof(x))
     return
 end
 
-function KA.launch_config(kernel::KA.Kernel{ROCBackend}, ndrange, workgroupsize)
-    if ndrange isa Integer
-        ndrange = (ndrange,)
-    end
-    if workgroupsize isa Integer
-        workgroupsize = (workgroupsize, )
-    end
-
-    # partition checked that the ndrange's agreed
-    if KA.ndrange(kernel) <: KA.StaticSize
-        ndrange = nothing
-    end
+## kernel launch
 
-    iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize &&
-        workgroupsize === nothing
-        # use ndrange as preliminary workgroupsize for autotuning
-        KA.partition(kernel, ndrange, ndrange)
-    else
-        KA.partition(kernel, ndrange, workgroupsize)
-    end
+KI.argconvert(::ROCBackend, arg) = rocconvert(arg)
 
-    return ndrange, workgroupsize, iterspace, dynamic
+# a compiled kernel, and the callable it was compiled from. the kernel only holds pointers
+# to the arrays the callable captures, so the callable has to be kept alive.
+struct ROCKernel{F, K <: AMDGPU.Runtime.HIPKernel}
+    f::F
+    kernel::K
 end
 
-threads_to_workgroupsize(threads, ndrange::Tuple) = _threads_to_workgroupsize(threads, 1, ndrange)
-_threads_to_workgroupsize(threads, total, ::Tuple{}) = ()
-function _threads_to_workgroupsize(threads, total, ndrange::Tuple)
-    x = min(div(threads, total), first(ndrange))
-    return (x, _threads_to_workgroupsize(threads, total * x, Base.tail(ndrange))...)
+function KI.kernel_function(backend::ROCBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
+    # kernels have to execute with the device's wavefront size, which is what `hipfunction`
+    # compiles for by default
+    ws = HIP.wavefrontsize(AMDGPU.device())
+    if haskey(kwargs, :wavefrontsize64) && kwargs[:wavefrontsize64] != (ws == 64)
+        throw(ArgumentError("`wavefrontsize64=$(kwargs[:wavefrontsize64])` conflicts with the wavefront size of the device, $ws"))
+    end
+    kernel = hipfunction(rocconvert(f), tt; name, kwargs...)
+    return KI.Kernel(backend, ROCKernel(f, kernel))
 end
 
-# forwards the arguments as a tuple, see `AMDGPU.Runtime.launch_tuple(::HIPFunction, ...)`
-(obj::KA.Kernel{ROCBackend})(args::Vararg{Any, N}) where N = launch_tuple(obj, args)
-Core.kwcall(kwargs::NamedTuple, obj::KA.Kernel{ROCBackend}, args::Vararg{Any, N}) where N =
-    launch_tuple(obj, args; kwargs...)
-
-# `(x, t...)`, without splatting
-@inline @generated prepend(x, t::Tuple) = :((x, $((:(t[$i]) for i in 1:fieldcount(t))...)))
-
-function launch_tuple(obj::KA.Kernel{ROCBackend}, args::Tuple; ndrange=nothing, workgroupsize=nothing)
-    ndrange, new_workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize)
-    ctx = KA.mkcontext(obj, ndrange, iterspace)
-    if KA.workgroupsize(obj) <: KA.StaticSize
-        maxthreads = prod(KA.get(KA.workgroupsize(obj)))
-    else
-        maxthreads = nothing
+function KI.launch(obj::KI.Kernel{ROCBackend}, groups::Dims{3}, items::Dims{3},
+                   args::Tuple; kwargs...)
+    # KernelInterface has validated the launch geometry
+    if haskey(kwargs, :groupsize) || haskey(kwargs, :gridsize)
+        throw(ArgumentError("KernelInterface kernels take `numgroups`, `workgroupsize` or `ndrange`, not `groupsize` or `gridsize`"))
     end
-    # `@roc` rebuilds a single splatted tuple as `(t...,)`, which Julia compiles to `t` itself
-    kernel = AMDGPU.@roc launch=false maxthreads=maxthreads obj.f(prepend(ctx, args)...)
-
-    # If dynamic, figure out the optimal groupsize automatically.
-    is_dynamic =
-        KA.workgroupsize(obj) <: KA.DynamicSize &&
-        isnothing(workgroupsize)
-    if is_dynamic
-        (; groupsize) = AMDGPU.launch_configuration(kernel)
-        new_workgroupsize = threads_to_workgroupsize(groupsize, ndrange)
-        iterspace, dynamic = KA.partition(obj, ndrange, new_workgroupsize)
-        ctx = KA.mkcontext(obj, ndrange, iterspace)
+    f = obj.kern.f
+    kernel = obj.kern.kernel
+    stream = get(kwargs, :stream, AMDGPU.stream())
+    GC.@preserve f begin
+        # convert the callable again, like the arguments, which makes the arrays it captures
+        # available to the stream
+        kernel = typeof(kernel)(rocconvert(f, stream), kernel.fun)
+        AMDGPU.Runtime.launch_tuple(kernel, args; groupsize=items, gridsize=groups, kwargs..., stream)
     end
-
-    nblocks = length(KA.blocks(iterspace))
-    nthreads = length(KA.workitems(iterspace))
-    nblocks == 0 && return
-
-    # what `kernel(ctx, args...; ...)` calls, without splatting
-    AMDGPU.Runtime.launch_tuple(kernel, prepend(ctx, args); groupsize=nthreads, gridsize=nblocks)
     return
 end
 
-function KA.mkcontext(kernel::KA.Kernel{ROCBackend}, _ndrange, iterspace)
-    metadata = KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace)
+function KI.max_work_group_size(kernel::KI.Kernel{ROCBackend})::Int
+    max_items = Ref{Cint}()
+    HIP.hipFuncGetAttribute(max_items, HIP.HIP_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, kernel.kern.kernel.fun)
+    return Int(max_items[])
+end
+function KI.launch_configuration(kernel::KI.Kernel{ROCBackend}; nitems::Union{Integer,Nothing}=nothing,
+                                 max_work_group_size::Integer=typemax(Int))
+    max_items = min(max_work_group_size, something(nitems, typemax(Int)),
+                    KI.max_work_group_size(kernel))
+    (; groupsize) = AMDGPU.launch_configuration(kernel.kern.kernel; max_block_size=max_items)
+    return (; workgroupsize=Int(min(groupsize, max_items)))
+end
+function KI.max_work_group_size(::ROCBackend)::Int
+    Int(HIP.attribute(AMDGPU.device(), HIP.hipDeviceAttributeMaxThreadsPerBlock))
+end
+# queried on every automatically-sized launch, so use the limits cached in the device
+KI.max_work_group_dims(::ROCBackend)::NTuple{3, Int} = HIP.max_workgroup_dims(AMDGPU.device())
+# HIP takes the grid size in workgroups, but the dispatch packet holds it in work-items
+# (as a UInt32 per dimension, which HIP checks), and the device code assumes workgroup
+# indices fit in an Int32 (see `Device._max_groups`). Report the number of workgroups
+# that can be launched with any valid workgroup size. HIP's `maxGridSize` isn't usable:
+# depending on the ROCm version it holds CUDA's block limits or the work-item limits.
+function KI.max_num_groups(backend::ROCBackend)::NTuple{3, Int}
+    dims = KI.max_work_group_dims(backend)
+    return ntuple(Val(3)) do d
+        Int(min(Device._max_groups[d], Device._max_grid_size[d] ÷ dims[d]))
+    end
 end
-function KA.mkcontext(kernel::KA.Kernel{ROCBackend}, I, _ndrange, iterspace, ::Dynamic) where Dynamic
-    metadata = KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
+function KI.multiprocessor_count(::ROCBackend)::Int
+    Int(HIP.attribute(AMDGPU.device(), HIP.hipDeviceAttributeMultiprocessorCount))
 end
 
-# Indexing.
-
-@device_override @inline function KA.__index_Local_Linear(ctx)
-    return AMDGPU.Device.threadIdx().x
-end
+## COV_EXCL_START
 
-@device_override @inline function KA.__index_Group_Linear(ctx)
-    return AMDGPU.Device.blockIdx().x
-end
+## indexing
 
-@device_override @inline function KA.__index_Global_Linear(ctx)
-    I =  @inbounds KA.expand(KA.__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x)
-    # TODO: This is unfortunate, can we get the linear index cheaper
-    @inbounds LinearIndices(KA.__ndrange(ctx))[I]
-end
+# computed with `% T`, which unlike `T(x)` has no error path
 
-@device_override @inline function KA.__index_Local_Cartesian(ctx)
-    @inbounds KA.workitems(KA.__iterspace(ctx))[AMDGPU.Device.threadIdx().x]
+@device_override @inline function KI.get_local_id(::Type{T}) where {T}
+    return (; x = Device.workitemIdx().x % T, y = Device.workitemIdx().y % T, z = Device.workitemIdx().z % T)
 end
 
-@device_override @inline function KA.__index_Group_Cartesian(ctx)
-    @inbounds KA.blocks(KA.__iterspace(ctx))[AMDGPU.Device.blockIdx().x]
+@device_override @inline function KI.get_group_id(::Type{T}) where {T}
+    return (; x = Device.workgroupIdx().x % T, y = Device.workgroupIdx().y % T, z = Device.workgroupIdx().z % T)
 end
 
-@device_override @inline function KA.__index_Global_Cartesian(ctx)
-    return @inbounds KA.expand(KA.__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x)
+@device_override @inline function KI.get_local_size(::Type{T}) where {T}
+    return (; x = Device.workgroupDim().x % T, y = Device.workgroupDim().y % T, z = Device.workgroupDim().z % T)
 end
 
-@device_override @inline function KA.__validindex(ctx)
-    if KA.__dynamic_checkbounds(ctx)
-        I = @inbounds KA.expand(KA.__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x)
-        return I in KA.__ndrange(ctx)
-    else
-        return true
-    end
+@device_override @inline function KI.get_num_groups(::Type{T}) where {T}
+    return (; x = Device.gridGroupDim().x % T, y = Device.gridGroupDim().y % T, z = Device.gridGroupDim().z % T)
 end
 
-# Shared memory.
+## shared memory
 
-@device_override @inline function KA.SharedMemory(::Type{T}, ::Val{Dims}, ::Val{Id}) where {T, Dims, Id}
-    ptr = AMDGPU.Device.alloc_special(Val(Id), T, Val(AMDGPU.AS.Local), Val(prod(Dims)))
+@device_override @inline function KI.localmemory(::Type{T}, ::Val{Dims}) where {T, Dims}
+    # every call site gets its own memory, see `alloc_special`
+    ptr = Device.alloc_special(Val(:localmemory), T, Val(AMDGPU.AS.Local), Val(prod(Dims)))
     AMDGPU.ROCDeviceArray(Dims, ptr)
 end
 
-@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
-    MArray{KA.__size(Dims), T}(undef)
-end
-
-# Other.
+## synchronization and printing
 
-@device_override @inline function KA.__synchronize()
-    AMDGPU.Device.sync_workgroup()
+@device_override @inline function KI.barrier()
+    Device.sync_workgroup()
 end
 
-@device_override @inline function KA.__print(args...)
-    # TODO
-end
+# not supported, see the `ROCBackend` docstring
+@device_override @inline KI._print(args...) = nothing
+
+## COV_EXCL_STOP
 
 end
diff --git a/src/utils.jl b/src/utils.jl
index 607e00440..7a3c7bbc7 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -118,7 +118,8 @@ function versioninfo(io::IO=stdout)
 
     println(io, "Julia packages: ")
     println(io, "- AMDGPU.jl: $(Base.pkgversion(AMDGPU))")
-    for pkg in [:GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"),
+    for pkg in [:GPUArrays, :GPUCompiler, :KernelInterface,
+                 ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"),
                  :LLVM, :AMDGPU_LLVM_Backend_jll, :LLVMDowngrader_jll]
         name, mod = get_module(pkg)
         isnothing(mod) || println(io, "- $(name): $(Base.pkgversion(mod))")
diff --git a/test/Project.toml b/test/Project.toml
index 8a24df04d..af3919c3a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -11,6 +11,7 @@ GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
diff --git a/test/kernelabstractions_tests.jl b/test/kernelabstractions_tests.jl
index b0de65da9..7fe826b42 100644
--- a/test/kernelabstractions_tests.jl
+++ b/test/kernelabstractions_tests.jl
@@ -2,14 +2,30 @@ using Test
 using AMDGPU
 
 import KernelAbstractions
+import KernelAbstractions as KA
+import KernelInterface as KI
 include(joinpath(pkgdir(KernelAbstractions), "test", "testsuite.jl"))
 
 AMDGPU.allowscalar(false)
 
+KA.@kernel function store_global_linear!(A)
+    I = KA.@index(Global, Linear)
+    @inbounds A[I] = I
+end
+
+KA.@kernel function store_last_index!(A)
+    I = KA.@index(Global, Linear)
+    if I == prod(KA.@ndrange())
+        @inbounds A[1] = I
+        @inbounds A[2] = KA.@index(Global, Cartesian)[2]
+    end
+end
+
 @testset "kernelabstractions" begin
 
 # TODO fix Printing
-skip_tests = ["Printing", "sparse"]
+# sparse is tested by rocSPARSE; the others run kernels on KA's POCL-based CPU back-end
+skip_tests = ["Printing", "sparse", "CPU synchronization", "fallback test: callable types"]
 if Sys.iswindows()
     # TODO
     # We do not support hostcalls on Windows yet.
@@ -55,4 +71,42 @@ if Sys.islinux()
     AMDGPU.synchronize(; stop_hostcalls=true)
 end
 
+@testset "launch configuration" begin
+    backend = ROCBackend()
+    function select(kernel, ndrange, workgroupsize=nothing)
+        ndrange, workgroupsize, iterspace, _ = KA.launch_config(kernel, ndrange, workgroupsize)
+        KA.select_launch(kernel, workgroupsize, iterspace)
+    end
+
+    # kernels are launched on an N-d grid, computing indices in 32 bits
+    kernel = store_global_linear!(backend)
+    @test select(kernel, (64, 32, 16)) === KA.NDLaunch{Int32}()
+    @test select(kernel, (4, 4, 4, 4)) === KA.LinearLaunch{Int32}()
+
+    # which doesn't need divisions to compute the index of a dynamic N-d range
+    A = AMDGPU.zeros(Int, 64, 32, 16)
+    ir = sprint(io -> AMDGPU.@device_code_llvm io=io kernel(A; ndrange=size(A)))
+    @test !occursin(r"\b[su](div|rem) ", ir)
+    @test Array(A) == LinearIndices(A)
+
+    # iteration spaces that don't fit 32 bits use 64-bit indices
+    kernel = store_last_index!(backend)
+    A = AMDGPU.zeros(Int, 2)
+    for (dims, launch) in (((2^16 + 1, 2^15), KA.NDLaunch{Int}()),
+                           ((2^11 + 1, 2^10, 2^10, 1), KA.LinearLaunch{Int}()))
+        @test select(kernel, dims) === launch
+        kernel(A; ndrange=dims)
+        @test Array(A) == [prod(dims), dims[2]]
+    end
+end
+
+@testset "compiler options" begin
+    # a static workgroup size bounds the number of work-items per workgroup
+    A = AMDGPU.zeros(Int, 1024)
+    kernel = store_global_linear!(ROCBackend(), 256)
+    ir = sprint(io -> AMDGPU.@device_code_llvm io=io dump_module=true kernel(A; ndrange=length(A)))
+    @test occursin("\"amdgpu-flat-work-group-size\"=\"1,256\"", ir)
+    @test Array(A) == 1:1024
+end
+
 end
diff --git a/test/kernelinterface_tests.jl b/test/kernelinterface_tests.jl
new file mode 100644
index 000000000..f4fda6600
--- /dev/null
+++ b/test/kernelinterface_tests.jl
@@ -0,0 +1,65 @@
+using Test
+using AMDGPU
+
+import KernelInterface
+import KernelInterface as KI
+include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl"))
+
+AMDGPU.allowscalar(false)
+
+function ki_fill!(A)
+    i = KI.get_global_id().x
+    if i <= length(A)
+        @inbounds A[i] = i
+    end
+    return
+end
+
+@testset "kernelinterface" begin
+
+backend = ROCBackend()
+Testsuite.testsuite(backend, ROCArray)
+
+@testset "copyto!" begin
+    # host to host
+    a = zeros(Float32, 4)
+    @test KI.copyto!(backend, a, ones(Float32, 4)) === a
+    @test a == ones(Float32, 4)
+
+    # contiguous views of host arrays (those of a `ROCArray` are `ROCArray`s)
+    dev = AMDGPU.zeros(Float32, 4)
+    host = Float32[1, 2, 3, 4, 5, 6]
+    @test KI.copyto!(backend, dev, view(host, 2:5)) === dev
+    KI.synchronize(backend)
+    @test Array(dev) == [2, 3, 4, 5]
+    KI.copyto!(backend, view(host, 1:4), AMDGPU.ones(Float32, 4))
+    KI.synchronize(backend)
+    @test host == [1, 1, 1, 1, 5, 6]
+
+    # only contiguous arrays
+    @test_throws ArgumentError KI.copyto!(backend, view(AMDGPU.zeros(Float32, 8), 1:2:8), AMDGPU.ones(Float32, 4))
+end
+
+@testset "launch keywords" begin
+    A = AMDGPU.zeros(Int, 4)
+    kernel = KI.@launch backend launch=false ki_fill!(A)
+
+    # AMDGPU's launch options are passed on
+    kernel(A; ndrange=4, stream=AMDGPU.stream())
+    @test Array(A) == 1:4
+
+    # but not ones that would override the launch geometry
+    @test_throws ArgumentError kernel(A; ndrange=4, groupsize=8)
+    @test_throws ArgumentError kernel(A; ndrange=4, gridsize=2)
+end
+
+@testset "wavefront size" begin
+    # kernels are compiled for the device's wavefront size
+    ws = AMDGPU.HIP.wavefrontsize(AMDGPU.device())
+    A = AMDGPU.zeros(Int, 4)
+    tt = Tuple{typeof(KI.argconvert(backend, A))}
+    @test KI.kernel_function(backend, ki_fill!, tt; wavefrontsize64 = ws == 64) isa KI.Kernel
+    @test_throws ArgumentError KI.kernel_function(backend, ki_fill!, tt; wavefrontsize64 = ws == 32)
+end
+
+end

From cd197d65208736225bb09ffadb382fec62f1d2d7 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:14:44 +0200
Subject: [PATCH 04/13] Add `sync_wavefront`

A barrier for the lanes of a wavefront, fenced at wavefront scope like
`sync_workgroup` is at workgroup scope, so that memory accesses before it are
visible to the other lanes afterwards. `llvm.amdgcn.wave.barrier` alone doesn't
generate any code; it only keeps the compiler from moving code across it.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 docs/src/api/intrinsics.md        |  1 +
 src/AMDGPU.jl                     |  6 ++++--
 src/device/gcn/synchronization.jl | 16 ++++++++++++++++
 test/device/synchronization.jl    | 18 ++++++++++++++++++
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/docs/src/api/intrinsics.md b/docs/src/api/intrinsics.md
index 5a957110b..07b91f39f 100644
--- a/docs/src/api/intrinsics.md
+++ b/docs/src/api/intrinsics.md
@@ -23,6 +23,7 @@ AMDGPU.Device.blockDim
 
 ```@docs
 AMDGPU.sync_workgroup
+AMDGPU.sync_wavefront
 AMDGPU.sync_workgroup_count
 AMDGPU.sync_workgroup_and
 AMDGPU.sync_workgroup_or
diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 2b97bc11e..d8694bded 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -80,9 +80,11 @@ include("memory.jl")
 
 Base.Experimental.@MethodTable(method_table)
 
-# Needs to be before Device since sync uses this. GPUCompiler emits `device` as LLVM's `agent`.
+# Needs to be before Device since sync uses this. GPUCompiler emits `device` as LLVM's
+# `agent`, and `subgroup` as `wavefront`.
 const syncscope_agent = UnsafeAtomics.device
 const syncscope_workgroup = UnsafeAtomics.workgroup
+const syncscope_wavefront = UnsafeAtomics.subgroup
 
 # Referenced by the generated `kernel_state()`, and generators run in the world they are
 # defined in, so this has to precede the device code.
@@ -103,7 +105,7 @@ import .Device: ROCDeviceArray, AS, HostCall, HostCallHolder, hostcall!
 import .Device: @ROCDynamicLocalArray, @ROCStaticLocalArray
 import .Device: workitemIdx, workgroupIdx, workgroupDim, gridItemDim, gridGroupDim
 import .Device: threadIdx, blockIdx, blockDim
-import .Device: sync_workgroup, sync_workgroup_count, sync_workgroup_and, sync_workgroup_or
+import .Device: sync_workgroup, sync_wavefront, sync_workgroup_count, sync_workgroup_and, sync_workgroup_or
 import .Device: @rocprint, @rocprintln, @rocprintf
 
 export ROCDeviceArray, @ROCDynamicLocalArray, @ROCStaticLocalArray
diff --git a/src/device/gcn/synchronization.jl b/src/device/gcn/synchronization.jl
index a9d663807..b900372fa 100644
--- a/src/device/gcn/synchronization.jl
+++ b/src/device/gcn/synchronization.jl
@@ -13,6 +13,22 @@ Waits until all wavefronts in a workgroup have reached this call and that their
     UnsafeAtomics.fence(UnsafeAtomics.seq_cst, AMDGPU.syncscope_workgroup)
 end
 
+"""
+    sync_wavefront()
+
+Waits until all lanes of the wavefront have reached this call, and makes their memory
+accesses before it visible to the other lanes of the wavefront.
+"""
+@device_function @inline function sync_wavefront()
+    # the lanes of a wavefront execute in lockstep, so the barrier doesn't generate any
+    # code (https://github.com/llvm/llvm-project/blob/88b77d5eaa66747538a12c9876eeffdce31ddb71/openmp/device/src/Synchronization.cpp#L136-L140),
+    # but it keeps the compiler from moving code across it. like `sync_workgroup`, it needs
+    # fences to order memory.
+    UnsafeAtomics.fence(UnsafeAtomics.seq_cst, AMDGPU.syncscope_wavefront)
+    ccall("llvm.amdgcn.wave.barrier", llvmcall, Cvoid, ())
+    UnsafeAtomics.fence(UnsafeAtomics.seq_cst, AMDGPU.syncscope_wavefront)
+end
+
 """
     sync_workgroup_count(predicate::Cint)::Cint
 
diff --git a/test/device/synchronization.jl b/test/device/synchronization.jl
index 6d3fe15ba..080897712 100644
--- a/test/device/synchronization.jl
+++ b/test/device/synchronization.jl
@@ -189,4 +189,22 @@ end
     AMDGPU.unsafe_free!(all_workitems_minus_one)
 end
 
+function test_sync_wavefront!(out)
+    i = workitemIdx().x
+    ws = Device.wavefrontsize()
+    shmem = @ROCStaticLocalArray(Int32, 64, false)
+    shmem[i] = i
+    AMDGPU.sync_wavefront()
+    # the value written by the next lane of the wavefront
+    out[i] = shmem[mod1(i + 1, ws)]
+    return
+end
+
+@testset "sync_wavefront" begin
+    ws = Int(AMDGPU.HIP.wavefrontsize(AMDGPU.device()))
+    out = ROCArray{Int32}(undef, ws)
+    @roc groupsize=ws test_sync_wavefront!(out)
+    @test Array(out) == [mod1(i + 1, ws) for i in 1:ws]
+end
+
 end

From 8f27bbba9b12b56a0b33c180a314dbd62ddaea58 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:20:32 +0200
Subject: [PATCH 05/13] Support KernelInterface's sub-groups

Sub-groups are wavefronts: implement the sub-group queries,
`sub_group_barrier` with `sync_wavefront`, and `shfl_down`, and report their
support to KernelInterface. Kernels execute with the device's wavefront size,
which `KI.kernel_function` enforces.

KernelInterface leaves unspecified how work-items are grouped into sub-groups;
AMD GPUs form wavefronts from consecutive linear work-item indices, so the last
wavefront of a workgroup can be partial. The lane is the hardware lane
(`mbcnt`), which unlike `activelane` doesn't change in divergent code. Tests
check both.

`shfl_down` of a `Complex{Int32}` didn't compile, as `Int32` components had no
shuffle method.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 src/ROCKernels.jl             | 48 +++++++++++++++++++++++++++++
 src/device/gcn/wavefront.jl   |  1 +
 test/kernelinterface_tests.jl | 58 ++++++++++++++++++++++++++++++++++-
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index a1f69186a..0221dbb5d 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -42,6 +42,10 @@ KI.synchronize(::ROCBackend) = AMDGPU.synchronize()
 
 KI.supports_float64(::ROCBackend) = true
 KI.supports_atomics(::ROCBackend) = true
+KI.supports_subgroups(::ROCBackend) = true
+# `shfl_down` decomposes other types into 32-bit shuffles
+KI.supports_shuffle(::ROCBackend, ::Type{T}) where {T} =
+    T <: Union{Bool, Base.BitInteger, Base.IEEEFloat, Complex{<:Union{Base.BitInteger, Base.IEEEFloat}}}
 
 function KI.priority!(::ROCBackend, priority::Symbol)
     priority ∉ (:high, :normal, :low) && error(
@@ -153,6 +157,10 @@ function KI.max_num_groups(backend::ROCBackend)::NTuple{3, Int}
         Int(min(Device._max_groups[d], Device._max_grid_size[d] ÷ dims[d]))
     end
 end
+# `KI.kernel_function` compiles for the device's wavefront size
+function KI.sub_group_size(::ROCBackend)::Int
+    Int(HIP.wavefrontsize(AMDGPU.device()))
+end
 function KI.multiprocessor_count(::ROCBackend)::Int
     Int(HIP.attribute(AMDGPU.device(), HIP.hipDeviceAttributeMultiprocessorCount))
 end
@@ -179,6 +187,38 @@ end
     return (; x = Device.gridGroupDim().x % T, y = Device.gridGroupDim().y % T, z = Device.gridGroupDim().z % T)
 end
 
+## sub-groups
+
+# wavefronts are formed from consecutive linear work-item indices
+@inline function linear_workitem_id()
+    return (Device.workitemIdx().x - 0x1) +
+           (Device.workitemIdx().y - 0x1) * Device.workgroupDim().x +
+           (Device.workitemIdx().z - 0x1) * Device.workgroupDim().x * Device.workgroupDim().y
+end
+
+@inline workgroup_items() = Device.workgroupDim().x * Device.workgroupDim().y * Device.workgroupDim().z
+
+# the 0-based lane of the work-item in its wavefront. unlike `Device.activelane()`, this
+# doesn't depend on which lanes are active. `mbcnt.hi` adds nothing in wave32 mode.
+@inline function hardware_lane()
+    lo = ccall("llvm.amdgcn.mbcnt.lo", llvmcall, UInt32, (UInt32, UInt32), typemax(UInt32), 0x0)
+    return ccall("llvm.amdgcn.mbcnt.hi", llvmcall, UInt32, (UInt32, UInt32), typemax(UInt32), lo)
+end
+
+# the last wavefront of a workgroup can be partial
+@device_override @inline function KI.get_sub_group_size(::Type{T}) where {T}
+    ws = Device.wavefrontsize()
+    return min(ws, workgroup_items() - (linear_workitem_id() ÷ ws) * ws) % T
+end
+
+@device_override KI.get_max_sub_group_size(::Type{T}) where {T} = Device.wavefrontsize() % T
+
+@device_override KI.get_num_sub_groups(::Type{T}) where {T} = cld(workgroup_items(), Device.wavefrontsize()) % T
+
+@device_override KI.get_sub_group_id(::Type{T}) where {T} = (linear_workitem_id() ÷ Device.wavefrontsize() + 0x1) % T
+
+@device_override KI.get_sub_group_local_id(::Type{T}) where {T} = (hardware_lane() + 0x1) % T
+
 ## shared memory
 
 @device_override @inline function KI.localmemory(::Type{T}, ::Val{Dims}) where {T, Dims}
@@ -193,6 +233,14 @@ end
     Device.sync_workgroup()
 end
 
+@device_override @inline function KI.sub_group_barrier()
+    Device.sync_wavefront()
+end
+
+@device_override function KI.shfl_down(val::T, offset::Integer) where T
+    @inline Device.shfl_down(val, offset % Cint)
+end
+
 # not supported, see the `ROCBackend` docstring
 @device_override @inline KI._print(args...) = nothing
 
diff --git a/src/device/gcn/wavefront.jl b/src/device/gcn/wavefront.jl
index be3cdd110..7f6522bb6 100644
--- a/src/device/gcn/wavefront.jl
+++ b/src/device/gcn/wavefront.jl
@@ -222,6 +222,7 @@ _shfl(op, x::UInt128) =
     (UInt128(_shfl(op, (x >>> 64) % UInt64)) << 64) |
     UInt128(_shfl(op, ((x & typemax(UInt64)) % UInt64)))
 
+_shfl(op, x::Cint) = op(x)
 _shfl(op, x::Int8) = reinterpret(Int8, _shfl(op, reinterpret(UInt8, x)))
 _shfl(op, x::Int16) = reinterpret(Int16, _shfl(op, reinterpret(UInt16, x)))
 _shfl(op, x::Int64) = reinterpret(Int64, _shfl(op, reinterpret(UInt64, x)))
diff --git a/test/kernelinterface_tests.jl b/test/kernelinterface_tests.jl
index f4fda6600..c4177d4bf 100644
--- a/test/kernelinterface_tests.jl
+++ b/test/kernelinterface_tests.jl
@@ -15,6 +15,33 @@ function ki_fill!(A)
     return
 end
 
+function ki_subgroup_kernel(num, sizes, id, lane)
+    l = KI.get_local_id()
+    s = KI.get_local_size()
+    i = l.x + (l.y - 1) * s.x
+    @inbounds begin
+        num[i] = KI.get_num_sub_groups()
+        sizes[i] = KI.get_sub_group_size()
+        id[i] = KI.get_sub_group_id()
+        lane[i] = KI.get_sub_group_local_id()
+    end
+    return
+end
+
+# only the odd lanes are active when querying the lane id
+function ki_divergent_lane_kernel(lane)
+    i = KI.get_local_id().x
+    if isodd(i)
+        @inbounds lane[i] = KI.get_sub_group_local_id()
+    end
+    return
+end
+
+function ki_wavefront_size_kernel(ws)
+    @inbounds ws[1] = KI.get_max_sub_group_size()
+    return
+end
+
 @testset "kernelinterface" begin
 
 backend = ROCBackend()
@@ -54,12 +81,41 @@ end
 end
 
 @testset "wavefront size" begin
-    # kernels are compiled for the device's wavefront size
+    # kernels are compiled for, and execute with, the device's wavefront size
     ws = AMDGPU.HIP.wavefrontsize(AMDGPU.device())
+    @test KI.sub_group_size(backend) == ws
+    out = AMDGPU.zeros(Int, 1)
+    KI.@launch backend ki_wavefront_size_kernel(out)
+    @test Array(out)[1] == ws
     A = AMDGPU.zeros(Int, 4)
     tt = Tuple{typeof(KI.argconvert(backend, A))}
     @test KI.kernel_function(backend, ki_fill!, tt; wavefrontsize64 = ws == 64) isa KI.Kernel
     @test_throws ArgumentError KI.kernel_function(backend, ki_fill!, tt; wavefrontsize64 = ws == 32)
 end
 
+# KernelInterface leaves the formation of sub-groups unspecified; AMD GPUs form wavefronts
+# from consecutive linear work-item indices
+@testset "partial sub-groups" begin
+    ws = KI.sub_group_size(backend)
+    # a (ws + 1)x2 workgroup is made up of 3 wavefronts, the last one only partially filled
+    workgroupsize = (ws + 1, 2)
+    n = prod(workgroupsize)
+    num = ROCArray{UInt32}(undef, n)
+    sizes = ROCArray{UInt32}(undef, n)
+    id = ROCArray{UInt32}(undef, n)
+    lane = ROCArray{UInt32}(undef, n)
+    KI.@launch backend workgroupsize=workgroupsize ki_subgroup_kernel(num, sizes, id, lane)
+    @test all(==(3), Array(num))
+    @test Array(sizes) == [i < 2ws ? ws : 2 for i in 0:n-1]
+    @test Array(id) == [div(i, ws) + 1 for i in 0:n-1]
+    @test Array(lane) == [rem(i, ws) + 1 for i in 0:n-1]
+end
+
+@testset "lane ids under divergence" begin
+    ws = KI.sub_group_size(backend)
+    lane = AMDGPU.zeros(UInt32, 2ws)
+    KI.@launch backend workgroupsize=2ws ki_divergent_lane_kernel(lane)
+    @test Array(lane) == [isodd(i) ? mod1(i, ws) : 0 for i in 1:2ws]
+end
+
 end

From 2ee737e90127b2ed190812de14b4e1341642a848 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:21:53 +0200
Subject: [PATCH 06/13] Order work across tasks with HIP events

Implement `KI.record_event` and `KI.wait_event` with a `HIPEvent` recorded on,
and waited for by, the task's stream. `KernelAbstractions.@spawn` uses them to
order a new task's work after the work its parent had queued, without
synchronizing the parent; the default is a full synchronization.

KernelInterface's testsuite checks them once `record_event` returns an event.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 src/ROCKernels.jl | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 0221dbb5d..8ddd50783 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -40,6 +40,13 @@ KI.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSR) = ROCBackend()
 
 KI.synchronize(::ROCBackend) = AMDGPU.synchronize()
 
+# an event recorded on, and waited for by, the task's stream
+KI.record_event(::ROCBackend) = HIP.HIPEvent(AMDGPU.stream())
+function KI.wait_event(::ROCBackend, ev::HIP.HIPEvent)
+    HIP.hipStreamWaitEvent(AMDGPU.stream(), ev, 0)
+    return
+end
+
 KI.supports_float64(::ROCBackend) = true
 KI.supports_atomics(::ROCBackend) = true
 KI.supports_subgroups(::ROCBackend) = true

From 925a3305191ab7c7162609504759061cf857e468 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:22:59 +0200
Subject: [PATCH 07/13] Report AMDGPU's version info through KernelInterface

`KI.versioninfo(ROCBackend())` prints `AMDGPU.versioninfo()`.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
---
 src/ROCKernels.jl             | 1 +
 test/kernelinterface_tests.jl | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 8ddd50783..1603c8df7 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -21,6 +21,7 @@ Printing from a kernel (`KernelAbstractions.@print`) is not supported: it does n
 struct ROCBackend <: KI.Backend end
 
 KI.functional(::ROCBackend) = AMDGPU.functional()
+KI.versioninfo(io::IO, ::ROCBackend) = AMDGPU.versioninfo(io)
 KI.ndevices(::ROCBackend) = AMDGPU.HIP.ndevices()
 KI.device(::ROCBackend) = AMDGPU.device_id()
 function KI.device!(kab::ROCBackend, id::Int)
diff --git a/test/kernelinterface_tests.jl b/test/kernelinterface_tests.jl
index c4177d4bf..675b8ae23 100644
--- a/test/kernelinterface_tests.jl
+++ b/test/kernelinterface_tests.jl
@@ -118,4 +118,8 @@ end
     @test Array(lane) == [isodd(i) ? mod1(i, ws) : 0 for i in 1:2ws]
 end
 
+@testset "versioninfo" begin
+    @test occursin("AMDGPU.jl", sprint(KI.versioninfo, backend))
+end
+
 end

From 5dd48c7dfe6fd5907f36a87b741343fbd2b46ef1 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Sep 2026 18:24:25 +0200
Subject: [PATCH 08/13] [TEMP] Get KernelAbstractions 0.10 and KernelInterface
 0.4 from git

KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet.
KernelAbstractions is only a weak dependency of AMDGPU, so it is listed in
[extras] to be allowed in [sources].

Julia 1.10 doesn't support [sources], and without workspaces (Julia 1.11) only
the test project picks up KernelAbstractions from them, so Buildkite runs the
tests in the test project there, developing the packages from git on 1.10. The
GPU-less job moves to Julia 1.12.

Drop this commit once they are registered.
---
 .buildkite/pipeline.yml | 49 ++++++++++++++++++++++++++++++++++++-----
 Project.toml            |  7 ++++++
 test/Project.toml       |  2 ++
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 15a2079a5..04695e31f 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -70,25 +70,63 @@ steps:
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
       JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
+  # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet. Without
+  #      workspaces (Julia 1.11), only the test project, which depends on KernelAbstractions,
+  #      can take it from [sources], and Julia 1.10 doesn't support [sources] at all. So run
+  #      the tests in the test project, developing the packages from git on Julia 1.10, in
+  #      one Pkg operation because AMDGPU's compat excludes the registered versions.
+  #
   # the oldest supported release
-  - <<: *julia
+  - &julia_test_project
+    label: "Julia {{matrix.version}}"
     matrix:
       setup:
         version:
           - "1.10"
-        test_args:
-          - "--verbose"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "rocm"
+      rocmgpu: "*"
     if: |
       build.message !~ /\[skip [^\]]*(tests|julia)/ &&
       (build.message =~ /\[only [^\]]*(tests|julia)/ ||
        build.message !~ /\[only / && !build.pull_request.draft)
+    commands: |
+      git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka
+      julia --project=test -e '
+        using Pkg
+        if VERSION < v"1.11"
+          Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface")])
+        else
+          Pkg.instantiate()
+        end'
+      julia --project=test --check-bounds=yes --code-coverage=@. test/runtests.jl --verbose
+    timeout_in_minutes: 90
+    env:
+      JULIA_NUM_THREADS: 4
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
   # intermediate releases, not tested on pull requests
-  - <<: *julia
+  - <<: *julia_test_project
     matrix:
       setup:
         version:
           - "1.11"
+    if: |
+      build.message !~ /\[skip [^\]]*(tests|julia)/ &&
+      (build.message =~ /\[only [^\]]*(tests|julia)/ ||
+       build.message !~ /\[only / && build.pull_request.id == null)
+
+  - <<: *julia
+    matrix:
+      setup:
+        version:
           - "1.12"
         test_args:
           - "--verbose"
@@ -145,7 +183,8 @@ steps:
   - label: "GPU-less environment"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.10"
+          # XXX: needs Julia 1.12+ to pick up KernelAbstractions 0.10 from [sources]
+          version: "1.12"
       - JuliaCI/julia-test#v1:
           run_tests: false
     command: |
diff --git a/Project.toml b/Project.toml
index 0b1278866..80e1ed2e1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -49,6 +49,10 @@ AMDGPUKernelAbstractionsExt = "KernelAbstractions"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
+[sources]
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}
+
 [compat]
 AMDGPU_LLVM_Backend_jll = "23"
 AbstractFFTs = "1.0"
@@ -83,3 +87,6 @@ StaticArraysCore = "1"
 Statistics = "1"
 UnsafeAtomics = "0.4"
 julia = "1.10"
+
+[extras]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
diff --git a/test/Project.toml b/test/Project.toml
index af3919c3a..16c56d2c0 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,3 +28,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 AMDGPU = {path = ".."}
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}

From 7b974914e410003c6ea759fb3d3eccce0b3981db Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Tue, 6 Oct 2026 11:00:48 +0200
Subject: [PATCH 09/13] Use KernelAbstractions' @private storage

KernelAbstractions now backs @private with a PrivateArray over
GPUCompiler.alloca, defined as an overlay in GPUCompiler's shared method
table, which GPUCompiler stacks under AMDGPU's own method table. Drop our
MArray-based Scratchpad override, which would otherwise take precedence,
along with the StaticArraysCore dependency. The shared table needs
GPUCompiler 2.12, which we already require.
---
 Project.toml                       | 2 --
 ext/AMDGPUKernelAbstractionsExt.jl | 9 ---------
 2 files changed, 11 deletions(-)

diff --git a/Project.toml b/Project.toml
index 80e1ed2e1..000005021 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,7 +31,6 @@ Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
@@ -83,7 +82,6 @@ Scratch = "1"
 SparseArrays = "1"
 SparseMatricesCSR = "0.6.9"
 SpecialFunctions = "2"
-StaticArraysCore = "1"
 Statistics = "1"
 UnsafeAtomics = "0.4"
 julia = "1.10"
diff --git a/ext/AMDGPUKernelAbstractionsExt.jl b/ext/AMDGPUKernelAbstractionsExt.jl
index 177628f28..3f35f8ca7 100644
--- a/ext/AMDGPUKernelAbstractionsExt.jl
+++ b/ext/AMDGPUKernelAbstractionsExt.jl
@@ -1,15 +1,12 @@
 module AMDGPUKernelAbstractionsExt
 
 import AMDGPU
-import AMDGPU.Device: @device_override
 using AMDGPU: GPUArrays, ROCBackend
 
 import Adapt
 import KernelAbstractions as KA
 import LLVM
 
-using StaticArraysCore: MArray
-
 Adapt.adapt_storage(::KA.CPU, a::Union{AMDGPU.ROCArray, GPUArrays.AbstractGPUSparseArray}) =
     Adapt.adapt(Array, a)
 
@@ -25,12 +22,6 @@ function KA.compiler_options(obj::KA.Kernel{ROCBackend})
     end
 end
 
-## scratch memory
-
-@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
-    MArray{Tuple{Dims...}, T}(undef)
-end
-
 ## other
 
 # `@Const` arrays are read through the constant address space

From 9aaba4222f589476f8ebaf16b19bf1b6f13de006 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sat, 10 Oct 2026 18:48:27 +0200
Subject: [PATCH 10/13] [TEMP] Test against KernelAbstractions'
 vc/ki-subgroup-ops branch

Co-authored-by: Tim Besard <tim@juliahub.com>
---
 .buildkite/pipeline.yml | 2 +-
 Project.toml            | 4 ++--
 test/Project.toml       | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 04695e31f..04ffb88fe 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -96,7 +96,7 @@ steps:
       (build.message =~ /\[only [^\]]*(tests|julia)/ ||
        build.message !~ /\[only / && !build.pull_request.draft)
     commands: |
-      git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka
+      git clone --depth 1 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka
       julia --project=test -e '
         using Pkg
         if VERSION < v"1.11"
diff --git a/Project.toml b/Project.toml
index 000005021..cd28d814e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -49,8 +49,8 @@ AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
 [sources]
-KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"}
-KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"}
 
 [compat]
 AMDGPU_LLVM_Backend_jll = "23"
diff --git a/test/Project.toml b/test/Project.toml
index 16c56d2c0..dd64f57b0 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,5 +28,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 AMDGPU = {path = ".."}
-KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"}
-KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"}

From f91d52fce6a24a946b6ce2d3f8dd425ec4eeb00b Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sat, 10 Oct 2026 18:48:40 +0200
Subject: [PATCH 11/13] Implement KernelInterface's sub-group shuffles and
 votes

Implement `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` with `ds_bpermute` on the hardware
lane, for the primitive types it supports by decomposing them into 32-bit words; other types
use KernelInterface's fallbacks. Past the wavefront, `shfl_down` and `shfl_up` return the
work-item's own value, as on CUDA, comparing the offset before narrowing it. The votes use
the wavefront ballot, which only has bits for the work-items of the wavefront.

Wavefronts are formed from consecutive linear work-item indices and execute independently,
so report both `supports_linear_subgroups` and `supports_independent_subgroups`.

Co-authored-by: Tim Besard <tim@juliahub.com>
---
 src/ROCKernels.jl | 58 +++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 53 insertions(+), 5 deletions(-)

diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 1603c8df7..0df8ee723 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -51,9 +51,13 @@ end
 KI.supports_float64(::ROCBackend) = true
 KI.supports_atomics(::ROCBackend) = true
 KI.supports_subgroups(::ROCBackend) = true
-# `shfl_down` decomposes other types into 32-bit shuffles
-KI.supports_shuffle(::ROCBackend, ::Type{T}) where {T} =
-    T <: Union{Bool, Base.BitInteger, Base.IEEEFloat, Complex{<:Union{Base.BitInteger, Base.IEEEFloat}}}
+# wavefronts are formed from consecutive linear work-item indices, and execute independently
+KI.supports_linear_subgroups(::ROCBackend) = true
+KI.supports_independent_subgroups(::ROCBackend) = true
+# the types the shuffles below support, by decomposing them into 32-bit `ds_bpermute`s;
+# KernelInterface shuffles other `isbits` types (e.g. `Complex`) field by field
+const ShuffleTypes = Union{Bool, Base.BitInteger, Base.IEEEFloat}
+KI.supports_shuffle(::ROCBackend, ::Type{T}) where {T <: ShuffleTypes} = true
 
 function KI.priority!(::ROCBackend, priority::Symbol)
     priority ∉ (:high, :normal, :low) && error(
@@ -219,6 +223,8 @@ end
     return min(ws, workgroup_items() - (linear_workitem_id() ÷ ws) * ws) % T
 end
 
+# a constant: `wavefrontsize` is folded to the wavefront size the kernel is compiled for, see
+# `fold_wavefrontsize!`
 @device_override KI.get_max_sub_group_size(::Type{T}) where {T} = Device.wavefrontsize() % T
 
 @device_override KI.get_num_sub_groups(::Type{T}) where {T} = cld(workgroup_items(), Device.wavefrontsize()) % T
@@ -245,10 +251,52 @@ end
     Device.sync_wavefront()
 end
 
-@device_override function KI.shfl_down(val::T, offset::Integer) where T
-    @inline Device.shfl_down(val, offset % Cint)
+## communication
+
+# read `val` from the work-item in the 0-based hardware lane `lane`. `ds_bpermute` only uses
+# the low bits of the address, so lanes out of range read some other lane instead of
+# trapping. unlike `Device.shfl` etc., the lanes are the hardware lanes (see
+# `hardware_lane`), not `activelane`.
+@inline function bpermute_lane(val, lane::Cint)
+    return Device._shfl(x -> Device.bpermute(lane << 0x2, x), val)
+end
+
+@inline lane_id() = hardware_lane() % Cint
+
+@device_override @inline KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} =
+    bpermute_lane(val, (lane % Cint) - Cint(1))
+
+# where the source lane is past the wavefront, `shfl_down` and `shfl_up` read from the
+# work-item itself, like CUDA's shuffles (rather than from the lane `ds_bpermute` wraps around
+# to). The offset is compared before it is narrowed, so that large offsets do too.
+# `wavefrontsize` is folded to a constant, see `fold_wavefrontsize!`.
+@device_override @inline function KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes}
+    lane = lane_id()
+    ws = Device.wavefrontsize() % Cint
+    return bpermute_lane(val, ifelse(offset < ws - lane, lane + (offset % Cint), lane))
+end
+
+@device_override @inline function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes}
+    lane = lane_id()
+    return bpermute_lane(val, ifelse(offset <= lane, lane - (offset % Cint), lane))
 end
 
+# `mask` is below the wavefront size, so the source lane is in the wavefront
+@device_override @inline KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} =
+    bpermute_lane(val, lane_id() ⊻ (mask % Cint))
+
+# the shuffles with a `width` use KernelInterface's fallbacks, a `ds_bpermute` from a lane
+# computed with a few integer operations, like `Device.shfl` etc. (which use `activelane`).
+
+# `ballot` only sets the bits of active lanes, i.e. of the work-items of the sub-group, and
+# its result is uniform. `wavefrontsize`, which selects the 32- or 64-bit ballot, is folded
+# to the wavefront size the kernel is compiled for (see `fold_wavefrontsize!`).
+@device_override @inline KI.sub_group_ballot(pred::Bool) = Device.ballot(pred)
+
+@device_override @inline KI.sub_group_any(pred::Bool) = Device.ballot(pred) != 0
+
+@device_override @inline KI.sub_group_all(pred::Bool) = Device.ballot(!pred) == 0
+
 # not supported, see the `ROCBackend` docstring
 @device_override @inline KI._print(args...) = nothing
 

From c3d2f0df6a4bb92a9eb93fc746ca0660eeb39ab4 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Sun, 4 Oct 2026 15:51:41 +0200
Subject: [PATCH 12/13] Use Julia's min and max for floats

The device overrides of `Base.min` and `Base.max` for floats called OCML's
`__ocml_min`/`__ocml_max`, which, like C's `fmin`/`fmax`, return the other
argument for a NaN, while Julia's `min` and `max` return NaN. Drop them: Julia's
own definitions (`llvm.minimum`/`llvm.maximum` on Julia 1.12+, arithmetic before)
compile for AMDGPU.

Assisted-by: Claude Code (Opus 5.5)
Co-authored-by: Tim Besard <tim@juliahub.com>
---
 src/device/gcn/math.jl |  6 ------
 test/device/math.jl    | 11 +++++++++++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/device/gcn/math.jl b/src/device/gcn/math.jl
index df1cc231f..d45b2e6f5 100644
--- a/src/device/gcn/math.jl
+++ b/src/device/gcn/math.jl
@@ -60,12 +60,6 @@ for jltype in (Float64, Float32, Float16)
     @eval @device_override Base.fma(x::$jltype, y::$jltype, z::$jltype) = ccall(
         $("extern __ocml_fma_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype, $jltype), x, y, z)
 
-    @eval @device_override Base.min(x::$jltype, y::$jltype) = ccall(
-        $("extern __ocml_min_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y)
-
-    @eval @device_override Base.max(x::$jltype, y::$jltype) = ccall(
-        $("extern __ocml_max_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y)
-
     @eval @device_override Base.copysign(x::$jltype, y::$jltype) = ccall(
         $("extern __ocml_copysign_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y)
 
diff --git a/test/device/math.jl b/test/device/math.jl
index 618eada48..a2e433e85 100644
--- a/test/device/math.jl
+++ b/test/device/math.jl
@@ -20,6 +20,17 @@ using Base.FastMath
     end
 end
 
+@testset "min/max" begin
+    # Julia's semantics: NaN propagates, and -0.0 < 0.0
+    a = [NaN, 1, -0.0, 0.0, 2, NaN, Inf, -Inf]
+    b = [1, NaN, 0.0, -0.0, 3, NaN, NaN, 1]
+    for T in (Float16, Float32, Float64)
+        x, y = T.(a), T.(b)
+        @test isequal(Array(max.(ROCArray(x), ROCArray(y))), max.(x, y))
+        @test isequal(Array(min.(ROCArray(x), ROCArray(y))), min.(x, y))
+    end
+end
+
 @testset "Fast min/max" begin
     function ker!(x)
         x[1] = @fastmath max(x[1], zero(eltype(x)))

From b49da45f8862af65835c10436c82592aee3bf537 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Sun, 11 Oct 2026 11:22:43 +0200
Subject: [PATCH 13/13] Report that sub-groups are partitionable

Report `KI.supports_partitionable_subgroups`: the functions of KernelInterface
with a `width` only need the lanes of the caller's segment on AMD GPUs:
`ds_bpermute` only involves the lanes it reads from, `llvm.amdgcn.ballot` only
reports the lanes that execute it, and the sub-group barrier is wavefront-scope
fences around `llvm.amdgcn.wave.barrier`, which HIP's tile synchronization also
uses in divergent code. KernelInterface's fallbacks for the width variants build
on those, so KernelAbstractions' tiles can now pack several narrow tiles into a
wavefront.
---
 src/ROCKernels.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 0df8ee723..7347a3c9c 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -54,8 +54,13 @@ KI.supports_subgroups(::ROCBackend) = true
 # wavefronts are formed from consecutive linear work-item indices, and execute independently
 KI.supports_linear_subgroups(::ROCBackend) = true
 KI.supports_independent_subgroups(::ROCBackend) = true
+# the functions with a `width` only need the lanes of the caller's segment: a `ds_bpermute`
+# only involves the lanes it reads from, `llvm.amdgcn.ballot` only reports the lanes that
+# execute it, and `sub_group_barrier` is wavefront-scope fences around
+# `llvm.amdgcn.wave.barrier`, which (like HIP's tile `sync`) part of a wavefront can execute
+KI.supports_partitionable_subgroups(::ROCBackend) = true
 # the types the shuffles below support, by decomposing them into 32-bit `ds_bpermute`s;
-# KernelInterface shuffles other `isbits` types (e.g. `Complex`) field by field
+# KernelInterface shuffles other `isbits` types (e.g. `Complex`) as `UInt32` words
 const ShuffleTypes = Union{Bool, Base.BitInteger, Base.IEEEFloat}
 KI.supports_shuffle(::ROCBackend, ::Type{T}) where {T <: ShuffleTypes} = true
 
@@ -287,6 +292,8 @@ end
 
 # the shuffles with a `width` use KernelInterface's fallbacks, a `ds_bpermute` from a lane
 # computed with a few integer operations, like `Device.shfl` etc. (which use `activelane`).
+# so do the votes, collectives and barrier with a `width`, which use the functions below
+# and the ones without a `width`.
 
 # `ballot` only sets the bits of active lanes, i.e. of the work-items of the sub-group, and
 # its result is uniform. `wavefrontsize`, which selects the 32- or 64-bit ballot, is folded