From 895e104648d5c9d1b1faa334e8dff969f8747e94 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:03:23 +0200 Subject: [PATCH 1/8] 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 b57937911..97b02c6b7 100644 --- a/src/device/gcn/memory_static.jl +++ b/src/device/gcn/memory_static.jl @@ -24,8 +24,10 @@ 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 + # every call site gets its own allocation. local memory can't be + # initialized, so use an `undef` initializer, which the backend accepts. + linkage!(gv, LLVM.API.LLVMInternalLinkage) + initializer!(gv, UndefValue(gv_typ)) elseif as == AS.Private linkage!(gv, LLVM.API.LLVMInternalLinkage) initializer!(gv, null(gv_typ)) From 9ba7b6b2269df00414f25fa5a6ad4a69770caf57 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:03:51 +0200 Subject: [PATCH 2/8] 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 20194f57f..167a7dfad 100644 --- a/docs/src/api/devices.md +++ b/docs/src/api/devices.md @@ -41,6 +41,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 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 From 3d6d5db7fc45df49032ae4f9c48489e7be7d75ab Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:13:15 +0200 Subject: [PATCH 3/8] 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. A `HIPKernel` can't be launched with a tuple yet (JuliaGPU/AMDGPU.jl#1115), so it splats them for now, which is slow for more than 32 arguments. 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 | 262 +++++++++++++++-------------- src/utils.jl | 3 +- test/Project.toml | 1 + test/kernelabstractions_tests.jl | 56 +++++- test/kernelinterface_tests.jl | 65 +++++++ 8 files changed, 304 insertions(+), 134 deletions(-) create mode 100644 ext/AMDGPUKernelAbstractionsExt.jl create mode 100644 test/kernelinterface_tests.jl diff --git a/Project.toml b/Project.toml index 8bc9180ef..163e36572 100644 --- a/Project.toml +++ b/Project.toml @@ -18,7 +18,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" @@ -39,12 +39,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" @@ -62,7 +64,8 @@ ExprTools = "0.1" GPUArrays = "11.5.14" GPUCompiler = "2.9" GPUToolbox = "3" -KernelAbstractions = "0.9.2" +KernelAbstractions = "0.10" +KernelInterface = "0.4" LLVM = "9" LLVMDowngrader_jll = "0.11" 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 636b3bf27..89698bce0 100644 --- a/src/AMDGPU.jl +++ b/src/AMDGPU.jl @@ -12,6 +12,7 @@ using Preferences using Printf import AcceleratedKernels as AK +import KernelInterface import UnsafeAtomics import Atomix import Atomix: @atomic, @atomicswap, @atomicreplace @@ -147,6 +148,7 @@ function Atomix.modify!(ref::ROCIndexableRef, op::OP, x, ord) where OP <: Union{ GC.@preserve root UnsafeAtomics.modify!(ptr, op, x, ord, syncscope_agent) end +# KernelInterface include("ROCKernels.jl") import .ROCKernels: ROCBackend export ROCBackend diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl index 2315222fe..b47aff960 100644 --- a/src/ROCKernels.jl +++ b/src/ROCKernels.jl @@ -4,198 +4,200 @@ 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 +## kernel launch - # 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 +# 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 -function threads_to_workgroupsize(threads, ndrange) - total = 1 - return map(ndrange) do n - x = min(div(threads, total), n) - total *= x - return x +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 -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 +# XXX: this splats the arguments, which is slow for more than 32 of them. pass the tuple on +# once a `HIPKernel` can be launched with one (JuliaGPU/AMDGPU.jl#1115). +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 - 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) + 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) + kernel(args...; groupsize=items, gridsize=groups, kwargs..., stream) 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) +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. +## COV_EXCL_START -@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 +## 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 +## synchronization and printing -# Other. - -@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 2c8220429..b582b4425 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. @@ -27,4 +43,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 c7aa3f2108bca8709c302ffd124d0c54a27aaad8 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:14:44 +0200 Subject: [PATCH 4/8] 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 | 3 ++- src/device/gcn/synchronization.jl | 18 +++++++++++++++++- test/device/synchronization.jl | 18 ++++++++++++++++++ 4 files changed, 38 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 89698bce0..93f7b98e5 100644 --- a/src/AMDGPU.jl +++ b/src/AMDGPU.jl @@ -83,6 +83,7 @@ Base.Experimental.@MethodTable(method_table) #needs to be before Device since sync uses this const syncscope_agent = UnsafeAtomics.Internal.LLVMSyncScope{:agent}() const syncscope_workgroup = UnsafeAtomics.Internal.LLVMSyncScope{:workgroup}() +const syncscope_wavefront = UnsafeAtomics.Internal.LLVMSyncScope{:wavefront}() # 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 +104,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..bdad1c110 100644 --- a/src/device/gcn/synchronization.jl +++ b/src/device/gcn/synchronization.jl @@ -1,6 +1,6 @@ for ord in UnsafeAtomics.Internal.orderings - for sync in (AMDGPU.syncscope_agent, AMDGPU.syncscope_workgroup) + for sync in (AMDGPU.syncscope_agent, AMDGPU.syncscope_workgroup, AMDGPU.syncscope_wavefront) @eval @device_function function UnsafeAtomics.fence(::$(typeof(ord)), ::$(typeof(sync))) Base.llvmcall( $(""" @@ -29,6 +29,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 9f29c793f3c97c30a6c9b7675d6b79573bd28b93 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:20:32 +0200 Subject: [PATCH 5/8] 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 b47aff960..b7cea55be 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( @@ -155,6 +159,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 @@ -181,6 +189,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} @@ -195,6 +235,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 3264aa54f067a527241a4e691f27ad3fd954a7de Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:21:53 +0200 Subject: [PATCH 6/8] 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 b7cea55be..1b52eb990 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 65db95d93dfe4242aaa87b4ead6fd54fa4a9d2b3 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:22:59 +0200 Subject: [PATCH 7/8] 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 1b52eb990..9aefe1fbb 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 a3f1d4c2b69e31d0ee1c1970d2633d8c925c7626 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:24:25 +0200 Subject: [PATCH 8/8] [TEMP] Get KernelAbstractions 0.10, KernelInterface 0.4 and AcceleratedKernels from git KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet, and the registered AcceleratedKernels doesn't support KernelAbstractions 0.10. 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 and Enzyme jobs move to Julia 1.12. Drop this commit once they are registered. --- .buildkite/pipeline.yml | 51 +++++++++++++++++++++++++++++++++++++---- Project.toml | 8 +++++++ test/Project.toml | 2 ++ 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 30727c02b..a0c2da48a 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -73,8 +73,6 @@ steps: matrix: setup: version: - - "1.10" - - "1.11" - "1.12" test_args: - "--verbose" @@ -83,6 +81,48 @@ steps: (build.message =~ /\[only [^\]]*(tests|julia)/ || build.message !~ /\[only / && !build.pull_request.draft) + # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet, and the + # registered AcceleratedKernels doesn't support KernelAbstractions 0.10. 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. + - label: "Julia {{matrix.version}}" + matrix: + setup: + version: + - "1.10" + - "1.11" + 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 + git clone --depth 1 https://github.com/JuliaGPU/AcceleratedKernels.jl ak + julia --project=test -e ' + using Pkg + if VERSION < v"1.11" + Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface", "ak")]) + 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" + - <<: *julia matrix: setup: @@ -99,8 +139,8 @@ steps: matrix: setup: version: - - "1.10" - - "1.11" + # XXX: needs Julia 1.12+ to pick up KernelAbstractions 0.10 from [sources] + - "1.12" plugins: - JuliaCI/julia#v1: version: "{{matrix.version}}" @@ -125,7 +165,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 163e36572..747de04da 100644 --- a/Project.toml +++ b/Project.toml @@ -50,6 +50,11 @@ AMDGPUKernelAbstractionsExt = "KernelAbstractions" AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR" AMDGPUSpecialFunctionsExt = "SpecialFunctions" +[sources] +AcceleratedKernels = {url = "https://github.com/JuliaGPU/AcceleratedKernels.jl", rev = "main"} +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" @@ -85,3 +90,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..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"}