From 92c58378917dab2abe65af474e22c62de2219b57 Mon Sep 17 00:00:00 2001
From: William F Godoy <williamfgc@yahoo.com>
Date: Wed, 17 Apr 2024 14:01:52 -0400
Subject: [PATCH] Add suport for Atomix.@atomic

Atomix is supported by Threads, CUDA and AMDGPU back ends
In the roadmap for oneAPI
Added unit test
---
 Project.toml          |  1 +
 src/JACC.jl           |  3 ++-
 test/tests_amdgpu.jl  | 20 ++++++++++++++++++++
 test/tests_cuda.jl    | 20 ++++++++++++++++++++
 test/tests_threads.jl | 24 +++++++++++++++++++++---
 5 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index 1b8ec56..512a77b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,7 @@ authors = ["pedrovalerolara <valerolarap@ornl.gov>", "williamfgc <williamfgc@yah
 version = "0.0.3"
 
 [deps]
+Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 
 [weakdeps]
diff --git a/src/JACC.jl b/src/JACC.jl
index 7e5a02d..5e0453c 100644
--- a/src/JACC.jl
+++ b/src/JACC.jl
@@ -1,11 +1,12 @@
 __precompile__(false)
 module JACC
 
+import Atomix: @atomic
 # module to set back end preferences 
 include("JACCPreferences.jl")
 include("helper.jl")
 
-export Array
+export Array, @atomic
 export parallel_for
 
 global Array
diff --git a/test/tests_amdgpu.jl b/test/tests_amdgpu.jl
index fff64f1..dc2093b 100644
--- a/test/tests_amdgpu.jl
+++ b/test/tests_amdgpu.jl
@@ -52,3 +52,23 @@ end
 
 	@test Array(x_device) ≈ x_expected rtol = 1e-1
 end
+
+@testset "AtomicCounter" begin
+
+	function axpy_counter!(i, alpha, x, y, counter)
+		@inbounds x[i] += alpha * y[i]
+		JACC.@atomic counter[1] += 1
+	end
+
+	N = Int32(10)
+	# Generate random vectors x and y of length N for the interval [0, 100]
+	alpha = 2.5
+
+	x = JACC.Array(round.(rand(Float32, N) * 100))
+	y = JACC.Array(round.(rand(Float32, N) * 100))
+	counter = JACC.Array{Int32}([0])
+	JACC.parallel_for(N, axpy_counter!, alpha, x, y, counter)
+
+	@test Array(counter)[1] == N
+end
+
diff --git a/test/tests_cuda.jl b/test/tests_cuda.jl
index 1a27386..81bfe95 100644
--- a/test/tests_cuda.jl
+++ b/test/tests_cuda.jl
@@ -53,6 +53,26 @@ end
 	@test Array(x_device) ≈ x_expected rtol = 1e-1
 end
 
+@testset "AtomicCounter" begin
+
+	function axpy_counter!(i, alpha, x, y, counter)
+		@inbounds x[i] += alpha * y[i]
+		JACC.@atomic counter[1] += 1
+	end
+
+	N = Int32(10)
+	# Generate random vectors x and y of length N for the interval [0, 100]
+	alpha = 2.5
+
+	x = JACC.Array(round.(rand(Float32, N) * 100))
+	y = JACC.Array(round.(rand(Float32, N) * 100))
+	counter = JACC.Array{Int32}([0])
+	JACC.parallel_for(N, axpy_counter!, alpha, x, y, counter)
+
+	@test Array(counter)[1] == N
+end
+
+
 
 
 
diff --git a/test/tests_threads.jl b/test/tests_threads.jl
index 831e451..b64ac38 100644
--- a/test/tests_threads.jl
+++ b/test/tests_threads.jl
@@ -52,6 +52,27 @@ end
 	@test x_host_JACC ≈ x_expected rtol = 1e-1
 end
 
+@testset "AtomicCounter" begin
+
+	function axpy_counter!(i, alpha, x, y, counter)
+		@inbounds x[i] += alpha * y[i]
+		JACC.@atomic counter[1] += 1
+	end
+
+	N = Int32(10)
+	# Generate random vectors x and y of length N for the interval [0, 100]
+	alpha = 2.5
+	counter = zeros(Int32, 1)
+
+	x_device = JACC.Array(round.(rand(Float32, N) * 100))
+	y_device = JACC.Array(round.(rand(Float32, N) * 100))
+	counter = JACC.Array{Int32}([0])
+	JACC.parallel_for(N, axpy_counter!, alpha, x_device, y_device, counter)
+
+	@test counter[1] == N
+
+end
+
 @testset "CG" begin
 
 	function matvecmul(i, a1, a2, a3, x, y, SIZE)
@@ -112,9 +133,6 @@ end
 		ccond = JACC.parallel_reduce(SIZE, dot, r, r)
 		global cond = ccond
 		p = copy(r_aux)
-
-		println(cond)
-
 	end
 	@test cond[1, 1] <= 1e-14
 end