Add atomics support for ROCKernels (#304)

Author: jpsamaroo

examples/histogram.jl

@@ -1,3 +1,4 @@
+# INCLUDE ROCM
 using KernelAbstractions, Test
 using KernelAbstractions: @atomic, @atomicswap, @atomicreplace
 include(joinpath(@__DIR__, "utils.jl")) # Load backend
@@ -65,8 +66,10 @@ function histogram!(histogram_output, input;
 
     if isa(input, Array)
         kernel! = histogram_kernel!(CPU(), numcores)
-    else
+    elseif has_cuda
         kernel! = histogram_kernel!(CUDADevice(), numthreads)
+    elseif has_rocm
+        kernel! = histogram_kernel!(ROCDevice(), numthreads)
     end
 
     kernel!(histogram_output, input, ndrange=size(input))
@@ -96,16 +99,23 @@ end
         @test isapprox(CPU_2_histogram, histogram_2_baseline)
     end
 
-    if has_cuda_gpu()
+    if has_cuda && has_cuda_gpu()
         CUDA.allowscalar(false)
-
-        GPU_rand_input = CuArray(rand_input)
-        GPU_linear_input = CuArray(linear_input)
-        GPU_2_input = CuArray(all_2)
-
-        GPU_rand_histogram = CuArray(zeros(Int, 128))
-        GPU_linear_histogram = CuArray(zeros(Int, 1024))
-        GPU_2_histogram = CuArray(zeros(Int, 2))
+        GPUArray = CuArray
+        has_gpu = true
+    elseif has_rocm && AMDGPU.functional()
+        AMDGPU.allowscalar(false)
+        GPUArray = ROCArray
+        has_gpu = true
+    end
+    if has_gpu
+        GPU_rand_input = GPUArray(rand_input)
+        GPU_linear_input = GPUArray(linear_input)
+        GPU_2_input = GPUArray(all_2)
+
+        GPU_rand_histogram = GPUArray(zeros(Int, 128))
+        GPU_linear_histogram = GPUArray(zeros(Int, 1024))
+        GPU_2_histogram = GPUArray(zeros(Int, 2))
 
         wait(histogram!(GPU_rand_histogram, GPU_rand_input))
         wait(histogram!(GPU_linear_histogram, GPU_linear_input))

lib/ROCKernels/Project.toml

@@ -8,6 +8,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 
 [compat]
 AMDGPU = "0.3.2"

lib/ROCKernels/src/ROCKernels.jl

@@ -7,6 +7,8 @@ import StaticArrays: MArray
 import Adapt
 import KernelAbstractions
 
+import UnsafeAtomicsLLVM
+
 export ROCDevice
 
 KernelAbstractions.get_device(::AMDGPU.ROCArray) = ROCDevice()

test/examples.jl

@@ -9,11 +9,14 @@ function find_sources(path::String, sources=String[])
     sources
 end
 
-function examples_testsuite()
+function examples_testsuite(backend_str)
 @testset "examples" begin
     examples_dir = joinpath(@__DIR__, "..", "examples")
     examples = find_sources(examples_dir)
     filter!(file -> readline(file) != "# EXCLUDE FROM TESTING", examples)
+    if backend_str == "ROCM"
+        filter!(file -> occursin("# INCLUDE ROCM", String(read(file))), examples)
+    end
 
     @testset "$(basename(example))" for example in examples
         code = """

test/testsuite.jl

@@ -69,9 +69,9 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT)
         convert_testsuite(backend, AT)
     end
 
-    if backend_str == "CUDA"
+    if backend_str == "CUDA" || backend_str == "ROCM"
         @testset "Examples" begin
-            examples_testsuite()
+            examples_testsuite(backend_str)
         end
     end
 end