From e9fa6a342486fd0455d76e27758db3889fe692ed Mon Sep 17 00:00:00 2001
From: James Osborn <osborn@clswl142.cels.anl.gov>
Date: Tue, 9 Jul 2024 13:34:12 -0500
Subject: [PATCH] add force solver statistics

---
 src/experimental/stagag.nim | 33 +++++++++++++++++++++++++--------
 src/hmc/agradOps.nim        | 11 +++++++----
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/experimental/stagag.nim b/src/experimental/stagag.nim
index 75ae886..e4a18cf 100644
--- a/src/experimental/stagag.nim
+++ b/src/experimental/stagag.nim
@@ -168,12 +168,19 @@ var spa = initSolverParams()
 spa.r2req = arsq
 spa.maxits = 10000
 #spa.backend = sbQex
-var spf = initSolverParams()
-#spf.subsetName = "even"
-spf.r2req = frsq
-spf.maxits = 10000
-spf.verbosity = 0
-#spf.backend = sbQex
+var spf = newSeq[type spa](hmasses.len+1)  # fermion force forward
+var spfb = newSeq[type spa](hmasses.len+1)  # fermion force backward
+for i in 0..<spf.len:
+  spf[i] = initSolverParams()
+  spf[i].r2req = frsq
+  spf[i].maxits = 10000
+  spf[i].verbosity = 0
+  #spf.backend = sbQex
+  spfb[i] = initSolverParams()
+  spfb[i].r2req = frsq
+  spfb[i].maxits = 10000
+  spfb[i].verbosity = 0
+  #spf.backend = sbQex
 
 proc norm2*(x: seq): float =
   for i in 0..<x.len:
@@ -343,9 +350,9 @@ proc addFf(g: GaugeV, i = 0) =
   pushCvec()
   let cv = cvecvs[^1]
   if i == 0:
-    stag.agradSolve(g, cv, phiv[i], mass, spf)
+    stag.agradSolve(g, cv, phiv[i], mass, spf[i], spfb[i])
   else:
-    stag.agradSolve(g, cv, phiv[i], vhmasses[i-1], spf)
+    stag.agradSolve(g, cv, phiv[i], vhmasses[i-1], spf[i], spfb[i])
   pushMom()
   stag.agradStagDeriv(momvs[^1], cv)
   pushMom()
@@ -1968,6 +1975,9 @@ alwaysAccept = false
 #gutime = 0.0
 #gftime = 0.0
 #fftime = 0.0
+for i in 0..<spf.len:
+  spf[i].resetStats
+  spfb[i].resetStats
 block:
   tic()
   for n in 1..ntrain:
@@ -1977,6 +1987,8 @@ block:
     m.update
     getGrad(m)
     let tup = getElapsedTime()
+    for i in 0..<spf.len:
+      echo &"FFits{i}: ", spf[i].getAveStats
     measure()
     if upit > 0:
       if n mod upit == 0:
@@ -1992,6 +2004,9 @@ block:
   #echo &"gu: {gutime}  gf: {gftime}  ff: {fftime}  ot: {et-at}  tt: {et}"
 
 resetMeasure()
+for i in 0..<spf.len:
+  spf[i].resetStats
+  spfb[i].resetStats
 if trajs > 0:
   m.clearStats
   pacc.clear
@@ -2006,6 +2021,8 @@ if trajs > 0:
     #echo "cost: ", nff/(vtau.obj*vtau.obj*m.avgPAccept)
     echo "cost: ", getCost0(m)
     let tup = getElapsedTime()
+    for i in 0..<spf.len:
+      echo &"FFits{i}: ", spf[i].getAveStats
     measure()
     let ttot = getElapsedTime()
     echo "End inference update: ", tup, "  measure: ", ttot-tup, "  total: ", ttot
diff --git a/src/hmc/agradOps.nim b/src/hmc/agradOps.nim
index 4514fb9..3dc6843 100644
--- a/src/hmc/agradOps.nim
+++ b/src/hmc/agradOps.nim
@@ -558,7 +558,7 @@ proc agradSolvebck[I,O](op: AgOp[I,O]) {.nimcall.} =
   let g = op.inputs[1]
   let x = op.inputs[2]
   let m = op.inputs[3]
-  let p = op.inputs[4]
+  let p = op.inputs[5]
   let r = op.outputs
   var c = r.grad.newOneOf
   r.grad.odd *= -1
@@ -583,12 +583,15 @@ proc agradSolvebck[I,O](op: AgOp[I,O]) {.nimcall.} =
   when m is AgVar:
     if m.doGrad:
       m.grad += redot(r.obj, c)
-proc agradSolve(c: var AgTape, s,g,r,x,m,p: auto) =
-  var op = newAgOp((s,g,x,m,p), r, agradSolvefwd, agradSolvebck)
+proc agradSolve(c: var AgTape, s,g,r,x,m,pf,pb: auto) =
+  var op = newAgOp((s,g,x,m,pf,pb), r, agradSolvefwd, agradSolvebck)
   c.add op
+template agradSolve*(s: Staggered, g,r,x,m,pf,pb: auto) =
+  ## g: gauge, r: result, x: src, m: mass, p: solve params
+  r.ctx.agradSolve(s, g, r, x, m, addr pf, addr pb)
 template agradSolve*(s: Staggered, g,r,x,m,p: auto) =
   ## g: gauge, r: result, x: src, m: mass, p: solve params
-  r.ctx.agradSolve(s, g, r, x, m, addr p)
+  r.ctx.agradSolve(s, g, r, x, m, addr p, addr p)
 
 when isMainModule:
   import qex, physics/stagSolve