From 86ff07e306262f284636b389c8449f00f2b754f9 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Wed, 17 Dec 2025 12:56:50 -0400
Subject: [PATCH 01/14] Thread-local memory without `id` argument

---
 src/device/gcn/memory_static.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/device/gcn/memory_static.jl b/src/device/gcn/memory_static.jl
index b57937911..be0eb4df3 100644
--- a/src/device/gcn/memory_static.jl
+++ b/src/device/gcn/memory_static.jl
@@ -2,7 +2,7 @@
 @generated function alloc_special(
     ::Val{id}, ::Type{T}, ::Val{as}, ::Val{len}, ::Val{zeroinit} = Val{false}(),
 ) where {id,T,as,len,zeroinit}
-    @dispose ctx=Context() begin
+    Context() do ctx
         eltyp = convert(LLVMType, T)
 
         # old versions of GPUArrays invoke _shmem with an integer id; make sure those are unique
@@ -24,8 +24,8 @@
         gv = GlobalVariable(mod, gv_typ, string(id), as)
         if len > 0
             if as == AS.Local
-                linkage!(gv, LLVM.API.LLVMExternalLinkage)
-                # NOTE: Backend doesn't support initializer for local AS
+                linkage!(gv, LLVM.API.LLVMInternalLinkage)
+                initializer!(gv, UndefValue(gv_typ))
             elseif as == AS.Private
                 linkage!(gv, LLVM.API.LLVMInternalLinkage)
                 initializer!(gv, null(gv_typ))
@@ -38,7 +38,7 @@
         alignment!(gv, Base.max(32, Base.datatype_alignment(T)))
 
         # generate IR
-        @dispose builder=IRBuilder() begin
+        IRBuilder() do builder
             entry = BasicBlock(llvm_f, "entry")
             position!(builder, entry)
 

From 972aa4a8298ad6430dfe30f69cb2a2efd43306c1 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Sat, 22 Aug 2026 15:43:52 -0300
Subject: [PATCH 02/14] sync_wavefront

---
 src/AMDGPU.jl                     |  2 +-
 src/device/gcn/synchronization.jl | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 636b3bf27..94a6147f4 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -102,7 +102,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 6d8b682dc..5e40d9cac 100644
--- a/src/device/gcn/synchronization.jl
+++ b/src/device/gcn/synchronization.jl
@@ -29,6 +29,17 @@ 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 wavefronts in a workgroup have reached this call and that their memory accesses are visible to other threads in the workgroup.
+"""
+@inline function sync_wavefront()
+    # This is a no-op https://github.com/llvm/llvm-project/blob/88b77d5eaa66747538a12c9876eeffdce31ddb71/openmp/device/src/Synchronization.cpp#L136-L140
+    ccall("llvm.amdgcn.wave.barrier", llvmcall, Cvoid, ())
+end
+
+
 """
     sync_workgroup_count(predicate::Cint)::Cint
 

From 122af31f7c99e073a5ef55cc921fe0574e6ecd9a Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Sat, 22 Aug 2026 15:57:39 -0300
Subject: [PATCH 03/14] KernelInterface

---
 Project.toml                  |   2 +
 src/AMDGPU.jl                 |   7 +-
 src/ROCKernels.jl             | 205 ++++++++++++++--------------------
 src/ROCKernelsOld.jl          | 201 +++++++++++++++++++++++++++++++++
 test/Project.toml             |   1 +
 test/kernelinterface_tests.jl |  15 +++
 6 files changed, 311 insertions(+), 120 deletions(-)
 create mode 100644 src/ROCKernelsOld.jl
 create mode 100644 test/kernelinterface_tests.jl

diff --git a/Project.toml b/Project.toml
index 8bc9180ef..0d22a3edd 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,6 +19,7 @@ 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"
@@ -63,6 +64,7 @@ GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
 KernelAbstractions = "0.9.2"
+KernelInterface = "0.1.0"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
 Libdl = "1"
diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 94a6147f4..79fda4125 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -147,10 +147,15 @@ function Atomix.modify!(ref::ROCIndexableRef, op::OP, x, ord) where OP <: Union{
     GC.@preserve root UnsafeAtomics.modify!(ptr, op, x, ord, syncscope_agent)
 end
 
-include("ROCKernels.jl")
+# KernelAbstractions
+include("ROCKernelsOld.jl")
 import .ROCKernels: ROCBackend
 export ROCBackend
 
+# KernelInterface
+include("ROCKernels.jl")
+import .ROCInterface
+
 include("precompile.jl")
 
 function __init__()
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 2315222fe..f9d1ad037 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -1,13 +1,14 @@
-module ROCKernels
+module ROCInterface
 
 export ROCBackend
 
 import AMDGPU
+import AMDGPU: rocconvert, hipfunction
 import AMDGPU.Device: @device_override
-using AMDGPU: GPUArrays, rocSPARSE
+using AMDGPU: GPUArrays, rocSPARSE, HIP, Device
 
 import Adapt
-import KernelAbstractions as KA
+import KernelInterface as KI
 import LLVM
 
 using StaticArraysCore: MArray
@@ -19,183 +20,149 @@ KernelAbstractions 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)`.
 """
-struct ROCBackend <: KA.GPU end
+struct ROCBackend <: KI.GPU 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.versioninfo(io::IO, ::ROCBackend) = AMDGPU.versioninfo(io)
+
+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
 
 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.unsafe_free!(x::AMDGPU.ROCArray) = AMDGPU.unsafe_free!(x)
+KI.allocate(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ROCArray{T}(undef, dims)
+KI.zeros(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.zeros(T, dims)
+KI.ones(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ones(T, dims)
 
-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)
 end
 
-function KA.copyto!(::ROCBackend, A, B)
+function KI.copyto!(::ROCBackend, A, B)
     GC.@preserve A B begin
         copyto!(A, 1, B, 1, length(A))
     end
     return
 end
 
-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
-
-    iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize && workgroupsize === nothing
-        workgroupsize = ntuple(
-            i -> i == 1 ? min(prod(ndrange), AMDGPU.Device._max_group_size) : 1,
-            length(ndrange))
-        KA.partition(kernel, ndrange, workgroupsize)
-    else
-        KA.partition(kernel, ndrange, workgroupsize)
-    end
+KI.argconvert(::ROCBackend, arg) = rocconvert(arg)
 
-    return ndrange, workgroupsize, iterspace, dynamic
+function KI.kernel_function(::ROCBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
+    kern = hipfunction(f, tt; name, kwargs...)
+    KI.Kernel{ROCBackend, typeof(kern)}(ROCBackend(), kern)
 end
 
-function threads_to_workgroupsize(threads, ndrange)
-    total = 1
-    return map(ndrange) do n
-        x = min(div(threads, total), n)
-        total *= x
-        return x
-    end
-end
+function (obj::KI.Kernel{ROCBackend})(args...; numworkgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int))
+    KI.check_launch_args(numworkgroups, workgroupsize, ndrange)
+    prod(ndrange) == 0 && return nothing
 
-function (obj::KA.Kernel{ROCBackend})(args...; 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
-    end
-    kernel = AMDGPU.@roc launch=false maxthreads=maxthreads obj.f(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)
-    end
+    numworkgroups, workgroupsize = KI.auto_launch_sizes(obj, numworkgroups, workgroupsize, ndrange, max_work_group_size)
+    obj.kern(args...; groupsize = workgroupsize, gridsize = numworkgroups)
+    return nothing
+end
 
-    nblocks = length(KA.blocks(iterspace))
-    nthreads = length(KA.workitems(iterspace))
-    nblocks == 0 && return
+function KI.kernel_max_work_group_size(kikern::KI.Kernel{<:ROCBackend}; max_work_items::Int=Int(typemax(Int32)))::Int
+    (; groupsize) = AMDGPU.launch_configuration(kikern.kern; max_block_size = max_work_items)
 
-    kernel(ctx, args...; groupsize=nthreads, gridsize=nblocks)
-    return
+    return Int(min(max_work_items, groupsize))
 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(::ROCBackend)::Int
+    Int(HIP.attribute(AMDGPU.HIP.device(), AMDGPU.HIP.hipDeviceAttributeMaxThreadsPerBlock))
 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.sub_group_size(::ROCBackend)::Int
+    HIP.wavefrontsize(HIP.device())
+end
+function KI.multiprocessor_count(::ROCBackend)::Int
+    Int(HIP.attribute(AMDGPU.HIP.device(), AMDGPU.HIP.hipDeviceAttributeMultiprocessorCount))
 end
 
-# Indexing.
+KI.shfl_down_types(::ROCBackend) = DataType[Bool,
+                                             UInt8, UInt16, UInt32, UInt64, UInt128,
+                                             Int8, Int16, Int32, Int64, Int128,
+                                             Float16, Float32, Float64,
+                                             ComplexF16, ComplexF32, ComplexF64]
 
-@device_override @inline function KA.__index_Local_Linear(ctx)
-    return AMDGPU.Device.threadIdx().x
+# Indexing.
+## COV_EXCL_START
+@device_override @inline function KI.get_local_id()
+    return (; x = Int(AMDGPU.Device.workitemIdx().x), y = Int(AMDGPU.Device.workitemIdx().y), z = Int(AMDGPU.Device.workitemIdx().z))
 end
 
-@device_override @inline function KA.__index_Group_Linear(ctx)
-    return AMDGPU.Device.blockIdx().x
+@device_override @inline function KI.get_group_id()
+    return (; x = Int(AMDGPU.Device.workgroupIdx().x), y = Int(AMDGPU.Device.workgroupIdx().y), z = Int(AMDGPU.Device.workgroupIdx().z))
 end
 
-@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]
+@device_override @inline function KI.get_global_id()
+    return (; x = Int((AMDGPU.Device.workgroupIdx().x-1)*AMDGPU.Device.blockDim().x + AMDGPU.Device.workitemIdx().x), y = Int((AMDGPU.Device.workgroupIdx().y-1)*AMDGPU.Device.blockDim().y + AMDGPU.Device.workitemIdx().y), z = Int((AMDGPU.Device.workgroupIdx().z-1)*AMDGPU.Device.blockDim().z + AMDGPU.Device.workitemIdx().z))
 end
 
-@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_size()
+    return (; x = Int(AMDGPU.Device.workgroupDim().x), y = Int(AMDGPU.Device.workgroupDim().y), z = Int(AMDGPU.Device.workgroupDim().z))
 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_num_groups()
+    return (; x = Int(AMDGPU.Device.gridGroupDim().x), y = Int(AMDGPU.Device.gridGroupDim().y), z = Int(AMDGPU.Device.gridGroupDim().z))
 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_global_size()
+    return (; x = Int(AMDGPU.Device.gridItemDim().x), y = Int(AMDGPU.Device.gridItemDim().y), z = Int(AMDGPU.Device.gridItemDim().z))
 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
-end
+@device_override KI.get_sub_group_size() = UInt32(Device.wavefrontsize())
+
+@device_override KI.get_max_sub_group_size() = UInt32(Device.wavefrontsize())
+
+@device_override KI.get_num_sub_groups() = UInt32(prod(Device.blockDim()) ÷ Device.wavefrontsize())
+
+@device_override KI.get_sub_group_id() = UInt32(((Device.threadIdx().x - 1) + Device.blockDim().x * (Device.threadIdx().y - 1) + Device.blockDim().x * Device.blockDim().y * (Device.threadIdx().z - 1)) ÷ Device.wavefrontsize()) + 0x1
+
+@device_override KI.get_sub_group_local_id() = UInt32(Device.activelane() + 0x1)
 
 # 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}
+    ptr = AMDGPU.Device.alloc_special(Val(:shmem), 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.
 
-@device_override @inline function KA.__synchronize()
+@device_override @inline function KI.barrier()
     AMDGPU.Device.sync_workgroup()
 end
 
-@device_override @inline function KA.__print(args...)
+@device_override @inline function KI.sub_group_barrier()
+    AMDGPU.Device.sync_wavefront()
+end
+
+@device_override function KI.shfl_down(val::T, offset::Integer) where T
+    @inline AMDGPU.Device.shfl_down(val, Cint(offset))
+end
+
+@device_override @inline function KI._print(args...)
     # TODO
 end
+## COV_EXCL_STOP
 
 end
diff --git a/src/ROCKernelsOld.jl b/src/ROCKernelsOld.jl
new file mode 100644
index 000000000..2315222fe
--- /dev/null
+++ b/src/ROCKernelsOld.jl
@@ -0,0 +1,201 @@
+module ROCKernels
+
+export ROCBackend
+
+import AMDGPU
+import AMDGPU.Device: @device_override
+using AMDGPU: GPUArrays, rocSPARSE
+
+import Adapt
+import KernelAbstractions as KA
+import LLVM
+
+using StaticArraysCore: MArray
+
+"""
+    ROCBackend <: KernelAbstractions.GPU
+
+KernelAbstractions 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)`.
+"""
+struct ROCBackend <: KA.GPU 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."))
+    AMDGPU.device_id!(id)
+    return
+end
+
+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()
+
+KA.argconvert(::KA.Kernel{ROCBackend}, arg) = AMDGPU.rocconvert(arg)
+KA.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)
+
+function KA.priority!(::ROCBackend, priority::Symbol)
+    priority ∉ (:high, :normal, :low) && error(
+        "Priority `$priority` must be one of `:high`, `:normal`, `:low`.")
+    AMDGPU.priority!(priority)
+end
+
+function KA.copyto!(::ROCBackend, A, B)
+    GC.@preserve A B begin
+        copyto!(A, 1, B, 1, length(A))
+    end
+    return
+end
+
+function KA.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
+
+    iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize && workgroupsize === nothing
+        workgroupsize = ntuple(
+            i -> i == 1 ? min(prod(ndrange), AMDGPU.Device._max_group_size) : 1,
+            length(ndrange))
+        KA.partition(kernel, ndrange, workgroupsize)
+    else
+        KA.partition(kernel, ndrange, workgroupsize)
+    end
+
+    return ndrange, workgroupsize, iterspace, dynamic
+end
+
+function threads_to_workgroupsize(threads, ndrange)
+    total = 1
+    return map(ndrange) do n
+        x = min(div(threads, total), n)
+        total *= x
+        return x
+    end
+end
+
+function (obj::KA.Kernel{ROCBackend})(args...; 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
+    end
+    kernel = AMDGPU.@roc launch=false maxthreads=maxthreads obj.f(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)
+    end
+
+    nblocks = length(KA.blocks(iterspace))
+    nthreads = length(KA.workitems(iterspace))
+    nblocks == 0 && return
+
+    kernel(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)
+end
+function KA.mkcontext(kernel::KA.Kernel{ROCBackend}, I, _ndrange, iterspace, ::Dynamic) where Dynamic
+    metadata = KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
+end
+
+# Indexing.
+
+@device_override @inline function KA.__index_Local_Linear(ctx)
+    return AMDGPU.Device.threadIdx().x
+end
+
+@device_override @inline function KA.__index_Group_Linear(ctx)
+    return AMDGPU.Device.blockIdx().x
+end
+
+@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
+
+@device_override @inline function KA.__index_Local_Cartesian(ctx)
+    @inbounds KA.workitems(KA.__iterspace(ctx))[AMDGPU.Device.threadIdx().x]
+end
+
+@device_override @inline function KA.__index_Group_Cartesian(ctx)
+    @inbounds KA.blocks(KA.__iterspace(ctx))[AMDGPU.Device.blockIdx().x]
+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)
+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
+end
+
+# 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)))
+    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.
+
+@device_override @inline function KA.__synchronize()
+    AMDGPU.Device.sync_workgroup()
+end
+
+@device_override @inline function KA.__print(args...)
+    # TODO
+end
+
+end
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/kernelinterface_tests.jl b/test/kernelinterface_tests.jl
new file mode 100644
index 000000000..2eba9d15f
--- /dev/null
+++ b/test/kernelinterface_tests.jl
@@ -0,0 +1,15 @@
+using Test
+using AMDGPU
+using AMDGPU: ROCInterface
+
+import KernelInterface
+include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl"))
+
+AMDGPU.allowscalar(false)
+
+@testset "kernelinterface" begin
+
+Testsuite.testsuite(
+    ROCInterface.ROCBackend, "ROCM", AMDGPU, ROCArray, AMDGPU.ROCDeviceArray)
+
+end

From b3d70de49196de4391ba5decebb534f33ae646e3 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:50:15 -0300
Subject: [PATCH 04/14] Support KI 0.2

---
 Project.toml      |  2 +-
 src/ROCKernels.jl | 24 ++++++++++++------------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/Project.toml b/Project.toml
index 0d22a3edd..da7b821ed 100644
--- a/Project.toml
+++ b/Project.toml
@@ -64,7 +64,7 @@ GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
 KernelAbstractions = "0.9.2"
-KernelInterface = "0.1.0"
+KernelInterface = "0.2"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
 Libdl = "1"
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index f9d1ad037..fa4d70e7b 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -105,28 +105,28 @@ KI.shfl_down_types(::ROCBackend) = DataType[Bool,
 
 # Indexing.
 ## COV_EXCL_START
-@device_override @inline function KI.get_local_id()
-    return (; x = Int(AMDGPU.Device.workitemIdx().x), y = Int(AMDGPU.Device.workitemIdx().y), z = Int(AMDGPU.Device.workitemIdx().z))
+@device_override @inline function KI.get_local_id(::Type{T}) where {T}
+    return (; x = T(AMDGPU.Device.workitemIdx().x), y = T(AMDGPU.Device.workitemIdx().y), z = T(AMDGPU.Device.workitemIdx().z))
 end
 
-@device_override @inline function KI.get_group_id()
-    return (; x = Int(AMDGPU.Device.workgroupIdx().x), y = Int(AMDGPU.Device.workgroupIdx().y), z = Int(AMDGPU.Device.workgroupIdx().z))
+@device_override @inline function KI.get_group_id(::Type{T}) where {T}
+    return (; x = T(AMDGPU.Device.workgroupIdx().x), y = T(AMDGPU.Device.workgroupIdx().y), z = T(AMDGPU.Device.workgroupIdx().z))
 end
 
-@device_override @inline function KI.get_global_id()
-    return (; x = Int((AMDGPU.Device.workgroupIdx().x-1)*AMDGPU.Device.blockDim().x + AMDGPU.Device.workitemIdx().x), y = Int((AMDGPU.Device.workgroupIdx().y-1)*AMDGPU.Device.blockDim().y + AMDGPU.Device.workitemIdx().y), z = Int((AMDGPU.Device.workgroupIdx().z-1)*AMDGPU.Device.blockDim().z + AMDGPU.Device.workitemIdx().z))
+@device_override @inline function KI.get_global_id(::Type{T}) where {T}
+    return (; x = T((AMDGPU.Device.workgroupIdx().x-1)*AMDGPU.Device.blockDim().x + AMDGPU.Device.workitemIdx().x), y = T((AMDGPU.Device.workgroupIdx().y-1)*AMDGPU.Device.blockDim().y + AMDGPU.Device.workitemIdx().y), z = T((AMDGPU.Device.workgroupIdx().z-1)*AMDGPU.Device.blockDim().z + AMDGPU.Device.workitemIdx().z))
 end
 
-@device_override @inline function KI.get_local_size()
-    return (; x = Int(AMDGPU.Device.workgroupDim().x), y = Int(AMDGPU.Device.workgroupDim().y), z = Int(AMDGPU.Device.workgroupDim().z))
+@device_override @inline function KI.get_local_size(::Type{T}) where {T}
+    return (; x = T(AMDGPU.Device.workgroupDim().x), y = T(AMDGPU.Device.workgroupDim().y), z = T(AMDGPU.Device.workgroupDim().z))
 end
 
-@device_override @inline function KI.get_num_groups()
-    return (; x = Int(AMDGPU.Device.gridGroupDim().x), y = Int(AMDGPU.Device.gridGroupDim().y), z = Int(AMDGPU.Device.gridGroupDim().z))
+@device_override @inline function KI.get_num_groups(::Type{T}) where {T}
+    return (; x = T(AMDGPU.Device.gridGroupDim().x), y = T(AMDGPU.Device.gridGroupDim().y), z = T(AMDGPU.Device.gridGroupDim().z))
 end
 
-@device_override @inline function KI.get_global_size()
-    return (; x = Int(AMDGPU.Device.gridItemDim().x), y = Int(AMDGPU.Device.gridItemDim().y), z = Int(AMDGPU.Device.gridItemDim().z))
+@device_override @inline function KI.get_global_size(::Type{T}) where {T}
+    return (; x = T(AMDGPU.Device.gridItemDim().x), y = T(AMDGPU.Device.gridItemDim().y), z = T(AMDGPU.Device.gridItemDim().z))
 end
 
 @device_override KI.get_sub_group_size() = UInt32(Device.wavefrontsize())

From e33ef3b0443752597560f99ff175c4dd4f51def7 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Mon, 21 Sep 2026 16:23:12 -0300
Subject: [PATCH 05/14] KI 0.2.1

---
 Project.toml      | 2 +-
 src/ROCKernels.jl | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index da7b821ed..d7d144f4b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -64,7 +64,7 @@ GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
 KernelAbstractions = "0.9.2"
-KernelInterface = "0.2"
+KernelInterface = "0.2.1"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
 Libdl = "1"
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index fa4d70e7b..137849c55 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -43,6 +43,15 @@ KI.get_backend(::AMDGPU.rocSPARSE.ROCSparseMatrixCSR) = ROCBackend()
 
 KI.synchronize(::ROCBackend) = AMDGPU.synchronize()
 
+function KI.record_event(::ROCBackend)
+    return HIP.HIPEvent(AMDGPU.stream())
+end
+
+function KI.wait_event(::ROCBackend, ev::HIP.HIPEvent)
+    HIP.hipStreamWaitEvent(AMDGPU.stream(), ev, 0)
+    return
+end
+
 KI.unsafe_free!(x::AMDGPU.ROCArray) = AMDGPU.unsafe_free!(x)
 KI.allocate(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ROCArray{T}(undef, dims)
 KI.zeros(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.zeros(T, dims)

From 9aef27bb506a799db779d43579d876d4359386e8 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Mon, 28 Sep 2026 00:24:01 -0300
Subject: [PATCH 06/14] Add KernelInterface to versioninfo

---
 src/AMDGPU.jl | 1 +
 src/utils.jl  | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 79fda4125..9401b175e 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -60,6 +60,7 @@ include("discovery/discovery.jl")
 using .ROCmDiscovery
 using .ROCmDiscovery: AMDGPU_LLVM_Backend_jll, LLVMDowngrader_jll
 
+using KernelInterface
 include("utils.jl")
 
 include("hsa/HSA.jl")
diff --git a/src/utils.jl b/src/utils.jl
index 11a92547d..799380df8 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -97,7 +97,7 @@ 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"),
-                 :LLVM, :AMDGPU_LLVM_Backend_jll, :LLVMDowngrader_jll]
+                 :KernelInterface, :LLVM, :AMDGPU_LLVM_Backend_jll, :LLVMDowngrader_jll]
         name, mod = get_module(pkg)
         isnothing(mod) || println(io, "- $(name): $(Base.pkgversion(mod))")
     end

From 651b5d56ecfa3cc58d5467cb2fc842de00f3b6f6 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Wed, 17 Dec 2025 15:16:14 -0400
Subject: [PATCH 07/14] Remove irrelevant KA tests

These always run on the CPU backend which is currently broken in 1.12, creating false negatives for the GPU tests
---
 test/kernelabstractions_tests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/kernelabstractions_tests.jl b/test/kernelabstractions_tests.jl
index 2c8220429..6601087ed 100644
--- a/test/kernelabstractions_tests.jl
+++ b/test/kernelabstractions_tests.jl
@@ -9,7 +9,7 @@ AMDGPU.allowscalar(false)
 @testset "kernelabstractions" begin
 
 # TODO fix Printing
-skip_tests = ["Printing", "sparse"]
+skip_tests = ["Printing", "sparse", "CPU synchronization", "fallback test: callable types",]
 if Sys.iswindows()
     # TODO
     # We do not support hostcalls on Windows yet.

From 67c6e3bc2aba79960d089840f879d775d7e8f3da Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Sat, 22 Aug 2026 16:10:36 -0300
Subject: [PATCH 08/14] Temp CI

---
 Project.toml     | 6 +++++-
 test/runtests.jl | 6 ++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index d7d144f4b..f0a1d4041 100644
--- a/Project.toml
+++ b/Project.toml
@@ -43,6 +43,10 @@ EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
+[sources]
+KernelAbstractions = {rev = "main", url = "https://github.com/JuliaGPU/KernelAbstractions.jl"}
+AcceleratedKernels = {rev = "main", url = "https://github.com/JuliaGPU/AcceleratedKernels.jl"}
+
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
@@ -63,7 +67,7 @@ ExprTools = "0.1"
 GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.9, 0.10"
 KernelInterface = "0.2.1"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
diff --git a/test/runtests.jl b/test/runtests.jl
index e87358e57..12e37b592 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,3 +1,9 @@
+@static if VERSION < v"1.11" && get(ENV, "BUILDKITE_PIPELINE_NAME", "AMDGPU.jl") == "AMDGPU.jl"
+    using Pkg
+    Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
+    Pkg.add(url="https://github.com/JuliaGPU/KernelAbstractions.jl", rev="main")
+end
+
 using AMDGPU
 using AMDGPU: Device, Runtime, @allowscalar
 import AMDGPU.Device: HostCallHolder, hostcall!

From a6902f94a71aeeabf4d14877a08761974965d70f Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:08:17 -0300
Subject: [PATCH 09/14] Turn KA into an extension

---
 Project.toml                                  |   5 +-
 .../AMDGPUKernelAbstractionsExt.jl            | 100 +-----------------
 src/AMDGPU.jl                                 |   7 +-
 src/ROCKernels.jl                             |   2 +-
 test/kernelinterface_tests.jl                 |   4 +-
 5 files changed, 9 insertions(+), 109 deletions(-)
 rename src/ROCKernelsOld.jl => ext/AMDGPUKernelAbstractionsExt.jl (50%)

diff --git a/Project.toml b/Project.toml
index f0a1d4041..0c75c824e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,7 +18,6 @@ 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"
@@ -40,6 +39,7 @@ 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"
 
@@ -50,6 +50,7 @@ AcceleratedKernels = {rev = "main", url = "https://github.com/JuliaGPU/Accelerat
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
+AMDGPUKernelAbstractionsExt = "KernelAbstractions"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
@@ -67,7 +68,7 @@ ExprTools = "0.1"
 GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
-KernelAbstractions = "0.9, 0.10"
+KernelAbstractions = "0.10"
 KernelInterface = "0.2.1"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
diff --git a/src/ROCKernelsOld.jl b/ext/AMDGPUKernelAbstractionsExt.jl
similarity index 50%
rename from src/ROCKernelsOld.jl
rename to ext/AMDGPUKernelAbstractionsExt.jl
index 2315222fe..b9e163e30 100644
--- a/src/ROCKernelsOld.jl
+++ b/ext/AMDGPUKernelAbstractionsExt.jl
@@ -1,10 +1,8 @@
-module ROCKernels
-
-export ROCBackend
+module AMDGPUKernelAbstractionsExt
 
 import AMDGPU
 import AMDGPU.Device: @device_override
-using AMDGPU: GPUArrays, rocSPARSE
+using AMDGPU: GPUArrays, rocSPARSE, ROCBackend
 
 import Adapt
 import KernelAbstractions as KA
@@ -12,26 +10,6 @@ import LLVM
 
 using StaticArraysCore: MArray
 
-"""
-    ROCBackend <: KernelAbstractions.GPU
-
-KernelAbstractions 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)`.
-"""
-struct ROCBackend <: KA.GPU 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."))
-    AMDGPU.device_id!(id)
-    return
-end
-
-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
@@ -39,36 +17,7 @@ function Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) whe
     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()
-
 KA.argconvert(::KA.Kernel{ROCBackend}, arg) = AMDGPU.rocconvert(arg)
-KA.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)
-
-function KA.priority!(::ROCBackend, priority::Symbol)
-    priority ∉ (:high, :normal, :low) && error(
-        "Priority `$priority` must be one of `:high`, `:normal`, `:low`.")
-    AMDGPU.priority!(priority)
-end
-
-function KA.copyto!(::ROCBackend, A, B)
-    GC.@preserve A B begin
-        copyto!(A, 1, B, 1, length(A))
-    end
-    return
-end
-
-function KA.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
@@ -140,34 +89,6 @@ function KA.mkcontext(kernel::KA.Kernel{ROCBackend}, I, _ndrange, iterspace, ::D
     metadata = KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
 end
 
-# Indexing.
-
-@device_override @inline function KA.__index_Local_Linear(ctx)
-    return AMDGPU.Device.threadIdx().x
-end
-
-@device_override @inline function KA.__index_Group_Linear(ctx)
-    return AMDGPU.Device.blockIdx().x
-end
-
-@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
-
-@device_override @inline function KA.__index_Local_Cartesian(ctx)
-    @inbounds KA.workitems(KA.__iterspace(ctx))[AMDGPU.Device.threadIdx().x]
-end
-
-@device_override @inline function KA.__index_Group_Cartesian(ctx)
-    @inbounds KA.blocks(KA.__iterspace(ctx))[AMDGPU.Device.blockIdx().x]
-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)
-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)
@@ -177,25 +98,8 @@ end
     end
 end
 
-# 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)))
-    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.
-
-@device_override @inline function KA.__synchronize()
-    AMDGPU.Device.sync_workgroup()
-end
-
-@device_override @inline function KA.__print(args...)
-    # TODO
-end
-
 end
diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl
index 9401b175e..9e42cb45d 100644
--- a/src/AMDGPU.jl
+++ b/src/AMDGPU.jl
@@ -148,14 +148,11 @@ function Atomix.modify!(ref::ROCIndexableRef, op::OP, x, ord) where OP <: Union{
     GC.@preserve root UnsafeAtomics.modify!(ptr, op, x, ord, syncscope_agent)
 end
 
-# KernelAbstractions
-include("ROCKernelsOld.jl")
+# KernelInterface
+include("ROCKernels.jl")
 import .ROCKernels: ROCBackend
 export ROCBackend
 
-# KernelInterface
-include("ROCKernels.jl")
-import .ROCInterface
 
 include("precompile.jl")
 
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 137849c55..61fd3f1b4 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -1,4 +1,4 @@
-module ROCInterface
+module ROCKernels
 
 export ROCBackend
 
diff --git a/test/kernelinterface_tests.jl b/test/kernelinterface_tests.jl
index 2eba9d15f..018754687 100644
--- a/test/kernelinterface_tests.jl
+++ b/test/kernelinterface_tests.jl
@@ -1,6 +1,5 @@
 using Test
 using AMDGPU
-using AMDGPU: ROCInterface
 
 import KernelInterface
 include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl"))
@@ -10,6 +9,5 @@ AMDGPU.allowscalar(false)
 @testset "kernelinterface" begin
 
 Testsuite.testsuite(
-    ROCInterface.ROCBackend, "ROCM", AMDGPU, ROCArray, AMDGPU.ROCDeviceArray)
-
+    ROCBackend, "ROCM", AMDGPU, ROCArray, AMDGPU.ROCDeviceArray)
 end

From 9492f7e278b51db7b2db7ac54837f77ef556f4e8 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Thu, 3 Sep 2026 12:08:29 -0300
Subject: [PATCH 10/14] [temp] source KA

---
 Project.toml      | 4 +++-
 test/Project.toml | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 0c75c824e..8bcff419e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -44,7 +44,6 @@ SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [sources]
-KernelAbstractions = {rev = "main", url = "https://github.com/JuliaGPU/KernelAbstractions.jl"}
 AcceleratedKernels = {rev = "main", url = "https://github.com/JuliaGPU/AcceleratedKernels.jl"}
 
 [extensions]
@@ -89,3 +88,6 @@ StaticArraysCore = "1"
 Statistics = "1"
 UnsafeAtomics = "0.3"
 julia = "1.10"
+
+[extras]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
diff --git a/test/Project.toml b/test/Project.toml
index af3919c3a..38ada0c60 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,3 +28,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 AMDGPU = {path = ".."}
+KernelAbstractions = {rev = "main", url = "https://github.com/JuliaGPU/KernelAbstractions.jl"}

From e2a30b565f93884d6734a9a7462b0bfb85ba0b7b Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Tue, 8 Sep 2026 00:00:44 -0300
Subject: [PATCH 11/14] [CI] Try new approach

---
 .buildkite/pipeline.yml | 3 ++-
 test/runtests.jl        | 6 ------
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 6a5489c49..07fb3bb61 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -61,7 +61,8 @@ steps:
     if: |
       build.message !~ /\[skip [^\]]*(tests|julia)/ &&
       (build.message =~ /\[only [^\]]*(tests|julia)/ || build.message !~ /\[only /)
-    command: "julia --project -e 'using Pkg; Pkg.update()'"
+    command:
+      julia --project -e 'using Pkg; Pkg.update(); Pkg.develop(["AcceleratedKernels", "KernelAbstractions"])'
     timeout_in_minutes: 90
     env:
       JULIA_NUM_THREADS: 4
diff --git a/test/runtests.jl b/test/runtests.jl
index 12e37b592..e87358e57 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,9 +1,3 @@
-@static if VERSION < v"1.11" && get(ENV, "BUILDKITE_PIPELINE_NAME", "AMDGPU.jl") == "AMDGPU.jl"
-    using Pkg
-    Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
-    Pkg.add(url="https://github.com/JuliaGPU/KernelAbstractions.jl", rev="main")
-end
-
 using AMDGPU
 using AMDGPU: Device, Runtime, @allowscalar
 import AMDGPU.Device: HostCallHolder, hostcall!

From 1db41e45d1c91c91285a37e43ece94d4542917a4 Mon Sep 17 00:00:00 2001
From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Date: Sat, 12 Sep 2026 13:11:44 -0300
Subject: [PATCH 12/14] No source

---
 test/Project.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/Project.toml b/test/Project.toml
index 38ada0c60..af3919c3a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,4 +28,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 AMDGPU = {path = ".."}
-KernelAbstractions = {rev = "main", url = "https://github.com/JuliaGPU/KernelAbstractions.jl"}

From cc502022c0772303d5541921c393ee1396ed3ec1 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Sun, 27 Sep 2026 14:53:40 +0200
Subject: [PATCH 13/14] KA: fix tuned launches with range and empty ndranges

Since KernelAbstractions 0.10, an ndrange may contain ranges or be a
CartesianIndices, so compute the preliminary and the tuned workgroup size
from the extents of the ndrange rather than from the raw ndrange. Also return
before compiling when the ndrange is empty, and keep every workgroup
dimension at least 1, which avoids dividing by zero when tuning the
workgroup size for a zero-sized dimension.
---
 ext/AMDGPUKernelAbstractionsExt.jl | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/ext/AMDGPUKernelAbstractionsExt.jl b/ext/AMDGPUKernelAbstractionsExt.jl
index b9e163e30..0ce2a21b0 100644
--- a/ext/AMDGPUKernelAbstractionsExt.jl
+++ b/ext/AMDGPUKernelAbstractionsExt.jl
@@ -33,9 +33,10 @@ function KA.launch_config(kernel::KA.Kernel{ROCBackend}, ndrange, workgroupsize)
     end
 
     iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize && workgroupsize === nothing
+        extents = KA.NDIteration.extents(ndrange)
         workgroupsize = ntuple(
-            i -> i == 1 ? min(prod(ndrange), AMDGPU.Device._max_group_size) : 1,
-            length(ndrange))
+            i -> i == 1 ? min(prod(extents), AMDGPU.Device._max_group_size) : 1,
+            length(extents))
         KA.partition(kernel, ndrange, workgroupsize)
     else
         KA.partition(kernel, ndrange, workgroupsize)
@@ -47,7 +48,7 @@ end
 function threads_to_workgroupsize(threads, ndrange)
     total = 1
     return map(ndrange) do n
-        x = min(div(threads, total), n)
+        x = max(min(div(threads, total), n), 1)
         total *= x
         return x
     end
@@ -55,6 +56,9 @@ end
 
 function (obj::KA.Kernel{ROCBackend})(args...; ndrange=nothing, workgroupsize=nothing)
     ndrange, new_workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize)
+    # nothing to launch (or compile) for an empty ndrange
+    length(KA.blocks(iterspace)) == 0 && return
+
     ctx = KA.mkcontext(obj, ndrange, iterspace)
     if KA.workgroupsize(obj) <: KA.StaticSize
         maxthreads = prod(KA.get(KA.workgroupsize(obj)))
@@ -69,14 +73,13 @@ function (obj::KA.Kernel{ROCBackend})(args...; ndrange=nothing, workgroupsize=no
         isnothing(workgroupsize)
     if is_dynamic
         (; groupsize) = AMDGPU.launch_configuration(kernel)
-        new_workgroupsize = threads_to_workgroupsize(groupsize, ndrange)
+        new_workgroupsize = threads_to_workgroupsize(groupsize, KA.NDIteration.extents(ndrange))
         iterspace, dynamic = KA.partition(obj, ndrange, new_workgroupsize)
         ctx = KA.mkcontext(obj, ndrange, iterspace)
     end
 
     nblocks = length(KA.blocks(iterspace))
     nthreads = length(KA.workitems(iterspace))
-    nblocks == 0 && return
 
     kernel(ctx, args...; groupsize=nthreads, gridsize=nblocks)
     return

From 52ad7b221c78fd8bfaa0f5903107214e6315fede Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Mon, 28 Sep 2026 07:09:09 +0200
Subject: [PATCH 14/14] KernelInterface: report per-dimension launch limits

Implement KI.max_work_group_dims from the device's maxThreadsDim, cached in
HIPDevice since it is queried on every automatically-sized launch, and
KI.max_num_groups as the number of workgroups that fits the work-item grid
limit for any valid workgroup size. Requires KernelInterface 0.2.3.
---
 Project.toml            |  2 +-
 docs/src/api/devices.md |  1 +
 src/ROCKernels.jl       | 13 +++++++++++++
 src/hip/device.jl       | 11 ++++++++++-
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8bcff419e..c6aab5664 100644
--- a/Project.toml
+++ b/Project.toml
@@ -68,7 +68,7 @@ GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
 KernelAbstractions = "0.10"
-KernelInterface = "0.2.1"
+KernelInterface = "0.2.3"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
 Libdl = "1"
diff --git a/docs/src/api/devices.md b/docs/src/api/devices.md
index 20194f57f..39447b2b1 100644
--- a/docs/src/api/devices.md
+++ b/docs/src/api/devices.md
@@ -42,6 +42,7 @@ AMDGPU.device_id!
 AMDGPU.HIP.name
 AMDGPU.HIP.wavefrontsize
 AMDGPU.HIP.gcn_arch
+AMDGPU.HIP.max_workgroup_dims
 AMDGPU.HIP.device_id
 AMDGPU.HIP.properties
 ```
diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl
index 61fd3f1b4..bff115d77 100644
--- a/src/ROCKernels.jl
+++ b/src/ROCKernels.jl
@@ -99,6 +99,19 @@ end
 function KI.max_work_group_size(::ROCBackend)::Int
     Int(HIP.attribute(AMDGPU.HIP.device(), AMDGPU.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 KI.sub_group_size(::ROCBackend)::Int
     HIP.wavefrontsize(HIP.device())
 end
diff --git a/src/hip/device.jl b/src/hip/device.jl
index 61b9408c3..479aea449 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