From e6d4daf194f603870a0fa173c13a77331a905f82 Mon Sep 17 00:00:00 2001
From: Peter Boyle <paboyle@ph.ed.ac.uk>
Date: Fri, 2 Oct 2026 12:00:12 -0400
Subject: [PATCH] Split operators

---
 Grid/algorithms/Algorithms.h                  |   1 +
 Grid/algorithms/LinearOperator.h              | 110 ++++++++++++++++++
 Grid/algorithms/SparseMatrix.h                |  11 ++
 .../ConjugateGradientMixedPrecBatched.h       | 110 +++++++++++++++++-
 4 files changed, 226 insertions(+), 6 deletions(-)

diff --git a/Grid/algorithms/Algorithms.h b/Grid/algorithms/Algorithms.h
index 6ebb1c58d..fb26692e3 100644
--- a/Grid/algorithms/Algorithms.h
+++ b/Grid/algorithms/Algorithms.h
@@ -33,6 +33,7 @@ NAMESPACE_CHECK(blas);
 #include <Grid/algorithms/blas/BatchedBlas.h>
 
 NAMESPACE_CHECK(algorithms);
+#include <Grid/algorithms/SplitOperator.h>
 #include <Grid/algorithms/SparseMatrix.h>
 #include <Grid/algorithms/LinearOperator.h>
 #include <Grid/algorithms/Preconditioner.h>
diff --git a/Grid/algorithms/LinearOperator.h b/Grid/algorithms/LinearOperator.h
index 986a46d02..e0eacc420 100644
--- a/Grid/algorithms/LinearOperator.h
+++ b/Grid/algorithms/LinearOperator.h
@@ -52,9 +52,31 @@ public:
   virtual void AdjOp  (const Field &in, Field &out) = 0; // Abstract base
   virtual void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2)=0;
   virtual void HermOp(const Field &in, Field &out)=0;
+
+  // As CheckerBoardedSparseMatrixBase::SplitClone; the returned bundle's Linop is the
+  // same kind of wrapper as this one. Collective. nullptr if not supported.
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    return nullptr;
+  }
+
   virtual ~LinearOperatorBase(){};
 };
 
+/////////////////////////////////////////////////////////////////////////////////////////////
+// The wrappers below are templates on an arbitrary Matrix, and their virtual SplitClone is
+// instantiated with the class, so it must compile for Matrix types with no SplitClone.
+/////////////////////////////////////////////////////////////////////////////////////////////
+template<class Field,class Matrix>
+SplitOperator<Field> *SplitCloneMatrix(Matrix &Mat,const Coordinate &mpi_split)
+{
+  if constexpr ( std::is_base_of<CheckerBoardedSparseMatrixBase<Field>,Matrix>::value ) {
+    return Mat.SplitClone(mpi_split);
+  } else {
+    return nullptr;
+  }
+}
+
 
 /////////////////////////////////////////////////////////////////////////////////////////////
 // By sharing the class for Sparse Matrix across multiple operator wrappers, we can share code
@@ -102,6 +124,17 @@ public:
   void HermOp(const Field &in, Field &out){
     _Mat.MdagM(in,out);
   }
+
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+    if ( split == nullptr ) {
+      return nullptr;
+    }
+    split->Linop     = new MdagMLinearOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+    split->FieldGrid = split->Matrix->Grid();
+    return split;
+  }
 };
 template<class Matrix,class Field>
 class MMdagLinearOperator : public LinearOperatorBase<Field> {
@@ -401,6 +434,17 @@ template<class Matrix,class Field>
       _Mat.MooeeDag(in,out);
       axpy(out,-1.0,tmp,out);
     }
+
+    virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+    {
+      SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+      if ( split == nullptr ) {
+        return nullptr;
+      }
+      split->Linop     = new SchurDiagMooeeOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+      split->FieldGrid = split->Matrix->RedBlackGrid();
+      return split;
+    }
 };
 template<class Matrix,class Field>
   class SchurDiagOneOperator :  public SchurOperatorBase<Field> {
@@ -427,6 +471,17 @@ template<class Matrix,class Field>
       _Mat.MeooeDag(out,tmp);
       axpy(out,-1.0,tmp,in);
     }
+
+    virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+    {
+      SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+      if ( split == nullptr ) {
+        return nullptr;
+      }
+      split->Linop     = new SchurDiagOneOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+      split->FieldGrid = split->Matrix->RedBlackGrid();
+      return split;
+    }
 };
 template<class Matrix,class Field>
   class SchurDiagTwoOperator :  public SchurOperatorBase<Field> {
@@ -455,6 +510,17 @@ template<class Matrix,class Field>
 
       axpy(out,-1.0,tmp,in);
     }
+
+    virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+    {
+      SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+      if ( split == nullptr ) {
+        return nullptr;
+      }
+      split->Linop     = new SchurDiagTwoOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+      split->FieldGrid = split->Matrix->RedBlackGrid();
+      return split;
+    }
 };
 
 template<class Field>
@@ -522,6 +588,17 @@ class NonHermitianSchurDiagMooeeOperator :  public NonHermitianSchurOperatorBase
     
     axpy(out, -1.0, tmp, out);
   }
+
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+    if ( split == nullptr ) {
+      return nullptr;
+    }
+    split->Linop     = new NonHermitianSchurDiagMooeeOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+    split->FieldGrid = split->Matrix->RedBlackGrid();
+    return split;
+  }
 };
     
 template<class Matrix,class Field>
@@ -552,6 +629,17 @@ class NonHermitianSchurDiagOneOperator : public NonHermitianSchurOperatorBase<Fi
     
     axpy(out, -1.0, tmp, in);
   }
+
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+    if ( split == nullptr ) {
+      return nullptr;
+    }
+    split->Linop     = new NonHermitianSchurDiagOneOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+    split->FieldGrid = split->Matrix->RedBlackGrid();
+    return split;
+  }
 };
 
 template<class Matrix, class Field>
@@ -583,6 +671,17 @@ class NonHermitianSchurDiagTwoOperator : public NonHermitianSchurOperatorBase<Fi
 
     axpy(out, -1.0, tmp, in);
   }
+
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+    if ( split == nullptr ) {
+      return nullptr;
+    }
+    split->Linop     = new NonHermitianSchurDiagTwoOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+    split->FieldGrid = split->Matrix->RedBlackGrid();
+    return split;
+  }
 };
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -636,6 +735,17 @@ class SchurStaggeredOperator :  public SchurOperatorBase<Field> {
   virtual void MpcDagMpc(const Field &in, Field &out) {
     GRID_ASSERT(0);// Never need with staggered
   }
+
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    SplitOperator<Field> *split = SplitCloneMatrix<Field>(_Mat,mpi_split);
+    if ( split == nullptr ) {
+      return nullptr;
+    }
+    split->Linop     = new SchurStaggeredOperator<CheckerBoardedSparseMatrixBase<Field>,Field>(*split->Matrix);
+    split->FieldGrid = split->Matrix->RedBlackGrid();
+    return split;
+  }
 };
 template<class Matrix,class Field> using SchurStagOperator = SchurStaggeredOperator<Matrix,Field>;
 
diff --git a/Grid/algorithms/SparseMatrix.h b/Grid/algorithms/SparseMatrix.h
index 2922b3e38..ccf750b04 100644
--- a/Grid/algorithms/SparseMatrix.h
+++ b/Grid/algorithms/SparseMatrix.h
@@ -78,6 +78,17 @@ public:
   virtual  void MeooeDag    (const Field &in, Field &out)=0;
   virtual  void MooeeDag    (const Field &in, Field &out)=0;
   virtual  void MooeeInvDag (const Field &in, Field &out)=0;
+
+  //////////////////////////////////////////////////////////////////////
+  // Copy of this operator on grids whose communicator is split into
+  // partitions of MPI layout mpi_split. Collective over the full
+  // communicator. nullptr if this operator cannot be cloned.
+  //////////////////////////////////////////////////////////////////////
+  virtual SplitOperator<Field> *SplitClone(const Coordinate &mpi_split)
+  {
+    return nullptr;
+  }
+
   virtual ~CheckerBoardedSparseMatrixBase() {};
 };
 
diff --git a/Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h b/Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h
index 32df03020..27d901e2c 100644
--- a/Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h
+++ b/Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h
@@ -51,6 +51,12 @@ public:
   LinearFunction<FieldF> *guesser;
   bool updateResidual;
   
+  // 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;
+
   MixedPrecisionConjugateGradientBatched(RealD tol, 
           Integer maxinnerit, 
           Integer maxouterit, 
@@ -61,7 +67,8 @@ public:
           bool _updateResidual=true) :
     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) { };
+    OuterLoopNormMult(100.), guesser(NULL), updateResidual(_updateResidual),
+    BatchedSplit(GridDefaultBatchedSolverSplit()), BatchedSplitNode(GridDefaultBatchedSolverSplitNode()) { };
 
   void useGuesser(LinearFunction<FieldF> &g){
     guesser = &g;
@@ -130,6 +137,20 @@ public:
     ConjugateGradient<FieldF> CG_f(inner_tol, MaxInnerIterations);
     CG_f.ErrorOnNoConverge = false;
     
+    //Optionally clone the single precision operator onto split-communicator partitions
+    GridStopWatch SplitTimer;
+    SplitOperator<FieldF> *split = nullptr;
+    int partitions;
+    Coordinate layout = BatchedSolverSplitLayout(SinglePrecGrid,BatchedSplit,BatchedSplitNode,NBatch,partitions);
+    if ( partitions > 1 ) {
+      SplitTimer.Start();
+      split = Linop_f.SplitClone(layout);
+      SplitTimer.Stop();
+      if ( split == nullptr ) {
+        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: operator cannot be split; serial inner solves" << std::endl;
+      }
+    }
+    
     Integer &outer_iter = TotalOuterIterations; //so it will be equal to the final iteration count
       
     for(outer_iter = 0; outer_iter < MaxOuterIterations; outer_iter++){
@@ -169,12 +190,18 @@ public:
         (*guesser)(src_f, sol_f);
       }
 
+      if ( split != nullptr ) {
+        InnerSplitSolves(*split, CG_f, src_f, sol_f, TotalInnerIterations, InnerCGtimer, SplitTimer);
+      }
+
       for (int i=0; i<NBatch; i++) {
         //Inner CG
-        InnerCGtimer.Start();
-        CG_f(Linop_f, src_f[i], sol_f[i]);
-        InnerCGtimer.Stop();
-        TotalInnerIterations[i] += CG_f.IterationsToComplete;
+        if ( split == nullptr ) {
+          InnerCGtimer.Start();
+          CG_f(Linop_f, src_f[i], sol_f[i]);
+          InnerCGtimer.Stop();
+          TotalInnerIterations[i] += CG_f.IterationsToComplete;
+        }
         
         //Convert sol back to double and add to double prec solution
         PrecChangeTimer.Start();
@@ -185,6 +212,8 @@ public:
       }
 
     }
+
+    delete split;
     
     //Final trial CG
     std::cout << GridLogMessage << std::endl;
@@ -203,9 +232,78 @@ public:
       std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: solve " << i << " Inner CG iterations " << TotalInnerIterations[i] << " Restarts " << TotalOuterIterations << " Final CG iterations " << TotalFinalStepIterations[i] << std::endl;
     }
     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: Total time " << TotalTimer.Elapsed() << " Precision change " << PrecChangeTimer.Elapsed() << " Inner CG total " << InnerCGtimer.Elapsed() << " Split setup and transfer " << SplitTimer.Elapsed() << std::endl;
     
   }
+
+private:
+
+  ////////////////////////////////////////////////////////////////////////////////////////
+  // One inner solve per partition, in groups of Partitions right-hand sides. The last
+  // group is zero-padded; a zero source returns at once from CG. Collective.
+  ////////////////////////////////////////////////////////////////////////////////////////
+  void InnerSplitSolves(SplitOperator<FieldF> &split,
+                        ConjugateGradient<FieldF> &CG_f,
+                        std::vector<FieldF> &src_f,
+                        std::vector<FieldF> &sol_f,
+                        std::vector<Integer> &TotalInnerIterations,
+                        GridStopWatch &InnerCGtimer,
+                        GridStopWatch &SplitTimer)
+  {
+    int NBatch = src_f.size();
+    int P      = split.Partitions;
+    int cb     = src_f[0].Checkerboard();
+
+    FieldF s_src(split.FieldGrid);
+    FieldF s_sol(split.FieldGrid);
+    std::vector<FieldF> group_src(P,SinglePrecGrid);
+    std::vector<FieldF> group_sol(P,SinglePrecGrid);
+    std::vector<uint64_t> iters(P);
+
+    for(int g=0;g<NBatch;g+=P){
+
+      // Gather the group; zero-pad past the end of the batch
+      for(int p=0;p<P;p++){
+        group_src[p].Checkerboard() = cb;
+        group_sol[p].Checkerboard() = cb;
+        if ( g+p < NBatch ) {
+          group_src[p] = src_f[g+p];
+          group_sol[p] = sol_f[g+p];
+        } else {
+          group_src[p] = Zero();
+          group_sol[p] = Zero();
+        }
+      }
+
+      // The initial guess (e.g. from the guesser) travels with the source
+      SplitTimer.Start();
+      Grid_split(group_src,s_src);
+      Grid_split(group_sol,s_sol);
+      SplitTimer.Stop();
+
+      InnerCGtimer.Start();
+      CG_f(*split.Linop,s_src,s_sol);
+      InnerCGtimer.Stop();
+
+      SplitTimer.Start();
+      Grid_unsplit(group_sol,s_sol);
+      SplitTimer.Stop();
+
+      // One iteration count per partition, contributed by the partition's rank 0 only
+      for(int p=0;p<P;p++){
+        iters[p] = 0;
+      }
+      if ( split.FieldGrid->ThisRank() == 0 ) {
+        iters[split.Partition] = CG_f.IterationsToComplete;
+      }
+      SinglePrecGrid->GlobalSumVector(&iters[0],P);
+
+      for(int p=0;p<P && g+p<NBatch;p++){
+        sol_f[g+p] = group_sol[p];
+        TotalInnerIterations[g+p] += iters[p];
+      }
+    }
+  }
 };
 
 NAMESPACE_END(Grid);