From b0650af33913f744c853322a8cfbf284089ce360 Mon Sep 17 00:00:00 2001
From: Juan Ignacio Polanco <juan-ignacio.polanco@cnrs.fr>
Date: Fri, 21 Nov 2025 10:41:46 +0100
Subject: [PATCH] Implement @atomicswap on all extensions

---
 ext/AtomixCUDAExt.jl       |  4 +++-
 ext/AtomixMetalExt.jl      |  4 +++-
 ext/AtomixOpenCLExt.jl     |  4 +++-
 ext/AtomixoneAPIExt.jl     |  4 +++-
 test/test_atomix_cuda.jl   | 12 ++++++++++++
 test/test_atomix_metal.jl  | 12 ++++++++++++
 test/test_atomix_oneapi.jl | 12 ++++++++++++
 test/test_atomix_opencl.jl | 12 ++++++++++++
 8 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/ext/AtomixCUDAExt.jl b/ext/AtomixCUDAExt.jl
index 5a79d62..2d82261 100644
--- a/ext/AtomixCUDAExt.jl
+++ b/ext/AtomixCUDAExt.jl
@@ -1,7 +1,7 @@
 # TODO: respect ordering
 module AtomixCUDAExt
 
-using Atomix: Atomix, IndexableRef
+using Atomix: Atomix, IndexableRef, right
 using CUDA: CUDA, CuDeviceArray
 
 const CuIndexableRef{Indexable<:CuDeviceArray} = IndexableRef{Indexable}
@@ -48,6 +48,8 @@ end
             CUDA.atomic_min!(ptr, x)
         elseif op === max
             CUDA.atomic_max!(ptr, x)
+        elseif op === right
+            CUDA.atomic_xchg!(ptr, x)
         else
             error("not implemented")
         end
diff --git a/ext/AtomixMetalExt.jl b/ext/AtomixMetalExt.jl
index af98afb..32a10e0 100644
--- a/ext/AtomixMetalExt.jl
+++ b/ext/AtomixMetalExt.jl
@@ -1,7 +1,7 @@
 # TODO: respect ordering
 module AtomixMetalExt
 
-using Atomix: Atomix, IndexableRef
+using Atomix: Atomix, IndexableRef, right
 using Metal: Metal, MtlDeviceArray
 
 const MtlIndexableRef{Indexable<:MtlDeviceArray} = IndexableRef{Indexable}
@@ -57,6 +57,8 @@ end
             Metal.atomic_fetch_min_explicit(ptr, x)
         elseif op === max
             Metal.atomic_fetch_max_explicit(ptr, x)
+        elseif op === right
+            Metal.atomic_exchange_explicit(ptr, x)
         else
             error("not implemented")
         end
diff --git a/ext/AtomixOpenCLExt.jl b/ext/AtomixOpenCLExt.jl
index a536ca8..6b68e51 100644
--- a/ext/AtomixOpenCLExt.jl
+++ b/ext/AtomixOpenCLExt.jl
@@ -1,7 +1,7 @@
 # TODO: respect ordering
 module AtomixOpenCLExt
 
-using Atomix: Atomix, IndexableRef
+using Atomix: Atomix, IndexableRef, right
 using OpenCL: SPIRVIntrinsics, CLDeviceArray
 
 const CLIndexableRef{Indexable<:CLDeviceArray} = IndexableRef{Indexable}
@@ -48,6 +48,8 @@ end
             SPIRVIntrinsics.atomic_min!(ptr, x)
         elseif op === max
             SPIRVIntrinsics.atomic_max!(ptr, x)
+        elseif op === right
+            SPIRVIntrinsics.atomic_xchg!(ptr, x)
         else
             error("not implemented")
         end
diff --git a/ext/AtomixoneAPIExt.jl b/ext/AtomixoneAPIExt.jl
index 1866a69..445471a 100644
--- a/ext/AtomixoneAPIExt.jl
+++ b/ext/AtomixoneAPIExt.jl
@@ -1,7 +1,7 @@
 # TODO: respect ordering
 module AtomixoneAPIExt
 
-using Atomix: Atomix, IndexableRef
+using Atomix: Atomix, IndexableRef, right
 using oneAPI: oneAPI, oneDeviceArray
 
 const oneIndexableRef{Indexable<:oneDeviceArray} = IndexableRef{Indexable}
@@ -48,6 +48,8 @@ end
             oneAPI.atomic_min!(ptr, x)
         elseif op === max
             oneAPI.atomic_max!(ptr, x)
+        elseif op === right
+            oneAPI.atomic_xchg!(ptr, x)
         else
             error("not implemented")
         end
diff --git a/test/test_atomix_cuda.jl b/test/test_atomix_cuda.jl
index 487eae2..c789f4a 100644
--- a/test/test_atomix_cuda.jl
+++ b/test/test_atomix_cuda.jl
@@ -79,3 +79,15 @@ end
     end
     @test collect(A) == [2, 1, 1]
 end
+
+@testset "AtomixCUDAExt:test_swap_sugar" begin
+    A = CUDA.ones(Int, 3)
+    B = CUDA.zeros(Int, 3)
+    cuda() do
+        GC.@preserve A B begin
+            B[begin] = @atomicswap A[begin] = 4
+        end
+    end
+    @test collect(A) == [4, 1, 1]
+    @test collect(B) == [1, 0, 0]
+end
diff --git a/test/test_atomix_metal.jl b/test/test_atomix_metal.jl
index 051c9e0..441e8c0 100644
--- a/test/test_atomix_metal.jl
+++ b/test/test_atomix_metal.jl
@@ -95,3 +95,15 @@ end
     end
     @test collect(A) == [2, 1, 1]
 end
+
+@testset "AtomixMetalExt:test_swap_sugar" begin
+    A = Metal.ones(Int32, 3)
+    B = Metal.zeros(Int32, 3)
+    metal() do
+        GC.@preserve A B begin
+            B[begin] = @atomicswap A[begin] = 4
+        end
+    end
+    @test collect(A) == [4, 1, 1]
+    @test collect(B) == [1, 0, 0]
+end
diff --git a/test/test_atomix_oneapi.jl b/test/test_atomix_oneapi.jl
index ca6dfba..0f4d8ae 100644
--- a/test/test_atomix_oneapi.jl
+++ b/test/test_atomix_oneapi.jl
@@ -79,3 +79,15 @@ end
     end
     @test collect(A) == [2, 1, 1]
 end
+
+@testset "AtomixoneAPIExt:test_swap_sugar" begin
+    A = oneAPI.ones(Int32, 3)
+    B = oneAPI.zeros(Int32, 3)
+    oneapi() do
+        GC.@preserve A B begin
+            B[begin] = @atomicswap A[begin] = 4
+        end
+    end
+    @test collect(A) == [4, 1, 1]
+    @test collect(B) == [1, 0, 0]
+end
diff --git a/test/test_atomix_opencl.jl b/test/test_atomix_opencl.jl
index df63338..fa5be64 100644
--- a/test/test_atomix_opencl.jl
+++ b/test/test_atomix_opencl.jl
@@ -79,3 +79,15 @@ end
     end
     @test collect(A) == [2, 1, 1]
 end
+
+@testset "AtomixOpenCLExt:test_swap_sugar" begin
+    A = OpenCL.ones(Int32, 3)
+    B = OpenCL.zeros(Int32, 3)
+    opencl() do
+        GC.@preserve A B begin
+            B[begin] = @atomicswap A[begin] = 4
+        end
+    end
+    @test collect(A) == [4, 1, 1]
+    @test collect(B) == [1, 0, 0]
+end