improve CG preconditioning

jcosborn · Mar 26, 2024 · b9788f3 · b9788f3
1 parent cfa2c8c
commit b9788f3
Show file tree

Hide file tree

Showing 6 changed files with 270 additions and 89 deletions.
diff --git a/src/physics/color.nim b/src/physics/color.nim
@@ -167,6 +167,8 @@ template iadd*(r: var Color, x: AsComplex) =
   iadd(r[], x)
 template iadd*(r: var Color, x: Color2) =
   iadd(r[], x[])
+template isub*(r: var Color, x: SomeNumber) =
+  isub(r[], x)
 template isub*(r: var Color, x: Color2) =
   isub(r[], x[])
 template imul*(r: var Color, x: SomeNumber) =

diff --git a/src/physics/stagD.nim b/src/physics/stagD.nim
@@ -582,6 +582,7 @@ proc eoReduce*(s:Staggered; r,b:Field; m:SomeNumber) =
 proc eoReconstruct*(s:Staggered; r,b:Field; m:SomeNumber) =
   # r.odd = (b.odd - Doe r.even)/m
   stagD(s.so, r, s.g, r, 0.0, -1.0/m)
+  threadBarrier()
   r.odd += b/m
 
 # (d/dg) redot[ (2D)*x, c ]

diff --git a/src/physics/stagSolve.nim b/src/physics/stagSolve.nim
@@ -66,7 +66,7 @@ proc solveXX*(s: Staggered; r,x: Field; m: SomeNumber; sp0: var SolverParams;
       #var oap = (apply: op, applyPrecon: oppre)
       #cg.solve(oap, sp)
     #else:
-    var oa = (apply: op)
+    var oa = (apply: op, precon: cpNone)
     cg.solve(oa, sp)
     toc("cg.solve")
     sp.calls = 1
@@ -81,6 +81,7 @@ proc solveXX*(s: Staggered; r,x: Field; m: SomeNumber; sp0: var SolverParams;
   of sbQuda:
     tic()
     if parEven:
+      #echo x.even.norm2, " ", sp.r2req
       s.qudaSolveEE(r,x,m,sp)
       toc("qudaSolveEE")
     else:
@@ -178,6 +179,7 @@ proc solveReconL(s:Staggered; x,b:Field; m:SomeNumber; sp: var SolverParams;
   toc("setup")
   s.solveEE(x, d, m, sp)
   toc("solveEE")
+  #echo "solveReconL ", d2e, "  ", sp.r2req
   threads:
     x.even *= 4
     threadBarrier()
@@ -247,6 +249,7 @@ proc solve*(s:Staggered; x,b:Field; m:SomeNumber; sp0: var SolverParams) =
       s.D(r, x, m)
       threadBarrier()
       r := b - r
+      threadBarrier()
       let
         r2et = r.even.norm2
         r2ot = r.odd.norm2
@@ -255,7 +258,7 @@ proc solve*(s:Staggered; x,b:Field; m:SomeNumber; sp0: var SolverParams) =
         r2o = r2ot
     r2 = r2e + r2o
     if sp.verbosity>0:
-      echo "stagSolve r2: ", r2/b2
+      echo "stagSolve r2/b2: ", r2/b2
 
   sp.r2.init r2/b2
   sp.calls = 1