KernelAbstractions 0.10

.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

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,17 @@ 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"
 
+[sources]
+AcceleratedKernels = {rev = "main", url = "https://github.com/JuliaGPU/AcceleratedKernels.jl"}
+
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
+AMDGPUKernelAbstractionsExt = "KernelAbstractions"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
@@ -62,7 +67,8 @@ ExprTools = "0.1"
 GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.10"
+KernelInterface = "0.2.3"
 LLVM = "9"
 LLVMDowngrader_jll = "0.11"
 Libdl = "1"
@@ -82,3 +88,6 @@ StaticArraysCore = "1"
 Statistics = "1"
 UnsafeAtomics = "0.3"
 julia = "1.10"
+
+[extras]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"

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
 ```

ext/AMDGPUKernelAbstractionsExt.jl

@@ -0,0 +1,108 @@
+module AMDGPUKernelAbstractionsExt
+
+import AMDGPU
+import AMDGPU.Device: @device_override
+using AMDGPU: GPUArrays, rocSPARSE, 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)
+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.argconvert(::KA.Kernel{ROCBackend}, arg) = AMDGPU.rocconvert(arg)
+
+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
+        extents = KA.NDIteration.extents(ndrange)
+        workgroupsize = ntuple(
+            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)
+    end
+
+    return ndrange, workgroupsize, iterspace, dynamic
+end
+
+function threads_to_workgroupsize(threads, ndrange)
+    total = 1
+    return map(ndrange) do n
+        x = max(min(div(threads, total), n), 1)
+        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)
+    # 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)))
+    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, 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))
+
+    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
+
+@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 @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
+    MArray{KA.__size(Dims), T}(undef)
+end
+
+end

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")
@@ -102,7 +103,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
@@ -147,10 +148,12 @@ 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
 
+
 include("precompile.jl")
 
 function __init__()