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/Project.toml b/Project.toml index 8bc9180ef..c6aab5664 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,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" 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/ext/AMDGPUKernelAbstractionsExt.jl b/ext/AMDGPUKernelAbstractionsExt.jl new file mode 100644 index 000000000..0ce2a21b0 --- /dev/null +++ b/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 diff --git a/src/AMDGPU.jl b/src/AMDGPU.jl index 636b3bf27..9e42cb45d 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") @@ -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__() diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl index 2315222fe..bff115d77 100644 --- a/src/ROCKernels.jl +++ b/src/ROCKernels.jl @@ -3,11 +3,12 @@ module ROCKernels 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,171 @@ 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) +function KI.record_event(::ROCBackend) + return HIP.HIPEvent(AMDGPU.stream()) +end -function KA.priority!(::ROCBackend, priority::Symbol) +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) +KI.ones(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ones(T, dims) + +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 +KI.argconvert(::ROCBackend, arg) = rocconvert(arg) - # partition checked that the ndrange's agreed - if KA.ndrange(kernel) <: KA.StaticSize - ndrange = nothing - end +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 - 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 +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 - return ndrange, workgroupsize, iterspace, dynamic + numworkgroups, workgroupsize = KI.auto_launch_sizes(obj, numworkgroups, workgroupsize, ndrange, max_work_group_size) + obj.kern(args...; groupsize = workgroupsize, gridsize = numworkgroups) + return nothing 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 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) -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) + return Int(min(max_work_items, groupsize)) +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 - - 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.sub_group_size(::ROCBackend)::Int + HIP.wavefrontsize(HIP.device()) 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.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(::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 KA.__index_Group_Linear(ctx) - return AMDGPU.Device.blockIdx().x +@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 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(::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 KA.__index_Local_Cartesian(ctx) - @inbounds KA.workitems(KA.__iterspace(ctx))[AMDGPU.Device.threadIdx().x] +@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 KA.__index_Group_Cartesian(ctx) - @inbounds KA.blocks(KA.__iterspace(ctx))[AMDGPU.Device.blockIdx().x] +@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 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(::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 @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/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) 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 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 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 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..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. diff --git a/test/kernelinterface_tests.jl b/test/kernelinterface_tests.jl new file mode 100644 index 000000000..018754687 --- /dev/null +++ b/test/kernelinterface_tests.jl @@ -0,0 +1,13 @@ +using Test +using AMDGPU + +import KernelInterface +include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl")) + +AMDGPU.allowscalar(false) + +@testset "kernelinterface" begin + +Testsuite.testsuite( + ROCBackend, "ROCM", AMDGPU, ROCArray, AMDGPU.ROCDeviceArray) +end