Skip to content

Commit d795583

Browse files
authored
Add atomics support for ROCKernels (#304)
1 parent 2619d61 commit d795583

File tree

5 files changed

+29
-13
lines changed

5 files changed

+29
-13
lines changed

examples/histogram.jl

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# INCLUDE ROCM
12
using KernelAbstractions, Test
23
using KernelAbstractions: @atomic, @atomicswap, @atomicreplace
34
include(joinpath(@__DIR__, "utils.jl")) # Load backend
@@ -65,8 +66,10 @@ function histogram!(histogram_output, input;
6566

6667
if isa(input, Array)
6768
kernel! = histogram_kernel!(CPU(), numcores)
68-
else
69+
elseif has_cuda
6970
kernel! = histogram_kernel!(CUDADevice(), numthreads)
71+
elseif has_rocm
72+
kernel! = histogram_kernel!(ROCDevice(), numthreads)
7073
end
7174

7275
kernel!(histogram_output, input, ndrange=size(input))
@@ -96,16 +99,23 @@ end
9699
@test isapprox(CPU_2_histogram, histogram_2_baseline)
97100
end
98101

99-
if has_cuda_gpu()
102+
if has_cuda && has_cuda_gpu()
100103
CUDA.allowscalar(false)
101-
102-
GPU_rand_input = CuArray(rand_input)
103-
GPU_linear_input = CuArray(linear_input)
104-
GPU_2_input = CuArray(all_2)
105-
106-
GPU_rand_histogram = CuArray(zeros(Int, 128))
107-
GPU_linear_histogram = CuArray(zeros(Int, 1024))
108-
GPU_2_histogram = CuArray(zeros(Int, 2))
104+
GPUArray = CuArray
105+
has_gpu = true
106+
elseif has_rocm && AMDGPU.functional()
107+
AMDGPU.allowscalar(false)
108+
GPUArray = ROCArray
109+
has_gpu = true
110+
end
111+
if has_gpu
112+
GPU_rand_input = GPUArray(rand_input)
113+
GPU_linear_input = GPUArray(linear_input)
114+
GPU_2_input = GPUArray(all_2)
115+
116+
GPU_rand_histogram = GPUArray(zeros(Int, 128))
117+
GPU_linear_histogram = GPUArray(zeros(Int, 1024))
118+
GPU_2_histogram = GPUArray(zeros(Int, 2))
109119

110120
wait(histogram!(GPU_rand_histogram, GPU_rand_input))
111121
wait(histogram!(GPU_linear_histogram, GPU_linear_input))

lib/ROCKernels/Project.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
88
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
99
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
1010
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
11+
UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
1112

1213
[compat]
1314
AMDGPU = "0.3.2"

lib/ROCKernels/src/ROCKernels.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import StaticArrays: MArray
77
import Adapt
88
import KernelAbstractions
99

10+
import UnsafeAtomicsLLVM
11+
1012
export ROCDevice
1113

1214
KernelAbstractions.get_device(::AMDGPU.ROCArray) = ROCDevice()

test/examples.jl

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@ function find_sources(path::String, sources=String[])
99
sources
1010
end
1111

12-
function examples_testsuite()
12+
function examples_testsuite(backend_str)
1313
@testset "examples" begin
1414
examples_dir = joinpath(@__DIR__, "..", "examples")
1515
examples = find_sources(examples_dir)
1616
filter!(file -> readline(file) != "# EXCLUDE FROM TESTING", examples)
17+
if backend_str == "ROCM"
18+
filter!(file -> occursin("# INCLUDE ROCM", String(read(file))), examples)
19+
end
1720

1821
@testset "$(basename(example))" for example in examples
1922
code = """

test/testsuite.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,9 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT)
6969
convert_testsuite(backend, AT)
7070
end
7171

72-
if backend_str == "CUDA"
72+
if backend_str == "CUDA" || backend_str == "ROCM"
7373
@testset "Examples" begin
74-
examples_testsuite()
74+
examples_testsuite(backend_str)
7575
end
7676
end
7777
end

0 commit comments

Comments
 (0)