Merge pull request #817

support cooperative groups

Author: luraess

src/device/gcn/synchronization.jl

@@ -69,3 +69,7 @@ and returns non-zero if and only if predicate evaluates to non-zero for any of t
 end
 
 @inline __not(x::Cint)::Cint = ifelse(iszero(x), one(x), zero(x))
+
+@inline function sync_grid()::Cvoid
+    ccall("extern __ockl_grid_sync", llvmcall, Cvoid, ())
+end

src/highlevel.jl

@@ -94,7 +94,7 @@ rocconvert(arg) = adapt(Runtime.Adaptor(), arg)
 
 const MACRO_KWARGS = [:launch]
 const COMPILER_KWARGS = [:name, :unsafe_fp_atomics]
-const LAUNCH_KWARGS = [:gridsize, :groupsize, :shmem, :stream]
+const LAUNCH_KWARGS = [:gridsize, :groupsize, :shmem, :stream, :cooperative]
 
 """
     @roc [kwargs...] func(args...)

src/runtime/hip-execution.jl

@@ -117,13 +117,20 @@ function launch(
     fun::HIP.HIPFunction, args::Vararg{Any, N};
     gridsize = 1, groupsize = 1,
     shmem::Integer = 0, stream::HIP.HIPStream,
+    cooperative = false,
 ) where N
     gd = gridsize isa ROCDim3 ? gridsize : ROCDim3(gridsize)
     bd = groupsize isa ROCDim3 ? groupsize : ROCDim3(groupsize)
     pack_arguments(args...) do kernel_params
-        HIP.hipModuleLaunchKernel(
-            fun, gd.x, gd.y, gd.z, bd.x, bd.y, bd.z,
-            shmem, stream, kernel_params, C_NULL)
+        if cooperative
+            HIP.hipModuleLaunchCooperativeKernel(
+                fun, gd.x, gd.y, gd.z, bd.x, bd.y, bd.z,
+                shmem, stream, kernel_params)
+        else
+            HIP.hipModuleLaunchKernel(
+                fun, gd.x, gd.y, gd.z, bd.x, bd.y, bd.z,
+                shmem, stream, kernel_params, C_NULL)
+        end
     end
 
     AMDGPU.LAUNCH_BLOCKING[] && AMDGPU.synchronize(stream)

test/device/launch.jl

@@ -130,3 +130,32 @@ end
     #     @test occ1 != occ2
     # end
 end
+
+if !iszero(AMDGPU.HIP.properties(AMDGPU.device()).cooperativeLaunch)
+    @testset "Cooperative Groups" begin
+        function test_kernel!(x)
+            block_row, block_col = workgroupIdx().x, workgroupIdx().y
+
+            for diag in 2:(gridGroupDim().x + gridGroupDim().y)
+                if block_row + block_col == diag
+                    if diag == 2
+                        x[block_col, block_row] = 1
+                    else
+                        x[block_col, block_row] = 0
+                        for I in CartesianIndices(x)
+                            i, j = Tuple(I)
+                            i + j == diag - 1 || continue
+                            x[block_col, block_row] += x[I]
+                        end
+                    end
+                end
+                AMDGPU.Device.sync_grid()
+            end
+        end
+
+        n = 4
+        x = ROCArray{Int}(undef, n, n)
+        @roc groupsize = (1, 1, 1) gridsize = (n, n, 1) cooperative = true test_kernel!(x)
+        @test Array(x) == [1 1 2 6; 1 2 6 24; 2 6 24 72; 6 24 72 144]
+    end
+end