Run ExpandAtomicModify when lowering for the JIT

Project.toml

@@ -45,7 +45,7 @@ EnzymeCore = "0.8.16"
 Enzyme_jll = "0.0.274"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"
-LLVM = "9.8.1"
+LLVM = "9.10"
 LogExpFunctions = "0.3, 1"
 ObjectFile = "0.4, 0.5"
 PrecompileTools = "1"

src/compiler/optimize.jl

@@ -362,6 +362,10 @@ function addJuliaLegalizationPasses!(mpm::LLVM.NewPMPassManager, lower_intrinsic
             if VERSION >= v"1.11.0-DEV.208"
                 add!(fpm, FinalLowerGCPass())
             end
+            if VERSION >= v"1.13.0-DEV.321"
+                # after LateLowerGCPass so that all IPO is valid
+                add!(fpm, ExpandAtomicModifyPass())
+            end
         end
         if VERSION < v"1.11.0-DEV.208"
             add!(mpm, FinalLowerGCPass())

test/locks.jl

@@ -13,3 +13,41 @@ end
     Enzyme.autodiff(Forward, my_lock, Const)
 end
 
+# https://github.com/EnzymeAD/Enzyme.jl/issues/3086
+mutable struct AtomicCounter
+    @atomic count::Int
+end
+
+const my_counter = AtomicCounter(0)
+
+function my_atomic_modify(x)
+    @atomic my_counter.count += 1
+    return x * x
+end
+
+@testset "Atomic modify" begin
+    @test Enzyme.autodiff(Reverse, my_atomic_modify, Active, Active(2.0))[1][1] ≈ 4.0
+    @test Enzyme.autodiff(Forward, my_atomic_modify, Duplicated(3.0, 1.0))[1] ≈ 6.0
+    @test (@atomic my_counter.count) == 2
+end
+
+mutable struct DupAtomicCounter
+    @atomic count::Int
+    x::Float64
+end
+
+function my_dup_atomic_modify(c)
+    @atomic c.count += 1
+    return c.x * c.x
+end
+
+@testset "Atomic modify duplicated" begin
+    c = DupAtomicCounter(0, 2.0)
+    dc = DupAtomicCounter(0, 0.0)
+    Enzyme.autodiff(Reverse, my_dup_atomic_modify, Active, Duplicated(c, dc))
+    @test dc.x ≈ 4.0
+    @test (@atomic c.count) == 1
+    # the modification is replicated on the shadow
+    @test (@atomic dc.count) == 1
+end
+