AD

Author: cometscome

README.md

@@ -227,6 +227,10 @@ mul!(M1, M2s, M3)                # M1 = (M2 shifted) * M3
 
 The shift is applied with periodic wrapping across the global lattice size.
 
+---
+
+
+
 ### 3) Multiplication with conjugate-transposed matrices
 
 ```julia
@@ -239,6 +243,27 @@ All combinations of shifted and adjoint operands are supported and tested in `te
 
 ---
 
+## Automatic differentiation (Enzyme)
+
+We provide Enzyme-based AD extensions and test cases. See `test/adtest/ad.jl` for a concrete comparison between
+automatic differentiation and numerical differentiation using `calc_action_loopfn`. The loop body is factored
+into a small helper function (`_calc_action_step!`), which makes Enzyme AD more reliable for loop-heavy code.
+
+Example (runs the AD vs numerical comparison with `calc_action_loopfn`):
+
+```julia
+using Enzyme
+using LatticeMatrices, MPI, JACC
+JACC.@init_backend
+MPI.Init()
+
+include("test/adtest/ad.jl") # runs main() in the script
+```
+
+Note: the AD result here follows Enzyme's complex differentiation convention. For a complex variable
+`U = X + iY`, the gradient reported by Enzyme is
+`dS/dUij = dS/dXij + i dS/dYij`.
+
 
 ---
 

ext/AD/AD.jl

@@ -43,17 +43,17 @@ function ER.augmented_primal(cfg::ER.RevConfig,
     mul_AshiftB!(C.val, A.val, B.val, shift.val)
 
     # Always tape A and parent(B) primals to survive workspace reuse.
-    # tapeA_obj, itA = get_block(A.val.temps)
-    # tapeA_obj .= A.val.A
-    # tapeA = (tapeA_obj, itA)
-    tapeA_obj = deepcopy(A.val.A)
-    tapeA = (tapeA_obj, nothing)
-
-    # tapeB_obj, itB = get_block(B.val.temps)
-    # tapeB_obj .= B.val.A
-    # tapeB = (tapeB_obj, itB)
-    tapeB_obj = deepcopy(B.val.A)
-    tapeB = (tapeB_obj, nothing)
+    tapeA_obj, itA = get_block(A.val.temps)
+    tapeA_obj .= A.val.A
+    tapeA = (tapeA_obj, itA)
+    #tapeA_obj = deepcopy(A.val.A)
+    #tapeA = (tapeA_obj, nothing)
+
+    tapeB_obj, itB = get_block(B.val.temps)
+    tapeB_obj .= B.val.A
+    tapeB = (tapeB_obj, itB)
+    #tapeB_obj = deepcopy(B.val.A)
+    #tapeB = (tapeB_obj, nothing)
 
     tape_shift = shift.val
     if get(ENV, "LM_DEBUG_MULASHIFTB", "") == "1"
@@ -188,10 +188,10 @@ function _rev_mul_AshiftB!(
 
     # Release tape blocks（早期 return があっても必ずやりたいなら try/finally 化推奨）
     if tapeA !== nothing
-        # unused!(A.val.temps, tapeA[2])
+        unused!(A.val.temps, tapeA[2])
     end
     if tapeB !== nothing
-        # unused!(B.val.temps, tapeB[2])
+        unused!(B.val.temps, tapeB[2])
     end
 
     if _should_zero_dC(dCout)