diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 04695e31f..04ffb88fe 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -96,7 +96,7 @@ steps: (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 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka julia --project=test -e ' using Pkg if VERSION < v"1.11" diff --git a/Project.toml b/Project.toml index 000005021..cd28d814e 100644 --- a/Project.toml +++ b/Project.toml @@ -49,8 +49,8 @@ AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR" AMDGPUSpecialFunctionsExt = "SpecialFunctions" [sources] -KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"} -KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"} [compat] AMDGPU_LLVM_Backend_jll = "23" diff --git a/src/ROCKernels.jl b/src/ROCKernels.jl index 1603c8df7..7011556e9 100644 --- a/src/ROCKernels.jl +++ b/src/ROCKernels.jl @@ -51,9 +51,18 @@ end 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}}} +# wavefronts are formed from consecutive linear work-item indices, and execute independently +KI.supports_linear_subgroups(::ROCBackend) = true +KI.supports_independent_subgroups(::ROCBackend) = true +# the functions with a `width` only need the lanes of the caller's segment: a `ds_bpermute` +# only involves the lanes it reads from, `llvm.amdgcn.ballot` only reports the lanes that +# execute it, and `sub_group_barrier` is wavefront-scope fences around +# `llvm.amdgcn.wave.barrier`, which (like HIP's tile `sync`) part of a wavefront can execute +KI.supports_independent_segments(::ROCBackend) = true +# the types the shuffles below support, by decomposing them into 32-bit `ds_bpermute`s; +# KernelInterface shuffles other `isbits` types (e.g. `Complex`) as `UInt32` words +const ShuffleTypes = Union{Bool, Base.BitInteger, Base.IEEEFloat} +KI.supports_shuffle(::ROCBackend, ::Type{T}) where {T <: ShuffleTypes} = true function KI.priority!(::ROCBackend, priority::Symbol) priority ∉ (:high, :normal, :low) && error( @@ -219,6 +228,8 @@ end return min(ws, workgroup_items() - (linear_workitem_id() ÷ ws) * ws) % T end +# a constant: `wavefrontsize` is folded to the wavefront size the kernel is compiled for, see +# `fold_wavefrontsize!` @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 @@ -245,10 +256,54 @@ end Device.sync_wavefront() end -@device_override function KI.shfl_down(val::T, offset::Integer) where T - @inline Device.shfl_down(val, offset % Cint) +## communication + +# read `val` from the work-item in the 0-based hardware lane `lane`. `ds_bpermute` only uses +# the low bits of the address, so lanes out of range read some other lane instead of +# trapping. unlike `Device.shfl` etc., the lanes are the hardware lanes (see +# `hardware_lane`), not `activelane`. +@inline function bpermute_lane(val, lane::Cint) + return Device._shfl(x -> Device.bpermute(lane << 0x2, x), val) +end + +@inline lane_id() = hardware_lane() % Cint + +@device_override @inline KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} = + bpermute_lane(val, (lane % Cint) - Cint(1)) + +# where the source lane is past the wavefront, `shfl_down` and `shfl_up` read from the +# work-item itself, like CUDA's shuffles (rather than from the lane `ds_bpermute` wraps around +# to). The offset is compared before it is narrowed, so that large offsets do too. +# `wavefrontsize` is folded to a constant, see `fold_wavefrontsize!`. +@device_override @inline function KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = lane_id() + ws = Device.wavefrontsize() % Cint + return bpermute_lane(val, ifelse(offset < ws - lane, lane + (offset % Cint), lane)) +end + +@device_override @inline function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = lane_id() + return bpermute_lane(val, ifelse(offset <= lane, lane - (offset % Cint), lane)) end +# `mask` is below the wavefront size, so the source lane is in the wavefront +@device_override @inline KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} = + bpermute_lane(val, lane_id() ⊻ (mask % Cint)) + +# the shuffles with a `width` use KernelInterface's fallbacks, a `ds_bpermute` from a lane +# computed with a few integer operations, like `Device.shfl` etc. (which use `activelane`). +# so do the votes, collectives and barrier with a `width`, which use the functions below +# and the ones without a `width`. + +# `ballot` only sets the bits of active lanes, i.e. of the work-items of the sub-group, and +# its result is uniform. `wavefrontsize`, which selects the 32- or 64-bit ballot, is folded +# to the wavefront size the kernel is compiled for (see `fold_wavefrontsize!`). +@device_override @inline KI.sub_group_ballot(pred::Bool) = Device.ballot(pred) + +@device_override @inline KI.sub_group_any(pred::Bool) = Device.ballot(pred) != 0 + +@device_override @inline KI.sub_group_all(pred::Bool) = Device.ballot(!pred) == 0 + # not supported, see the `ROCBackend` docstring @device_override @inline KI._print(args...) = nothing diff --git a/src/device/gcn/math.jl b/src/device/gcn/math.jl index df1cc231f..d45b2e6f5 100644 --- a/src/device/gcn/math.jl +++ b/src/device/gcn/math.jl @@ -60,12 +60,6 @@ for jltype in (Float64, Float32, Float16) @eval @device_override Base.fma(x::$jltype, y::$jltype, z::$jltype) = ccall( $("extern __ocml_fma_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype, $jltype), x, y, z) - @eval @device_override Base.min(x::$jltype, y::$jltype) = ccall( - $("extern __ocml_min_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y) - - @eval @device_override Base.max(x::$jltype, y::$jltype) = ccall( - $("extern __ocml_max_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y) - @eval @device_override Base.copysign(x::$jltype, y::$jltype) = ccall( $("extern __ocml_copysign_$(fntypes[jltype])"), llvmcall, $jltype, ($jltype, $jltype), x, y) diff --git a/test/Project.toml b/test/Project.toml index 16c56d2c0..dd64f57b0 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,5 +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"} +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"} diff --git a/test/device/math.jl b/test/device/math.jl index 618eada48..a2e433e85 100644 --- a/test/device/math.jl +++ b/test/device/math.jl @@ -20,6 +20,17 @@ using Base.FastMath end end +@testset "min/max" begin + # Julia's semantics: NaN propagates, and -0.0 < 0.0 + a = [NaN, 1, -0.0, 0.0, 2, NaN, Inf, -Inf] + b = [1, NaN, 0.0, -0.0, 3, NaN, NaN, 1] + for T in (Float16, Float32, Float64) + x, y = T.(a), T.(b) + @test isequal(Array(max.(ROCArray(x), ROCArray(y))), max.(x, y)) + @test isequal(Array(min.(ROCArray(x), ROCArray(y))), min.(x, y)) + end +end + @testset "Fast min/max" begin function ker!(x) x[1] = @fastmath max(x[1], zero(eltype(x)))