Make any split operators have lifetime the lifetime of the BatchedCG object

Grid/algorithms/SplitOperator.h

@@ -104,13 +104,11 @@ inline int GridSplitVectorIndex(GridBase *full,GridBase *split)
 // line (--batched-solver-split). Uses the trailing dimensions of grid's processor and shm
 // layouts, so it works for 4d and 5d grids. Returns grid's own processor layout (one
 // partition) when no split is requested. Asserts divisibility; warns when partitions
-// straddle nodes, when node boundaries are not visible, or when there are more partitions
-// than right-hand sides.
+// straddle nodes or when node boundaries are not visible.
 /////////////////////////////////////////////////////////////////////////////////////////////
 inline Coordinate BatchedSolverSplitLayout(GridBase *grid,
                                            const Coordinate &request,
                                            bool node,
-                                           int nbatch,
                                            int &partitions)
 {
   int nd  = GridDefaultMpi().size();
@@ -168,10 +166,6 @@ inline Coordinate BatchedSolverSplitLayout(GridBase *grid,
     std::cout << GridLogWarning << "BatchedSolverSplit: partitions " << split
               << " straddle node boundaries (node layout " << shm << "); inner solves will communicate off node" << std::endl;
   }
-  if ( partitions > nbatch ) {
-    std::cout << GridLogWarning << "BatchedSolverSplit: " << partitions << " partitions for "
-              << nbatch << " right-hand sides; the extra partitions only solve zero padding" << std::endl;
-  }
   return split;
 }
 

Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h

@@ -53,9 +53,17 @@ public:
   
   // Inner solves on independent partitions of the communicator (--batched-solver-split).
   // BatchedSplit is the partition MPI layout, empty for no split; BatchedSplitNode selects
-  // one partition per node. Default from the command line; callers may override.
-  Coordinate BatchedSplit;
-  bool       BatchedSplitNode;
+  // one partition per node. Fixed at construction, default from the command line.
+  const Coordinate BatchedSplit;
+  const bool       BatchedSplitNode;
+
+  // Copy of Linop_f on the partitions, made once at construction and owned; nullptr when
+  // there is no split or the operator cannot be split. It holds the gauge field as it was at
+  // construction: a solver must not outlive a change of the operator's gauge field. The
+  // outer defect correction uses Linop_d, so a stale copy would cost convergence, not
+  // correctness.
+  SplitOperator<FieldF> *split = nullptr;
+  int                    Partitions = 1;
 
   MixedPrecisionConjugateGradientBatched(RealD tol, 
           Integer maxinnerit, 
@@ -64,11 +72,37 @@ public:
           GridBase* _sp_grid, 
           LinearOperatorBase<FieldF> &_Linop_f, 
           LinearOperatorBase<FieldD> &_Linop_d,
-          bool _updateResidual=true) :
+          bool _updateResidual=true,
+          const Coordinate &_split=GridDefaultBatchedSolverSplit(),
+          bool _split_node=GridDefaultBatchedSolverSplitNode()) :
     Linop_f(_Linop_f), Linop_d(_Linop_d),
     Tolerance(tol), InnerTolerance(tol), MaxInnerIterations(maxinnerit), MaxOuterIterations(maxouterit), MaxPatchupIterations(maxpatchit), SinglePrecGrid(_sp_grid),
     OuterLoopNormMult(100.), guesser(NULL), updateResidual(_updateResidual),
-    BatchedSplit(GridDefaultBatchedSolverSplit()), BatchedSplitNode(GridDefaultBatchedSolverSplitNode()) { };
+    BatchedSplit(_split), BatchedSplitNode(_split_node)
+  {
+    Coordinate layout = BatchedSolverSplitLayout(SinglePrecGrid,BatchedSplit,BatchedSplitNode,Partitions);
+    if ( Partitions > 1 ) {
+      double t0 = usecond();
+      split = Linop_f.SplitClone(layout);
+      double t1 = usecond();
+      if ( split == nullptr ) {
+        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: operator cannot be split; serial inner solves" << std::endl;
+      } else {
+        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: split clone " << (t1-t0)/1.0e6 << " s" << std::endl;
+        HostMemoryReport(SinglePrecGrid,GridLogMessage,"MixedPrecisionConjugateGradientBatched: after split clone");
+      }
+    }
+  };
+
+  // Owns split
+  MixedPrecisionConjugateGradientBatched(const MixedPrecisionConjugateGradientBatched &) = delete;
+
+  MixedPrecisionConjugateGradientBatched &operator=(const MixedPrecisionConjugateGradientBatched &) = delete;
+
+  ~MixedPrecisionConjugateGradientBatched(void)
+  {
+    delete split;
+  }
 
   void useGuesser(LinearFunction<FieldF> &g){
     guesser = &g;
@@ -139,21 +173,10 @@ public:
     ConjugateGradient<FieldF> CG_f(inner_tol, MaxInnerIterations);
     CG_f.ErrorOnNoConverge = false;
     
-    //Optionally clone the single precision operator onto split-communicator partitions
-    GridStopWatch SplitCloneTimer;
-    SplitTimers   splitTimers;
-    SplitOperator<FieldF> *split = nullptr;
-    int partitions;
-    Coordinate layout = BatchedSolverSplitLayout(SinglePrecGrid,BatchedSplit,BatchedSplitNode,NBatch,partitions);
-    if ( partitions > 1 ) {
-      SplitCloneTimer.Start();
-      split = Linop_f.SplitClone(layout);
-      SplitCloneTimer.Stop();
-      if ( split == nullptr ) {
-        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: operator cannot be split; serial inner solves" << std::endl;
-      } else {
-        HostMemoryReport(DoublePrecGrid,GridLogMessage,"MixedPrecisionConjugateGradientBatched: after split clone");
-      }
+    SplitTimers splitTimers;
+    if ( split != nullptr && split->Partitions > NBatch ) {
+      std::cout << GridLogWarning << "MixedPrecisionConjugateGradientBatched: " << split->Partitions << " partitions for "
+                << NBatch << " right-hand sides; the extra partitions only solve zero padding" << std::endl;
     }
     
     Integer &outer_iter = TotalOuterIterations; //so it will be equal to the final iteration count
@@ -221,10 +244,6 @@ public:
 
     }
 
-    SplitCloneTimer.Start();
-    delete split;
-    SplitCloneTimer.Stop();
-    
     //Final trial CG
     std::cout << GridLogMessage << std::endl;
     std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Starting final patch-up double-precision solve"<<std::endl;
@@ -246,16 +265,15 @@ public:
     std::cout << GridLogMessage << std::endl;
     std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Total time " << TotalTimer.Elapsed() << " Precision change " << PrecChangeTimer.Elapsed() << " Inner CG total " << InnerCGtimer.Elapsed() << std::endl;
     std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Outer residual " << OuterResidualTimer.Elapsed() << " Patch-up " << PatchupTimer.Elapsed() << std::endl;
-    if ( partitions > 1 ) {
-      std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Split clone and delete " << SplitCloneTimer.Elapsed()
-               << " Grid_split " << splitTimers.Split.Elapsed() << " (" << splitTimers.SplitCalls << " calls)"
+    if ( split != nullptr ) {
+      std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Grid_split " << splitTimers.Split.Elapsed() << " (" << splitTimers.SplitCalls << " calls)"
                << " Grid_unsplit " << splitTimers.Unsplit.Elapsed() << " (" << splitTimers.UnsplitCalls << " calls)"
                << " Staging copies " << splitTimers.Staging.Elapsed()
                << " Host to device " << splitTimers.H2D.Elapsed()
                << " Iteration count sum " << splitTimers.Reduce.Elapsed() << std::endl;
     }
     double accounted = PrecChangeTimer.useconds() + InnerCGtimer.useconds() + OuterResidualTimer.useconds()
-                     + PatchupTimer.useconds() + SplitCloneTimer.useconds() + splitTimers.Split.useconds()
+                     + PatchupTimer.useconds() + splitTimers.Split.useconds()
                      + splitTimers.Unsplit.useconds() + splitTimers.Staging.useconds() + splitTimers.H2D.useconds()
                      + splitTimers.Reduce.useconds();
     std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Unaccounted " << (TotalTimer.useconds()-accounted)/1.0e6 << " s" << std::endl;

tests/solver/Test_split_mixedprec_batched.cc

@@ -145,16 +145,19 @@ void CheckSplitBatched(const std::string &name,
     sol_split[i] = Zero();
   }
 
-  MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(TOLERANCE,10000,50,1000,grid_f,Linop_f,Linop_d);
-
-  std::cout << GridLogMessage << name << ": unsplit batched solve" << std::endl;
-  mCG.BatchedSplit     = Coordinate();
-  mCG.BatchedSplitNode = false;
-  mCG(src,sol_ref);
-
-  std::cout << GridLogMessage << name << ": split batched solve, partition layout " << layout << std::endl;
-  mCG.BatchedSplit = layout;
-  mCG(src,sol_split);
+  {
+    std::cout << GridLogMessage << name << ": unsplit batched solve" << std::endl;
+    MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(TOLERANCE,10000,50,1000,grid_f,Linop_f,Linop_d,
+                                                              true,Coordinate(),false);
+    mCG(src,sol_ref);
+  }
+  {
+    std::cout << GridLogMessage << name << ": split batched solve, partition layout " << layout << std::endl;
+    MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(TOLERANCE,10000,50,1000,grid_f,Linop_f,Linop_d,
+                                                              true,layout,false);
+    GRID_ASSERT( (mCG.Partitions == 1) || (mCG.split != nullptr) );
+    mCG(src,sol_split);
+  }
 
   FieldD Msol(src[0].Grid());
   Msol.Checkerboard() = cb;

tests/solver/Test_split_mobius_batched.cc

@@ -213,25 +213,24 @@ int main (int argc, char ** argv)
 
   ReportMemory(UGrid_d,"operators and sources ready");
 
-  MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(tol,10000,50,10000,FrbGrid_f,Linop_f,Linop_d);
-
-  Coordinate split     = mCG.BatchedSplit;
-  bool       splitnode = mCG.BatchedSplitNode;
-
   if ( unsplit ) {
-    mCG.BatchedSplit     = Coordinate();
-    mCG.BatchedSplitNode = false;
+    MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(tol,10000,50,10000,FrbGrid_f,Linop_f,Linop_d,
+                                                              true,Coordinate(),false);
     SolveAndReport("UNSPLIT",mCG,Linop_d,src,sol);
     ReportMemory(UGrid_d,"after unsplit solve");
   }
 
-  mCG.BatchedSplit     = split;
-  mCG.BatchedSplitNode = splitnode;
-  // Repeated split solves expose allocations not released between calls
-  for(int r=0;r<repeat;r++){
-    SolveAndReport("SPLIT",mCG,Linop_d,src,sol);
-    ReportMemory(UGrid_d,"after split solve "+std::to_string(r));
+  if ( repeat > 0 ) {
+    // Split layout from the command line; the split operator is cloned here, once
+    MixedPrecisionConjugateGradientBatched<FieldD,FieldF> mCG(tol,10000,50,10000,FrbGrid_f,Linop_f,Linop_d);
+    ReportMemory(UGrid_d,"after split solver construction");
+    // Repeated split solves reuse the clone and expose allocations not released between calls
+    for(int r=0;r<repeat;r++){
+      SolveAndReport("SPLIT",mCG,Linop_d,src,sol);
+      ReportMemory(UGrid_d,"after split solve "+std::to_string(r));
+    }
   }
+  ReportMemory(UGrid_d,"after split solver destruction");
 
   Grid_finalize();
 }