Support KA groupreduce API

Author: pxl-th

src/ROCKernels.jl

@@ -166,4 +166,8 @@ end
     # TODO
 end
 
+@device_override @inline function KA.__shfl_down(val, offset)
+    AMDGPU.Device.shfl_down(val, offset)
+end
+
 end

src/device/gcn/wavefront.jl

@@ -57,22 +57,20 @@ end
 """
     wfred(op::Function, val::T) where T -> T
 
-Performs a wavefront-wide reduction on `val` in each lane, and returns the
-result. A limited subset of functions are available to be passed as `op`. When
-`op` is one of `(+, max, min, &, |, ⊻)`, `T` may be
-`<:Union{Cint, Clong, Cuint, Culong}`. When `op` is one of `(+, max, min)`,
-`T` may also be `<:Union{Float32, Float64}`.
+Performs a wavefront-wide reduction on `val` in each lane, and returns the result.
+A limited subset of functions are available to be passed as `op`.
+When `op` is one of `(+, max, min, &, |, ⊻)`, `T` may be `<:Union{Cint, Clong, Cuint, Culong}`.
+When `op` is one of `(+, max, min)`, `T` may also be `<:Union{Float32, Float64}`.
 """
 wfred
 
 """
     wfscan(op::Function, val::T) where T -> T
 
 Performs a wavefront-wide scan on `val` in each lane, and returns the
-result. A limited subset of functions are available to be passed as `op`. When
-`op` is one of `(+, max, min, &, |, ⊻)`, `T` may be
-`<:Union{Cint, Clong, Cuint, Culong}`. When `op` is one of `(+, max, min)`,
-`T` may also be `<:Union{Float32, Float64}`.
+result. A limited subset of functions are available to be passed as `op`.
+When `op` is one of `(+, max, min, &, |, ⊻)`, `T` may be `<:Union{Cint, Clong, Cuint, Culong}`.
+When `op` is one of `(+, max, min)`, `T` may also be `<:Union{Float32, Float64}`.
 """
 wfscan
 

t.jl

@@ -0,0 +1,24 @@
+using AMDGPU
+using KernelAbstractions
+import KernelAbstractions as KA
+
+@kernel function ker!(y, x, neutral, op)
+    i = @index(Global)
+    val = i > length(x) ? neutral : x[i]
+
+    res = KA.@groupreduce(:warp, op, val, neutral)
+    if i == 1
+        y[1] = res
+    end
+end
+
+function main()
+    n = 256
+    x = ROCArray(ones(Int64, n))
+    y = ROCArray(zeros(Int64, 1))
+
+    ker!(ROCBackend(), n)(y, x, 0, +; ndrange=length(x))
+    @show y
+    return
+end
+main()

test/runtests.jl

@@ -110,24 +110,25 @@ AMDGPU.versioninfo()
 
 @info "Test suite info"
 data = String["$np" "$(AMDGPU.device())" join(TARGET_TESTS, ", ");]
-PrettyTables.pretty_table(data; header=[
-    "Workers", "Device", "Tests"], crop=:none)
+PrettyTables.pretty_table(data; header=["Workers", "Device", "Tests"], crop=:none)
 
 runtests(AMDGPU; nworkers=np, nworker_threads=1, testitem_timeout=60 * 30) do ti
+    return ti.name == "kernelabstractions"
+
     for tt in TARGET_TESTS
         startswith(ti.name, tt) && return true
     end
     return false
 end
 
-if "core" in TARGET_TESTS && Sys.islinux()
-    @info "Testing `Hostcalls` on the main thread."
-    @testset "Hostcalls" begin
-        include("device/hostcall.jl")
-
-        # TODO 1.11 fails
-        if VERSION < v"1.11-"
-            include("device/output.jl")
-        end
-    end
-end
+# if "core" in TARGET_TESTS && Sys.islinux()
+#     @info "Testing `Hostcalls` on the main thread."
+#     @testset "Hostcalls" begin
+#         include("device/hostcall.jl")
+
+#         # TODO 1.11 fails
+#         if VERSION < v"1.11-"
+#             include("device/output.jl")
+#         end
+#     end
+# end