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 8bc9180ef..747de04da 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,15 +39,22 @@ 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" +[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" @@ -62,7 +69,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" @@ -82,3 +90,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..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/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/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..93f7b98e5 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 @@ -82,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. @@ -102,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 @@ -147,6 +149,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..9aefe1fbb 100644 --- a/src/ROCKernels.jl +++ b/src/ROCKernels.jl @@ -4,198 +4,256 @@ 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)`. -""" -struct ROCBackend <: KA.GPU end +obtain it from an array with `KernelInterface.get_backend(::ROCArray)`. -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.")) +Printing from a kernel (`KernelAbstractions.@print`) is not supported: it does nothing. +""" +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) + (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) +# 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 -function KA.priority!(::ROCBackend, priority::Symbol) +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( "Priority `$priority` must be one of `:high`, `:normal`, `:low`.") AMDGPU.priority!(priority) + return end -function KA.copyto!(::ROCBackend, A, B) - GC.@preserve A B begin - copyto!(A, 1, B, 1, length(A)) +## memory operations + +KI.unsafe_free!(x::AMDGPU.ROCArray) = AMDGPU.unsafe_free!(x) +KI.allocate(::ROCBackend, ::Type{T}, dims::Tuple) where T = AMDGPU.ROCArray{T}(undef, dims) + +# dense arrays, and contiguous views of them +const ContiguousArray{T} = Union{DenseArray{T}, Base.FastContiguousSubArray{T, <:Any, <:DenseArray}} +on_device(A::ContiguousArray) = parent(A) isa AMDGPU.ROCArray + +function KI.copyto!(::ROCBackend, A::ContiguousArray{T}, B::ContiguousArray{T}) where T + length(A) == length(B) || + throw(ArgumentError("Arrays must have the same length, got $(length(A)) and $(length(B))")) + if isbitstype(T) && (on_device(A) || on_device(B)) + # queued on the task's stream, after the work queued before it + GC.@preserve A B begin + AMDGPU.Mem.memcpy!(pointer(A), pointer(B), length(A) * AMDGPU.aligned_sizeof(T); + stream=AMDGPU.stream()) + end + else + # host-to-host copies, and bits unions, whose type tags are stored separately. + # queued work may still access host arrays (e.g. a copy from the device), so wait + # for it first. + on_device(A) || on_device(B) || AMDGPU.synchronize() + copyto!(A, B) end - return + return A end +KI.copyto!(::ROCBackend, A, B) = + throw(ArgumentError("KernelInterface.copyto! only supports contiguous arrays of the same element type, got $(typeof(A)) and $(typeof(B))")) -function KA.pagelock!(::ROCBackend, x::Array) +function KI.pagelock!(::ROCBackend, x::Array) AMDGPU.Mem.pin(pointer(x), sizeof(x)) return end -function KA.launch_config(kernel::KA.Kernel{ROCBackend}, ndrange, workgroupsize) - if ndrange isa Integer - ndrange = (ndrange,) - end - if workgroupsize isa Integer - workgroupsize = (workgroupsize, ) - end - - # partition checked that the ndrange's agreed - if KA.ndrange(kernel) <: KA.StaticSize - ndrange = nothing - end +## kernel launch - iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize && workgroupsize === nothing - 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) +# `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 + +## COV_EXCL_START + +## indexing -# Indexing. +# computed with `% T`, which unlike `T(x)` has no error path -@device_override @inline function KA.__index_Local_Linear(ctx) - return 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_Linear(ctx) - return 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_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_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.__index_Local_Cartesian(ctx) - @inbounds KA.workitems(KA.__iterspace(ctx))[AMDGPU.Device.threadIdx().x] +@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 -@device_override @inline function KA.__index_Group_Cartesian(ctx) - @inbounds KA.blocks(KA.__iterspace(ctx))[AMDGPU.Device.blockIdx().x] +## 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 -@device_override @inline function KA.__index_Global_Cartesian(ctx) - return @inbounds KA.expand(KA.__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x) +@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 -@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 +# 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 -# Shared memory. +@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 -@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))) +## shared memory + +@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 KI.barrier() + Device.sync_workgroup() +end -@device_override @inline function KA.__synchronize() - AMDGPU.Device.sync_workgroup() +@device_override @inline function KI.sub_group_barrier() + Device.sync_wavefront() end -@device_override @inline function KA.__print(args...) - # TODO +@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 + +## COV_EXCL_STOP + end 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)) 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/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/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 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..16c56d2c0 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" @@ -27,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"} 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 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..675b8ae23 --- /dev/null +++ b/test/kernelinterface_tests.jl @@ -0,0 +1,125 @@ +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 + +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() +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, 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 + +@testset "versioninfo" begin + @test occursin("AMDGPU.jl", sprint(KI.versioninfo, backend)) +end + +end