optimal torus placement for Tofu network (Fujitsu-specific)

Grid/GridCore.h

@@ -56,6 +56,7 @@ Author: paboyle <paboyle@ph.ed.ac.uk>
 #include <Grid/communicator/Communicator.h> 
 #include <Grid/communicator/RingAllReduce.h>
 #include <Grid/cartesian/Cartesian.h>
+#include <Grid/perfmon/HostMemory.h>
 #include <Grid/tensors/Tensors.h>      
 #include <Grid/lattice/Lattice.h>      
 #include <Grid/cshift/Cshift.h>       

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>

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>;
 

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() {};
 };
 

Grid/algorithms/SplitOperator.h

@@ -0,0 +1,172 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./lib/algorithms/SplitOperator.h
+
+    Copyright (C) 2026
+
+Author: Peter Boyle <paboyle@ph.ed.ac.uk>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+*************************************************************************************/
+/*  END LEGAL */
+#pragma once
+
+NAMESPACE_BEGIN(Grid);
+
+template<class Field> class LinearOperatorBase;
+template<class Field> class CheckerBoardedSparseMatrixBase;
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// A copy of an operator living on grids whose communicator is split into independent
+// partitions, together with those grids. Produced by SplitClone().
+//
+// Owns everything it points to. Destruction is the reverse of creation: linear operator,
+// matrix, then grids, so no object outlives a grid it references. For 4d operators the
+// fermion grids are the gauge grids and are deleted once.
+/////////////////////////////////////////////////////////////////////////////////////////////
+template<class Field>
+class SplitOperator
+{
+public:
+  GridCartesian         *GaugeGrid     = nullptr;
+  GridRedBlackCartesian *GaugeRBGrid   = nullptr;
+  GridCartesian         *FermionGrid   = nullptr;
+  GridRedBlackCartesian *FermionRBGrid = nullptr;
+
+  int Partition  = 0;  // Grid_split vector index held by this rank's partition
+  int Partitions = 1;  // number of partitions
+
+  CheckerBoardedSparseMatrixBase<Field> *Matrix    = nullptr;
+  LinearOperatorBase<Field>             *Linop     = nullptr;
+  GridBase                              *FieldGrid = nullptr;  // grid of Linop's fields
+
+  SplitOperator(void) {};
+
+  SplitOperator(const SplitOperator &) = delete;
+
+  SplitOperator &operator=(const SplitOperator &) = delete;
+
+  ~SplitOperator(void)
+  {
+    delete Linop;
+    delete Matrix;
+    if ( FermionRBGrid != GaugeRBGrid ) {
+      delete FermionRBGrid;
+    }
+    if ( FermionGrid != GaugeGrid ) {
+      delete FermionGrid;
+    }
+    delete GaugeRBGrid;
+    delete GaugeGrid;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Index of the vector that Grid_split(std::vector<Field> full, Field split) delivers to this
+// rank's partition. Grid_split orders partitions lexicographically with the first
+// dimension fastest; the split communicator's own rank (srank) uses the reversed MPI
+// convention, so the two differ whenever more than one dimension is split.
+/////////////////////////////////////////////////////////////////////////////////////////////
+inline int GridSplitVectorIndex(GridBase *full,GridBase *split)
+{
+  int nd = full->_ndimension;
+  GRID_ASSERT(split->_ndimension == nd);
+
+  Coordinate scoor(nd);
+  Coordinate ssize(nd);
+  for(int d=0;d<nd;d++){
+    scoor[d] = full->ThisProcessorCoor()[d] / split->ProcessorGrid()[d];
+    ssize[d] = full->ProcessorGrid()[d]     / split->ProcessorGrid()[d];
+  }
+  int index;
+  Lexicographic::IndexFromCoor(scoor,index,ssize);
+  return index;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Partition MPI layout for a batched solve on grid, from a request as given on the command
+// 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 or when node boundaries are not visible.
+/////////////////////////////////////////////////////////////////////////////////////////////
+inline Coordinate BatchedSolverSplitLayout(GridBase *grid,
+                                           const Coordinate &request,
+                                           bool node,
+                                           int &partitions)
+{
+  int nd  = GridDefaultMpi().size();
+  int pad = grid->_ndimension - nd;
+  GRID_ASSERT(pad >= 0);
+  GRID_ASSERT(grid->ShmGrid().size() == grid->_ndimension);
+
+  Coordinate processors(nd);
+  Coordinate shm(nd);
+  for(int d=0;d<nd;d++){
+    processors[d] = grid->ProcessorGrid()[pad+d];
+    shm[d]        = grid->ShmGrid()[pad+d];
+  }
+
+  Coordinate split(nd);
+  if ( node ) {
+    split = shm;
+  } else if ( request.size() == 0 ) {
+    split = processors;
+  } else {
+    GRID_ASSERT(request.size() == nd);
+    split = request;
+  }
+
+  partitions = 1;
+  for(int d=0;d<nd;d++){
+    GRID_ASSERT( (processors[d] % split[d]) == 0 );
+    partitions *= processors[d] / split[d];
+  }
+  if ( partitions == 1 ) {
+    return split;
+  }
+
+  std::cout << GridLogMessage << "BatchedSolverSplit: partition layout " << split
+            << " of " << processors << " : " << partitions << " partitions" << std::endl;
+
+  int inside_node  = 1;
+  int whole_nodes  = 1;
+  int shm_trivial  = 1;
+  for(int d=0;d<nd;d++){
+    if ( (shm[d] % split[d]) != 0 ) {
+      inside_node = 0;
+    }
+    if ( (split[d] % shm[d]) != 0 ) {
+      whole_nodes = 0;
+    }
+    if ( shm[d] != 1 ) {
+      shm_trivial = 0;
+    }
+  }
+  if ( shm_trivial ) {
+    std::cout << GridLogWarning << "BatchedSolverSplit: shm layout is trivial (one rank per node,"
+              << " or shared memory disabled); node locality of partitions not checked" << std::endl;
+  } else if ( !inside_node && !whole_nodes ) {
+    std::cout << GridLogWarning << "BatchedSolverSplit: partitions " << split
+              << " straddle node boundaries (node layout " << shm << "); inner solves will communicate off node" << std::endl;
+  }
+  return split;
+}
+
+NAMESPACE_END(Grid);

Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h

@@ -51,6 +51,20 @@ 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. 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, 
           Integer maxouterit, 
@@ -58,10 +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) { };
+    OuterLoopNormMult(100.), guesser(NULL), updateResidual(_updateResidual),
+    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;
@@ -91,6 +132,8 @@ public:
 
     GridStopWatch InnerCGtimer;
     GridStopWatch PrecChangeTimer;
+    GridStopWatch OuterResidualTimer;
+    GridStopWatch PatchupTimer;
     
     int cb = src_d_in[0].Checkerboard();
     
@@ -130,18 +173,27 @@ public:
     ConjugateGradient<FieldF> CG_f(inner_tol, MaxInnerIterations);
     CG_f.ErrorOnNoConverge = false;
     
+    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
       
     for(outer_iter = 0; outer_iter < MaxOuterIterations; outer_iter++){
       std::cout << GridLogMessage << std::endl;
       std::cout << GridLogMessage << "Outer iteration " << outer_iter << std::endl;
+      HostMemoryReport(DoublePrecGrid,GridLogMessage,"MixedPrecisionConjugateGradientBatched: outer iteration "+std::to_string(outer_iter));
       
       bool allConverged = true;
       
       for (int i=0; i<NBatch; i++) {
         //Compute double precision rsd and also new RHS vector.
+        OuterResidualTimer.Start();
         Linop_d.HermOp(sol_d[i], tmp_d);
         norm[i] = axpy_norm(src_d[i], -1., tmp_d, src_d_in[i]); //src_d is residual vector
+        OuterResidualTimer.Stop();
         
         std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Outer iteration " << outer_iter <<" solve " << i << " residual "<< norm[i] << " target "<< stop[i] <<std::endl;
 
@@ -169,12 +221,18 @@ public:
         (*guesser)(src_f, sol_f);
       }
 
+      if ( split != nullptr ) {
+        InnerSplitSolves(*split, CG_f, src_f, sol_f, TotalInnerIterations, InnerCGtimer, splitTimers);
+      }
+
       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,16 +243,18 @@ public:
       }
 
     }
-    
+
     //Final trial CG
     std::cout << GridLogMessage << std::endl;
     std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Starting final patch-up double-precision solve"<<std::endl;
     
+    PatchupTimer.Start();
     for (int i=0; i<NBatch; i++) {
       ConjugateGradient<FieldD> CG_d(Tolerance, MaxPatchupIterations);
       CG_d(Linop_d, src_d_in[i], sol_d[i]);
       TotalFinalStepIterations[i] += CG_d.IterationsToComplete;
     }
+    PatchupTimer.Stop();
 
     TotalTimer.Stop();
 
@@ -204,8 +264,118 @@ 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 ( 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() + 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;
     
   }
+
+private:
+
+  // Coarse breakdown of the split path outside the inner CG
+  struct SplitTimers {
+    GridStopWatch Split;
+    GridStopWatch Unsplit;
+    GridStopWatch Staging;
+    GridStopWatch H2D;
+    GridStopWatch Reduce;
+    int SplitCalls   = 0;
+    int UnsplitCalls = 0;
+  };
+
+  ////////////////////////////////////////////////////////////////////////////////////////
+  // 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,
+                        SplitTimers &timers)
+  {
+    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
+      timers.Staging.Start();
+      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();
+        }
+      }
+      timers.Staging.Stop();
+
+      // The initial guess (e.g. from the guesser) travels with the source
+      timers.Split.Start();
+      Grid_split(group_src,s_src);
+      Grid_split(group_sol,s_sol);
+      timers.Split.Stop();
+      timers.SplitCalls += 2;
+
+      // Grid_split leaves its output on the host; move it now so the copy is not in the CG time
+      timers.H2D.Start();
+      {
+        autoView(s_src_v, s_src, AcceleratorRead);
+        autoView(s_sol_v, s_sol, AcceleratorRead);
+      }
+      timers.H2D.Stop();
+
+      InnerCGtimer.Start();
+      CG_f(*split.Linop,s_src,s_sol);
+      InnerCGtimer.Stop();
+
+      timers.Unsplit.Start();
+      Grid_unsplit(group_sol,s_sol);
+      timers.Unsplit.Stop();
+      timers.UnsplitCalls += 1;
+
+      // One iteration count per partition, contributed by the partition's rank 0 only
+      timers.Reduce.Start();
+      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);
+      timers.Reduce.Stop();
+
+      // Includes the host to device copy of group_sol left by Grid_unsplit
+      timers.Staging.Start();
+      for(int p=0;p<P && g+p<NBatch;p++){
+        sol_f[g+p] = group_sol[p];
+        TotalInnerIterations[g+p] += iters[p];
+      }
+      timers.Staging.Stop();
+    }
+  }
 };
 
 NAMESPACE_END(Grid);

Grid/communicator/Communicator_base.cc

@@ -34,6 +34,9 @@ Author: Peter Boyle <paboyle@ph.ed.ac.uk>
 NAMESPACE_BEGIN(Grid);
 
 bool Stencil_force_mpi = true;
+#ifdef GRID_FJMPI
+bool Stencil_persistent = false;
+#endif
 bool Stencil_force_barrier = false;
 
 ///////////////////////////////////////////////////////////////

Grid/communicator/Communicator_base.h

@@ -38,6 +38,10 @@ Author: Peter Boyle <paboyle@ph.ed.ac.uk>
 NAMESPACE_BEGIN(Grid);
 
 extern bool Stencil_force_mpi ;
+#ifdef GRID_FJMPI
+// Fujitsu MPI persistent requests for the halo exchange with other nodes (set by --comms-overlap)
+extern bool Stencil_persistent ;
+#endif
 extern bool Stencil_force_barrier ;
 
 class CartesianCommunicator : public SharedMemory {
@@ -201,6 +205,11 @@ public:
 
   // Could do a PollHtoD and have a CommsMerge dependence
   void StencilSendToRecvFromPollDtoH (std::vector<CommsRequest_t> &list);
+#ifdef A64FX
+  // Make progress on outstanding requests without blocking (MPI only progresses
+  // inside MPI calls on Fugaku); safe to call repeatedly before Complete.
+  void StencilSendToRecvFromProgress(std::vector<CommsRequest_t> &list);
+#endif
   void StencilSendToRecvFromPollIRecv(std::vector<CommsRequest_t> &list);
 
   double StencilSendToRecvFromBegin(std::vector<CommsRequest_t> &list,

Grid/communicator/Communicator_mpi3.cc

@@ -28,6 +28,90 @@ Author: Peter Boyle <paboyle@ph.ed.ac.uk>
 #include <Grid/GridCore.h>
 #include <Grid/communicator/SharedMemory.h>
 
+// Fujitsu MPI extended persistent requests (FJMPI_Prequest_*, Fujitsu MPI User's Guide J2UL-2565 5.3).
+// Fujitsu MPI moves the body of a large non-blocking message only inside MPI calls (6.2), so the halo
+// exchange does not overlap the interior kernel. With these requests the partners exchange buffer
+// information once, at initialisation, and FJMPI_Prequest_start starts the transfer of the message body.
+// GRID_FJMPI is set by configure (--enable-fjmpi) when mpi-ext.h declares the Fujitsu extensions.
+#ifdef GRID_FJMPI
+#include <mpi-ext.h>
+
+NAMESPACE_BEGIN(Grid);
+namespace {
+// Persistent halo requests are used with --comms-overlap (Stencil_persistent); they bring nothing
+// without overlap. Only for peers on other nodes: persistent transfers to peers on the same node delay
+// the transfers to other nodes.
+bool PersistentHaloOn(void) { return Stencil_persistent; }
+bool PersistentHaloPeer(int grank) { return grank==MPI_UNDEFINED; }
+struct PersistentHaloReq {
+  void *buf; int count; MPI_Request req;
+};
+// Requests are keyed by (send, communicator, peer, tag): Fujitsu MPI allows one request per
+// (peer, tag, communicator). Several stencils can use the same communicator (e.g. even and odd
+// checkerboards), so each distinct pair of halo buffers of a direction gets a slot, in order of
+// first use, and the slot enters the tag. All ranks run the same sequence of stencils, so partners
+// agree on the slot.
+struct PersistentHaloCache {
+  std::map<std::pair<MPI_Comm,int>, std::vector<std::pair<void *,void *> > > slots;
+  std::map<std::tuple<int,MPI_Comm,int,int>, PersistentHaloReq> reqs;
+  int tag_ub = -1;
+};
+PersistentHaloCache &PersistentHalo(void) { static PersistentHaloCache c; return c; }
+
+int PersistentHaloSlot(MPI_Comm comm,int dir,void *xmit,void *recv)
+{
+  auto &v = PersistentHalo().slots[std::make_pair(comm,dir)];
+  for(int k=0;k<v.size();k++) if ( v[k].first==xmit && v[k].second==recv ) return k;
+  v.push_back(std::make_pair(xmit,recv));
+  return v.size()-1;
+}
+// Tag of a halo message: the usual tag (dir + 32 * sender rank) shifted by the slot; -1 if it
+// exceeds MPI_TAG_UB (both partners then fall back to MPI_Isend/MPI_Irecv)
+int PersistentHaloTag(int dir,int sender,int nprocs,int slot)
+{
+  auto &c = PersistentHalo();
+  if ( c.tag_ub < 0 ) {
+    int *ub, flag;
+    MPI_Comm_get_attr(MPI_COMM_WORLD,MPI_TAG_UB,&ub,&flag);
+    c.tag_ub = flag ? *ub : 32767;
+  }
+  int64_t tag = dir + 32*((int64_t)sender + (int64_t)nprocs*slot);
+  return (tag <= c.tag_ub) ? (int)tag : -1;
+}
+// The inactive persistent request of one message, created on first use and recreated if the
+// buffer or size changed (a stencil freed and another one allocated at the same address)
+MPI_Request PersistentHaloGet(int send,void *buf,int count,int peer,int tag,MPI_Comm comm)
+{
+  auto key = std::make_tuple(send,comm,peer,tag);
+  auto &reqs = PersistentHalo().reqs;
+  auto it = reqs.find(key);
+  if ( it != reqs.end() ) {
+    if ( it->second.buf==buf && it->second.count==count ) return it->second.req;
+    MPI_Request_free(&it->second.req);
+    reqs.erase(it);
+  }
+  PersistentHaloReq r = {buf,count,MPI_REQUEST_NULL};
+  int ierr = send ? FJMPI_Prequest_send_init(buf,count,MPI_INT32_T,peer,tag,comm,&r.req)
+                  : FJMPI_Prequest_recv_init(buf,count,MPI_INT32_T,peer,tag,comm,&r.req);
+  GRID_ASSERT(ierr==0);
+  reqs[key] = r;
+  return r.req;
+}
+void PersistentHaloFree(MPI_Comm comm)
+{
+  auto &c = PersistentHalo();
+  for(auto it=c.reqs.begin();it!=c.reqs.end();) {
+    if ( std::get<1>(it->first)==comm ) { MPI_Request_free(&it->second.req); it=c.reqs.erase(it); }
+    else ++it;
+  }
+  for(auto it=c.slots.begin();it!=c.slots.end();) {
+    if ( it->first.first==comm ) it=c.slots.erase(it); else ++it;
+  }
+}
+}
+NAMESPACE_END(Grid);
+#endif
+
 void GridAbort(void) { MPI_Abort(MPI_COMM_WORLD,SIGABRT); }
 extern void * Grid_backtrace_buffer[_NBACKTRACE];
 
@@ -191,6 +275,7 @@ CartesianCommunicator::CartesianCommunicator(const Coordinate &processors,const
   // Set up from the new split communicator
   //////////////////////////////////////////////////////////////////////////////////////////////////////
   InitFromMPICommunicator(processors,comm_split);
+  _shm_processors = shm_processors;
 
   //////////////////////////////////////////////////////////////////////////////////////////////////////
   // Take the right SHM buffers
@@ -261,6 +346,9 @@ CartesianCommunicator::~CartesianCommunicator()
   if (communicator && !MPI_is_finalised) {
     MPI_Comm_free(&communicator);
     for(int i=0;i<communicator_halo.size();i++){
+#ifdef GRID_FJMPI
+      PersistentHaloFree(communicator_halo[i]);
+#endif
       MPI_Comm_free(&communicator_halo[i]);
     }
   }
@@ -420,6 +508,18 @@ int CartesianCommunicator::IsOffNode(int rank)
 #ifdef ACCELERATOR_AWARE_MPI
 void CartesianCommunicator::StencilSendToRecvFromPollIRecv(std::vector<CommsRequest_t> &list) {};
 void CartesianCommunicator::StencilSendToRecvFromPollDtoH(std::vector<CommsRequest_t> &list) {};
+#ifdef A64FX
+void CartesianCommunicator::StencilSendToRecvFromProgress(std::vector<CommsRequest_t> &list)
+{
+  int nreq=list.size();
+  if (nreq==0) return;
+  int flag;
+  // completed requests become MPI_REQUEST_NULL (persistent ones become inactive), which the final
+  // Waitall handles
+  int ierr = MPI_Testall(nreq,&list[0],&flag,MPI_STATUSES_IGNORE);
+  GRID_ASSERT(ierr==0);
+}
+#endif
 double CartesianCommunicator::StencilSendToRecvFromPrepare(std::vector<CommsRequest_t> &list,
 							   void *xmit,
 							   int dest,int dox,
@@ -427,6 +527,34 @@ double CartesianCommunicator::StencilSendToRecvFromPrepare(std::vector<CommsRequ
 							   int from,int dor,
 							   uint64_t xbytes,uint64_t rbytes,int dir)
 {
+#ifdef GRID_FJMPI
+  // Persistent requests. Starting a persistent receive waits until the partner has created the
+  // matching send request, and starting a send waits until the partner has started the receive
+  // (measured; the manual documents the latter). So both requests of this packet are created here
+  // before its receive is started, and all receives are started before any send is started in
+  // StencilSendToRecvFromBegin. Partners process the packets in the same order, so no cycle forms.
+  if ( PersistentHaloOn() ) {
+    int ncomm  =communicator_halo.size();
+    int commdir=dir%ncomm;
+    int gfrom = ShmRanks[from];
+    int gdest = ShmRanks[dest];
+    MPI_Comm comm = communicator_halo[commdir];
+    int slot = PersistentHaloSlot(comm,dir,xmit,recv);
+    int rtag = PersistentHaloTag(dir,from,_Nprocessors,slot);
+    int xtag = PersistentHaloTag(dir,_processor,_Nprocessors,slot);
+    bool mpi_recv = dor && ( (gfrom ==MPI_UNDEFINED) || Stencil_force_mpi ) && rtag >= 0 && PersistentHaloPeer(gfrom);
+    bool mpi_send = dox && ( (gdest ==MPI_UNDEFINED) || Stencil_force_mpi ) && xtag >= 0 && PersistentHaloPeer(gdest);
+    MPI_Request rrq;
+    if ( mpi_recv ) rrq = PersistentHaloGet(0,recv,(int)(rbytes/sizeof(int32_t)),from,rtag,comm);
+    if ( mpi_send ) PersistentHaloGet(1,xmit,(int)(xbytes/sizeof(int32_t)),dest,xtag,comm);
+    if ( mpi_recv ) {
+      int ierr = FJMPI_Prequest_start(&rrq);
+      GRID_ASSERT(ierr==0);
+      list.push_back(rrq);
+      return rbytes;
+    }
+  }
+#endif
   return 0.0; // Do nothing -- no preparation required
 }
 double CartesianCommunicator::StencilSendToRecvFromBegin(std::vector<CommsRequest_t> &list,
@@ -453,7 +581,19 @@ double CartesianCommunicator::StencilSendToRecvFromBegin(std::vector<CommsReques
   double off_node_bytes=0.0;
   int tag;
   
+#ifdef GRID_FJMPI
+  // Persistent requests (see StencilSendToRecvFromPrepare): same slot and tags on both sides
+  int ptag_recv = -1, ptag_send = -1;
+  if ( PersistentHaloOn() ) {
+    MPI_Comm comm = communicator_halo[commdir];
+    int slot = PersistentHaloSlot(comm,dir,xmit_comp,recv_comp);
+    if ( PersistentHaloPeer(gfrom) ) ptag_recv = PersistentHaloTag(dir,from,_Nprocessors,slot);
+    if ( PersistentHaloPeer(gdest) ) ptag_send = PersistentHaloTag(dir,_processor,_Nprocessors,slot);
+  }
+  if ( dor && ptag_recv < 0 ) { // receive not started in StencilSendToRecvFromPrepare
+#else
   if ( dor ) {
+#endif
     if ( (gfrom ==MPI_UNDEFINED) || Stencil_force_mpi ) {
       tag= dir+from*32;
       //      std::cout << " StencilSendToRecvFrom "<<dir<<" MPI_Irecv "<<std::hex<<recv<<std::dec<<std::endl;
@@ -474,8 +614,18 @@ double CartesianCommunicator::StencilSendToRecvFromBegin(std::vector<CommsReques
   // This is a NVLINK PUT  
   if (dox) {
     if ( (gdest == MPI_UNDEFINED) || Stencil_force_mpi ) {
+#ifdef GRID_FJMPI
+      if ( ptag_send >= 0 ) {
+	xrq = PersistentHaloGet(1,xmit_comp,(int)(xbytes/sizeof(int32_t)),dest,ptag_send,communicator_halo[commdir]);
+	ierr = FJMPI_Prequest_start(&xrq);
+      } else {
+	tag= dir+_processor*32;
+	ierr =MPI_Isend(xmit_comp,(int)(xbytes/sizeof(int32_t)), MPI_INT32_T,dest,tag,communicator_halo[commdir],&xrq);
+      }
+#else
       tag= dir+_processor*32;
       ierr =MPI_Isend(xmit_comp,(int)(xbytes/sizeof(int32_t)), MPI_INT32_T,dest,tag,communicator_halo[commdir],&xrq);
+#endif
       GRID_ASSERT(ierr==0);
       list.push_back(xrq);
       off_node_bytes+=xbytes;
@@ -505,6 +655,9 @@ void CartesianCommunicator::StencilSendToRecvFromComplete(std::vector<CommsReque
 }
 
 #else /* NOT     ... ACCELERATOR_AWARE_MPI */
+#ifdef A64FX
+void CartesianCommunicator::StencilSendToRecvFromProgress(std::vector<CommsRequest_t> &list) {};
+#endif
 ///////////////////////////////////////////
 // Pipeline mode through host memory
 ///////////////////////////////////////////

Grid/communicator/Communicator_none.cc

@@ -167,6 +167,9 @@ double CartesianCommunicator::StencilSendToRecvFrom( void *xmit,
 }
 void CartesianCommunicator::StencilSendToRecvFromPollIRecv(std::vector<CommsRequest_t> &list) {};
 void CartesianCommunicator::StencilSendToRecvFromPollDtoH(std::vector<CommsRequest_t> &list) {};
+#ifdef A64FX
+void CartesianCommunicator::StencilSendToRecvFromProgress(std::vector<CommsRequest_t> &list) {};
+#endif
 double CartesianCommunicator::StencilSendToRecvFromPrepare(std::vector<CommsRequest_t> &list,
 							   void *xmit,
 							   int xmit_to_rank,int dox,

Grid/communicator/SharedMemory.cc

@@ -33,6 +33,7 @@ NAMESPACE_BEGIN(Grid);
 // static data
 
 int                 GlobalSharedMemory::HPEhypercube = 1;
+int                 GlobalSharedMemory::TofuPlacement = 0;
 uint64_t            GlobalSharedMemory::MAX_MPI_SHM_BYTES   = 1024LL*1024LL*1024LL; 
 int                 GlobalSharedMemory::Hugepages = 0;
 int                 GlobalSharedMemory::_ShmSetup;

Grid/communicator/SharedMemory.h

@@ -99,6 +99,10 @@ public:
   // HPE 8600 hypercube optimisation
   ///////////////////////////////////////
   static int HPEhypercube;
+  ///////////////////////////////////////
+  // Tofu (Fujitsu MPI) node placement, --tofu
+  ///////////////////////////////////////
+  static int TofuPlacement;
 
   static int      ShmSetup(void)      { return _ShmSetup; }
   static int      ShmAlloc(void)      { return _ShmAlloc; }
@@ -131,6 +135,7 @@ public:
   static void OptimalCommunicator            (const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &ShmDims); 
   static void OptimalCommunicatorHypercube   (const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &ShmDims); 
   static void OptimalCommunicatorSharedMemory(const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &ShmDims); 
+  static void OptimalCommunicatorTofu        (const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &ShmDims);
   static void GetShmDims(const Coordinate &WorldDims,Coordinate &ShmDims);
   ///////////////////////////////////////////////////
   // Provide shared memory facilities off comm world

Grid/communicator/SharedMemoryMPI.cc

@@ -32,6 +32,10 @@ Author: Christoph Lehner <christoph@lhnr.de>
 #include <Grid/GridCore.h>
 #include <pwd.h>
 
+#ifdef GRID_FJMPI
+#include <mpi-ext.h>
+#endif
+
 #ifdef GRID_CUDA
 #include <cuda_runtime_api.h>
 #endif
@@ -278,6 +282,18 @@ void GlobalSharedMemory::OptimalCommunicator(const Coordinate &processors,Grid_M
   gethostname(name,namelen);
   int nscan = sscanf(name,"r%di%dn%d",&R,&I,&N) ;
 
+  if ( TofuPlacement ) {
+#ifdef GRID_FJMPI
+    OptimalCommunicatorTofu(processors,optimal_comm,SHM);
+#else
+    if ( WorldRank==0 ) {
+      std::cerr << "Grid --tofu: not built with Fujitsu MPI" << std::endl;
+    }
+    MPI_Barrier(WorldComm);
+    MPI_Abort(WorldComm,EXIT_FAILURE);
+#endif
+    return;
+  }
   if(nscan==3 && HPEhypercube ) OptimalCommunicatorHypercube(processors,optimal_comm,SHM);
   else                          OptimalCommunicatorSharedMemory(processors,optimal_comm,SHM);
 }
@@ -450,6 +466,202 @@ void GlobalSharedMemory::OptimalCommunicatorSharedMemory(const Coordinate &proce
   int ierr= MPI_Comm_split(WorldComm,0,rank,&optimal_comm);
   GRID_ASSERT(ierr==0);
 }
+#ifdef GRID_FJMPI
+////////////////////////////////////////////////////////////////////////////////////////////
+// Tofu placement (--tofu): lay the node grid onto the job's logical torus (pjsub -L node=AxBxC:torus),
+// each direction along a job axis of the same extent. Errors are detected on all ranks alike.
+////////////////////////////////////////////////////////////////////////////////////////////
+static void TofuPlacementError(const std::string &msg)
+{
+  if ( GlobalSharedMemory::WorldRank==0 ) std::cerr << "Grid --tofu: " << msg << std::endl;
+  MPI_Barrier(GlobalSharedMemory::WorldComm);
+  MPI_Abort(GlobalSharedMemory::WorldComm,EXIT_FAILURE);
+}
+static std::string TofuJoin(const std::vector<int> &v,const char *sep)
+{
+  std::ostringstream os;
+  for(int i=0;i<v.size();i++) os << (i ? sep : "") << v[i];
+  return os.str();
+}
+void GlobalSharedMemory::OptimalCommunicatorTofu(const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &SHM)
+{
+  int ndimension = processors.size();
+  Coordinate WorldDims = processors; Coordinate ShmDims(ndimension);  Coordinate NodeDims (ndimension);
+  Coordinate ShmCoor(ndimension);    Coordinate NodeCoor(ndimension);   Coordinate WorldCoor(ndimension);
+
+  GetShmDims(WorldDims,ShmDims);
+  SHM=ShmDims;
+  int nodes=1;
+  for(int d=0;d<ndimension;d++){
+    NodeDims[d] = WorldDims[d]/ShmDims[d];
+    nodes *= NodeDims[d];
+  }
+  int Nprocessors=1;
+  for(int i=0;i<ndimension;i++){
+    Nprocessors*=processors[i];
+  }
+  GRID_ASSERT(WorldSize==Nprocessors);
+
+  std::vector<int> pvec(processors.begin(),processors.end());
+  std::vector<int> nvec(NodeDims.begin(),NodeDims.end());
+  std::vector<int> svec(ShmDims.begin(),ShmDims.end());
+  std::string layout = "--mpi " + TofuJoin(pvec,".") + " gives a node grid " + TofuJoin(nvec,".")
+    + " (" + std::to_string(nodes) + " nodes, ranks per node " + TofuJoin(svec,".") + ")";
+
+  ////////////////////////////////////////////////////////////////
+  // Shape of the job's logical node torus and this node's coordinate in it
+  ////////////////////////////////////////////////////////////////
+  int jdim=0;
+  if ( FJMPI_Topology_get_dimension(&jdim)!=FJMPI_SUCCESS || jdim<1 || jdim>3 ) {
+    TofuPlacementError("cannot read the job's node shape (FJMPI_Topology_get_dimension); "
+                       "is this a Fujitsu MPI job on whole nodes?");
+  }
+  std::vector<int> jshape(3,1), jcoor(3,0);
+  if ( FJMPI_Topology_get_shape(&jshape[0],&jshape[1],&jshape[2])!=FJMPI_SUCCESS ) {
+    TofuPlacementError("cannot read the job's node shape (FJMPI_Topology_get_shape)");
+  }
+  for(int a=jdim;a<3;a++) jshape[a]=1;
+  if ( FJMPI_Topology_get_coords(WorldComm,WorldRank,FJMPI_LOGICAL,jdim,&jcoor[0])!=FJMPI_SUCCESS ) {
+    TofuPlacementError("cannot read the node coordinates (FJMPI_Topology_get_coords)");
+  }
+  std::vector<int> jvec(jshape.begin(),jshape.begin()+jdim);
+  int jnodes = jshape[0]*jshape[1]*jshape[2];
+
+  ////////////////////////////////////////////////////////////////
+  // Lay each node grid direction of extent > 1 along an unused job axis of the same extent
+  ////////////////////////////////////////////////////////////////
+  std::vector<int> axis(ndimension,-1);
+  std::vector<int> used(3,0);
+  bool solvable = (nodes==WorldNodes) && (jnodes==WorldNodes);
+  for(int d=0;d<ndimension && solvable;d++){
+    if ( NodeDims[d]==1 ) continue;
+    for(int a=0;a<3;a++){
+      if ( !used[a] && jshape[a]==NodeDims[d] ) { axis[d]=a; used[a]=1; break; }
+    }
+    if ( axis[d]==-1 ) solvable=false;
+  }
+  for(int a=0;a<3;a++) if ( jshape[a]>1 && !used[a] ) solvable=false;
+  if ( !solvable ) {
+    std::vector<int> want;
+    for(int d=0;d<ndimension;d++) if ( NodeDims[d]>1 ) want.push_back(NodeDims[d]);
+    std::ostringstream os;
+    os << layout << ", which cannot be laid onto the job's node shape " << TofuJoin(jvec,"x")
+       << " (" << jnodes << " nodes, " << WorldNodes << " in use): every node grid extent larger than 1"
+       << " must equal a different extent of the job shape. ";
+    if ( want.size()>=1 && want.size()<=3 ) {
+      while ( want.size()<3 ) want.push_back(1);
+      os << "Request a job shape with the extents " << TofuJoin(want,", ") << " in any order"
+         << " (pjsub -L \"node=" << TofuJoin(want,"x") << ":torus\"), or change --mpi.";
+    } else {
+      os << "At most 3 node grid directions may have more than one node; change --mpi.";
+    }
+    TofuPlacementError(os.str());
+  }
+
+  ////////////////////////////////////////////////////////////////
+  // Processor coordinate: node block from the torus, on-node block as in the default layout
+  ////////////////////////////////////////////////////////////////
+  for(int d=0;d<ndimension;d++) NodeCoor[d] = (axis[d]>=0) ? jcoor[axis[d]] : 0;
+  Lexicographic::CoorFromIndexReversed(ShmCoor ,WorldShmRank,ShmDims);
+  for(int d=0;d<ndimension;d++) WorldCoor[d] = NodeCoor[d]*ShmDims[d]+ShmCoor[d];
+  int rank;
+  Lexicographic::IndexFromCoorReversed(WorldCoor,rank,WorldDims);
+
+  ////////////////////////////////////////////////////////////////
+  // Every processor coordinate must be taken exactly once
+  ////////////////////////////////////////////////////////////////
+  std::vector<int> ranks(WorldSize);
+  MPI_Allgather(&rank,1,MPI_INT,&ranks[0],1,MPI_INT,WorldComm);
+  std::vector<int> owner(WorldSize,-1);
+  for(int r=0;r<WorldSize;r++){
+    int k = ranks[r];
+    if ( owner[k]!=-1 ) {
+      std::ostringstream os;
+      os << layout << ": world ranks " << owner[k] << " and " << r << " get the same position; the job's"
+         << " logical node coordinates are not one per node, or the ranks per node differ from " << WorldShmSize;
+      TofuPlacementError(os.str());
+    }
+    owner[k]=r;
+  }
+
+  ////////////////////////////////////////////////////////////////
+  // Log the placement once per processor grid: summary, then (Debug) one line per rank
+  // and the Tofu hop counts between neighbouring ranks on different nodes
+  ////////////////////////////////////////////////////////////////
+  static std::vector<Coordinate> logged;
+  if ( std::find(logged.begin(),logged.end(),processors)==logged.end() ) {
+    logged.push_back(processors);
+    struct TofuRankInfo {
+      int tofu[6], logical[3];
+      char host[32];
+    } me;
+    for(int i=0;i<6;i++) me.tofu[i]=-1;
+    FJMPI_Topology_get_coords(WorldComm,WorldRank,FJMPI_TOFU_SYS,6,me.tofu);
+    for(int i=0;i<3;i++) me.logical[i]=jcoor[i];
+    memset(me.host,0,sizeof(me.host));
+    gethostname(me.host,sizeof(me.host)-1);
+    std::vector<TofuRankInfo> all(WorldSize);
+    MPI_Allgather(&me,sizeof(me),MPI_BYTE,&all[0],sizeof(me),MPI_BYTE,WorldComm);
+
+    std::vector<std::string> along;
+    for(int d=0;d<ndimension;d++){
+      if ( axis[d]>=0 ) along.push_back("dimension " + std::to_string(d) + " along job axis " + std::to_string(axis[d]));
+    }
+    std::ostringstream os;
+    for(int i=0;i<along.size();i++) os << (i ? ", " : "") << along[i];
+    std::cout << GridLogMessage << "Tofu placement: " << layout << ", job shape " << TofuJoin(jvec,"x")
+              << ": " << os.str() << std::endl;
+
+    // Tofu hops: sum over X,Y,Z,A,B,C, B a ring of 3; approximate, X,Y,Z are taken without wrap
+    auto hops = [](const TofuRankInfo &p,const TofuRankInfo &q) {
+      int h=0;
+      for(int i=0;i<6;i++){
+        int dd = std::abs(p.tofu[i]-q.tofu[i]);
+        if ( i==4 ) dd = std::min(dd,3-dd);
+        h += dd;
+      }
+      return h;
+    };
+    for(int k=0;k<WorldSize;k++){
+      const TofuRankInfo &r = all[owner[k]];
+      Coordinate c;
+      Lexicographic::CoorFromIndexReversed(c,k,WorldDims);
+      std::vector<int> cv(c.begin(),c.end());
+      char line[256];
+      snprintf(line,sizeof(line),"rank %5d coor %-14s world %5d host %-16s tofu %2d %2d %2d %d %d %d logical %2d %2d %2d",
+               k,TofuJoin(cv,".").c_str(),owner[k],r.host,r.tofu[0],r.tofu[1],r.tofu[2],r.tofu[3],r.tofu[4],r.tofu[5],
+               r.logical[0],r.logical[1],r.logical[2]);
+      std::cout << GridLogDebug << "Tofu placement: " << line << std::endl;
+    }
+    for(int d=0;d<ndimension;d++){
+      if ( WorldDims[d]==1 ) continue;
+      std::map<int,int> hist;
+      int onnode=0;
+      for(int k=0;k<WorldSize;k++){
+        Coordinate c;
+        Lexicographic::CoorFromIndexReversed(c,k,WorldDims);
+        c[d] = (c[d]+1)%WorldDims[d];
+        int n;
+        Lexicographic::IndexFromCoorReversed(c,n,WorldDims);
+        const TofuRankInfo &p = all[owner[k]], &q = all[owner[n]];
+        if ( strncmp(p.host,q.host,sizeof(p.host))==0 ) onnode++;
+        else hist[hops(p,q)]++;
+      }
+      std::ostringstream hs;
+      for(auto &h : hist) hs << " " << h.second << " at " << h.first << (h.first==1 ? " hop;" : " hops;");
+      std::cout << GridLogDebug << "Tofu placement: dimension " << d << " (" << WorldDims[d] << " ranks): "
+                << onnode << " of " << WorldSize << " +1 neighbours on the same node;" << hs.str()
+                << " (hops approximate, X/Y/Z not wrapped)" << std::endl;
+    }
+  }
+
+  /////////////////////////////////////////////////////////////////
+  // Build the new communicator
+  /////////////////////////////////////////////////////////////////
+  int ierr= MPI_Comm_split(WorldComm,0,rank,&optimal_comm);
+  GRID_ASSERT(ierr==0);
+}
+#endif
 ////////////////////////////////////////////////////////////////////////////////////////////
 // SHMGET
 ////////////////////////////////////////////////////////////////////////////////////////////

Grid/communicator/SharedMemoryNone.cc

@@ -50,6 +50,10 @@ void GlobalSharedMemory::Init(Grid_MPI_Comm comm)
 
 void GlobalSharedMemory::OptimalCommunicator(const Coordinate &processors,Grid_MPI_Comm & optimal_comm,Coordinate &SHM)
 {
+  if ( TofuPlacement ) {
+    std::cerr << "Grid --tofu: using a comms=none build" << std::endl;
+    exit(EXIT_FAILURE);
+  }
   optimal_comm = WorldComm;
   SHM = Coordinate(processors.size(),1);
 }

Grid/lattice/Lattice_transfer.h

@@ -1550,174 +1550,277 @@ void precisionChange(Lattice<VobjOut> &out, const Lattice<VobjIn> &in){
 ////////////////////////////////////////////////////////////////////////////////
 // Communicate between grids
 ////////////////////////////////////////////////////////////////////////////////
-///////////////////////////////////////////////////////////////////////////////////////////////////////////
-// SIMPLE CASE:
-///////////////////////////////////////////////////////////////////////////////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Grid_split / Grid_unsplit
 //
-// Mesh of nodes (2x2) ; subdivide to  1x1 subdivisions
+// The split grid divides the full grid's processor layout Pf into partitions of layout
+// Ps; ratio R = Pf/Ps per dimension, nvector = prod R partitions. A rank at full processor
+// coordinate c belongs to partition q = c/Ps at split coordinate s = c%Ps, and partition q
+// holds vector index lex(q,R) (GridSplitVectorIndex). Local volumes satisfy
+// Ls = Lf * R per dimension.
 //
-// Lex ord:   
-//          N0 va0 vb0 vc0 vd0       N1 va1 vb1 vc1 vd1  
-//          N2 va2 vb2 vc2 vd2       N3 va3 vb3 vc3 vd3 
+// Every full-grid rank's local block of field v lies wholly inside one split-grid rank's
+// local volume, in the sub-box at offset o = c%R (units of Lf), on the rank of partition
+// q(v) at split coordinate c/R. The redistribution is therefore nvector whole-block
+// point-to-point exchanges, nothing finer, and no data moves twice. Step k (a coordinate in
+// the R box) pairs every rank with one sender and one receiver:
 //
-// Ratio = full[dim] / split[dim]
+//   send  field q = (c%R + k)%R      to   full coordinate q*Ps + c/R
+//   recv  sub-box o = (c/Ps - k)%R   from full coordinate (c%Ps)*R + o
 //
-// For each dimension do an all to all; get Nvec -> Nvec / ratio
-//                                          Ldim -> Ldim * ratio
-//                                          LocalVol -> LocalVol * ratio
-// full AllToAll(0)
-//          N0 va0 vb0 va1 vb1       N1 vc0 vd0 vc1 vd1   
-//          N2 va2 vb2 va3 vb3       N3 vc2 vd2 vc3 vd3 
+// Within a step the send map is a permutation of the ranks and the receive map is its
+// inverse, so each step is one symmetric SendToRecvFrom per rank; a rank paired with
+// itself copies locally. Grid_unsplit runs the same steps with the roles exchanged.
 //
-// REARRANGE
-//          N0 va01 vb01      N1 vc01 vd01
-//          N2 va23 vb23      N3 vc23 vd23
-//
-// full AllToAll(1)           // Not what is wanted. FIXME
-//          N0 va01 va23      N1 vc01 vc23 
-//          N2 vb01 vb23      N3 vd01 vd23
-// 
-// REARRANGE
-//          N0 va0123      N1 vc0123
-//          N2 vb0123      N3 vd0123
-//
-// Must also rearrange data to get into the NEW lex order of grid at each stage. Some kind of "insert/extract".
-// NB: Easiest to programme if keep in lex order.
-/*
- *  Let chunk = (fvol*nvec)/sP be size of a chunk.         ( Divide lexico vol * nvec into fP/sP = M chunks )
- *  
- *  2nd A2A (over sP nodes; subdivide the fP into sP chunks of M)
- * 
- *     node 0     1st chunk of node 0M..(1M-1); 2nd chunk of node 0M..(1M-1)..   data chunk x M x sP = fL / sP * M * sP = fL * M growth
- *     node 1     1st chunk of node 1M..(2M-1); 2nd chunk of node 1M..(2M-1)..
- *     node 2     1st chunk of node 2M..(3M-1); 2nd chunk of node 2M..(3M-1)..
- *     node 3     1st chunk of node 3M..(3M-1); 2nd chunk of node 2M..(3M-1)..
- *  etc...
- */
+// Data stays on the device. A block on the full grid is the local volume in lexicographic
+// order; the split grid's local volume is staged box-major (sub-box index, then
+// lexicographic within the sub-box), so every exchanged block is contiguous. Packing and
+// unpacking are accelerator kernels. Buffers are deviceVectors: the staged split volume
+// plus one block. Without ACCELERATOR_AWARE_MPI each block is staged through host memory.
+/////////////////////////////////////////////////////////////////////////////////////////////
+
+// Scalar sites of in, written to buf in box-major order: sub-boxes of extent bdims taken
+// lexicographically, sites lexicographically within each. bdims = local dimensions gives
+// plain lexicographic order.
+template<class vobj>
+void GridSplitUnvectorise(const Lattice<vobj> &in,typename vobj::scalar_object *buf,const Coordinate &bdims)
+{
+  GridBase *grid = in.Grid();
+  int nd = grid->_ndimension;
+  Coordinate rdims = grid->_rdimensions;
+  Coordinate simd  = grid->_simd_layout;
+  Coordinate nbox(nd);
+  int64_t bvol = 1;
+  for(int d=0;d<nd;d++){
+    GRID_ASSERT( (grid->_ldimensions[d] % bdims[d]) == 0 );
+    nbox[d] = grid->_ldimensions[d] / bdims[d];
+    bvol   *= bdims[d];
+  }
+  const int Nsimd = grid->Nsimd();
+  autoView( in_v , in, AcceleratorRead);
+  accelerator_for(ss, grid->oSites(), 1, {
+    Coordinate ocoor(nd);
+    Coordinate icoor(nd);
+    Coordinate bcoor(nd);
+    Coordinate wcoor(nd);
+    Lexicographic::CoorFromIndex(ocoor,ss,rdims);
+    for(int lane=0;lane<Nsimd;lane++){
+      Lexicographic::CoorFromIndex(icoor,lane,simd);
+      for(int d=0;d<nd;d++){
+        int l = ocoor[d] + rdims[d]*icoor[d];
+        bcoor[d] = l / bdims[d];
+        wcoor[d] = l % bdims[d];
+      }
+      int64_t b;
+      int64_t w;
+      Lexicographic::IndexFromCoor(bcoor,b,nbox);
+      Lexicographic::IndexFromCoor(wcoor,w,bdims);
+      buf[b*bvol+w] = extractLane(lane,in_v[ss]);
+    }
+  });
+}
+
+// Inverse of GridSplitUnvectorise: every site of out is written
+template<class vobj>
+void GridSplitVectorise(const typename vobj::scalar_object *buf,Lattice<vobj> &out,const Coordinate &bdims)
+{
+  GridBase *grid = out.Grid();
+  int nd = grid->_ndimension;
+  Coordinate rdims = grid->_rdimensions;
+  Coordinate simd  = grid->_simd_layout;
+  Coordinate nbox(nd);
+  int64_t bvol = 1;
+  for(int d=0;d<nd;d++){
+    GRID_ASSERT( (grid->_ldimensions[d] % bdims[d]) == 0 );
+    nbox[d] = grid->_ldimensions[d] / bdims[d];
+    bvol   *= bdims[d];
+  }
+  const int Nsimd = grid->Nsimd();
+  autoView( out_v , out, AcceleratorWriteDiscard);
+  accelerator_for(ss, grid->oSites(), 1, {
+    Coordinate ocoor(nd);
+    Coordinate icoor(nd);
+    Coordinate bcoor(nd);
+    Coordinate wcoor(nd);
+    Lexicographic::CoorFromIndex(ocoor,ss,rdims);
+    for(int lane=0;lane<Nsimd;lane++){
+      Lexicographic::CoorFromIndex(icoor,lane,simd);
+      for(int d=0;d<nd;d++){
+        int l = ocoor[d] + rdims[d]*icoor[d];
+        bcoor[d] = l / bdims[d];
+        wcoor[d] = l % bdims[d];
+      }
+      int64_t b;
+      int64_t w;
+      Lexicographic::IndexFromCoor(bcoor,b,nbox);
+      Lexicographic::IndexFromCoor(wcoor,w,bdims);
+      insertLane(lane,out_v[ss],buf[b*bvol+w]);
+    }
+  });
+}
+
+// Step schedule shared by Grid_split and Grid_unsplit; see the description above
+class GridSplitSchedule
+{
+public:
+  int        nd;
+  int        nvector;
+  Coordinate ratio;      // R
+  Coordinate fldims;     // full grid local dimensions = sub-box extent on the split grid
+  GridBase  *full_grid;
+
+  // Step k: this rank sends field Field(k) to rank FieldRank(k) and receives sub-box
+  // Box(k) from rank BoxRank(k) (Grid_split); Grid_unsplit reverses both
+  std::vector<int> Field;
+  std::vector<int> FieldRank;
+  std::vector<int> Box;
+  std::vector<int> BoxRank;
+
+  GridSplitSchedule(GridBase *_full_grid,GridBase *split_grid) : full_grid(_full_grid)
+  {
+    nd = full_grid->_ndimension;
+    GRID_ASSERT(split_grid->_ndimension == nd);
+    ratio.resize(nd);
+    fldims = full_grid->_ldimensions;
+    Coordinate sprocs = split_grid->_processors;
+    Coordinate fcoor  = full_grid->_processor_coor;
+    nvector = 1;
+    for(int d=0;d<nd;d++){
+      GRID_ASSERT(full_grid->_gdimensions[d] == split_grid->_gdimensions[d]);
+      GRID_ASSERT(full_grid->_fdimensions[d] == split_grid->_fdimensions[d]);
+      GRID_ASSERT( (full_grid->_processors[d] % sprocs[d]) == 0 );
+      ratio[d] = full_grid->_processors[d] / sprocs[d];
+      GRID_ASSERT(split_grid->_ldimensions[d] == fldims[d]*ratio[d]);
+      nvector *= ratio[d];
+    }
+    GRID_ASSERT(nvector*split_grid->_Nprocessors == full_grid->_Nprocessors);
+
+    Field.resize(nvector);
+    FieldRank.resize(nvector);
+    Box.resize(nvector);
+    BoxRank.resize(nvector);
+    for(int k=0;k<nvector;k++){
+      Coordinate kc(nd);
+      Coordinate q(nd);
+      Coordinate dest(nd);
+      Coordinate o(nd);
+      Coordinate from(nd);
+      Lexicographic::CoorFromIndex(kc,k,ratio);
+      for(int d=0;d<nd;d++){
+        q[d]    = (fcoor[d] % ratio[d] + kc[d]) % ratio[d];
+        dest[d] = q[d]*sprocs[d] + fcoor[d] / ratio[d];
+        o[d]    = (fcoor[d] / sprocs[d] - kc[d] + ratio[d]) % ratio[d];
+        from[d] = (fcoor[d] % sprocs[d])*ratio[d] + o[d];
+      }
+      Lexicographic::IndexFromCoor(q,Field[k],ratio);
+      Lexicographic::IndexFromCoor(o,Box[k],ratio);
+      FieldRank[k] = full_grid->RankFromProcessorCoor(dest);
+      BoxRank[k]   = full_grid->RankFromProcessorCoor(from);
+    }
+  }
+
+  // One step's exchange; returns the bytes that left this rank
+  uint64_t Exchange(void *xmit,int dest,void *recv,int from,uint64_t bytes)
+  {
+    int me = full_grid->ThisRank();
+    // The send map is a permutation and the receive map its inverse: self pairs with self
+    GRID_ASSERT( (dest == me) == (from == me) );
+    if ( dest == me ) {
+      acceleratorCopyDeviceToDevice(xmit,recv,bytes);
+      return 0;
+    }
+    GRID_ASSERT( (bytes % sizeof(int32_t)) == 0 );
+    GRID_ASSERT( bytes/sizeof(int32_t) < (1ULL<<31) );   // SendToRecvFrom counts int32 words
+#ifdef ACCELERATOR_AWARE_MPI
+    full_grid->SendToRecvFrom(xmit,dest,recv,from,bytes);
+#else
+    std::vector<char> hxmit(bytes);
+    std::vector<char> hrecv(bytes);
+    acceleratorCopyFromDevice(xmit,&hxmit[0],bytes);
+    full_grid->SendToRecvFrom(&hxmit[0],dest,&hrecv[0],from,bytes);
+    acceleratorCopyToDevice(&hrecv[0],recv,bytes);
+#endif
+    return bytes;
+  }
+};
+
+// Common body: full[v] may alias one another (the single field overload)
 template<class Vobj>
-void Grid_split(std::vector<Lattice<Vobj> > & full,Lattice<Vobj>   & split)
+void Grid_split(const std::vector<const Lattice<Vobj> *> &full,Lattice<Vobj> &split)
 {
   typedef typename Vobj::scalar_object Sobj;
 
-  int full_vecs   = full.size();
-
-  GRID_ASSERT(full_vecs>=1);
-
-  GridBase * full_grid = full[0].Grid();
+  GRID_ASSERT(full.size()>=1);
+  GridBase *full_grid  = full[0]->Grid();
   GridBase *split_grid = split.Grid();
 
-  int       ndim  = full_grid->_ndimension;
-  int  full_nproc = full_grid->_Nprocessors;
-  int split_nproc =split_grid->_Nprocessors;
-
-  ////////////////////////////////
-  // Checkerboard management
-  ////////////////////////////////
-  int cb = full[0].Checkerboard();
+  int cb = full[0]->Checkerboard();
   split.Checkerboard() = cb;
-
-  //////////////////////////////
-  // Checks
-  //////////////////////////////
-  GRID_ASSERT(full_grid->_ndimension==split_grid->_ndimension);
-  for(int n=0;n<full_vecs;n++){
-    GRID_ASSERT(full[n].Checkerboard() == cb);
-    for(int d=0;d<ndim;d++){
-      GRID_ASSERT(full[n].Grid()->_gdimensions[d]==split.Grid()->_gdimensions[d]);
-      GRID_ASSERT(full[n].Grid()->_fdimensions[d]==split.Grid()->_fdimensions[d]);
-    }
+  for(int n=0;n<full.size();n++){
+    GRID_ASSERT(full[n]->Grid() == full_grid);
+    GRID_ASSERT(full[n]->Checkerboard() == cb);
   }
 
-  int   nvector   =full_nproc/split_nproc; 
-  GRID_ASSERT(nvector*split_nproc==full_nproc);
-  GRID_ASSERT(nvector == full_vecs);
+  double t_start = usecond();
+  GridSplitSchedule sched(full_grid,split_grid);
+  GRID_ASSERT(sched.nvector == full.size());
 
-  Coordinate ratio(ndim);
-  for(int d=0;d<ndim;d++){
-    ratio[d] = full_grid->_processors[d]/ split_grid->_processors[d];
+  uint64_t fvol  = full_grid->lSites();
+  uint64_t bytes = fvol*sizeof(Sobj);
+  deviceVector<Sobj> staged(split_grid->lSites());
+  deviceVector<Sobj> block(fvol);
+  Sobj *staged_p = &staged[0];
+  Sobj *block_p  = &block[0];
+  double t_alloc = usecond() - t_start;
+
+  double t_pack = 0;
+  double t_comm = 0;
+  uint64_t offrank = 0;
+  for(int k=0;k<sched.nvector;k++){
+    double t0 = usecond();
+    GridSplitUnvectorise(*full[sched.Field[k]],block_p,full_grid->_ldimensions);
+    double t1 = usecond();
+    offrank += sched.Exchange((void *)block_p,sched.FieldRank[k],
+                              (void *)&staged_p[sched.Box[k]*fvol],sched.BoxRank[k],bytes);
+    double t2 = usecond();
+    t_pack += t1-t0;
+    t_comm += t2-t1;
   }
 
-  uint64_t lsites = full_grid->lSites();
-  uint64_t     sz = lsites * nvector;
-  std::vector<Sobj> tmpdata(sz);
-  std::vector<Sobj> alldata(sz);
-  std::vector<Sobj> scalardata(lsites); 
+  double t0 = usecond();
+  GridSplitVectorise(staged_p,split,sched.fldims);
+  double t_unpack = usecond() - t0;
+  double t_total  = usecond() - t_start;
 
-  for(int v=0;v<nvector;v++){
-    unvectorizeToLexOrdArray(scalardata,full[v]);    
-    thread_for(site,lsites,{
-      alldata[v*lsites+site] = scalardata[site];
-    });
+  std::cout << GridLogPerformance << "Grid_split: " << sched.nvector << " x " << bytes/1.0e6
+            << " MB/rank total " << t_total/1.0e6 << " s : alloc " << t_alloc/1.0e6
+            << " pack " << t_pack/1.0e6
+            << " exchange " << t_comm/1.0e6 << " (" << offrank/1.0e6 << " MB off rank, "
+            << offrank/(t_comm+1.0e-9)/1.0e3 << " GB/s)"
+            << " unpack " << t_unpack/1.0e6 << std::endl;
+  // Staging buffers are still live here, so this is the high-water point of the call
+  if ( GridLogPerformance.isActive() ) {
+    HostMemoryReport(full_grid,GridLogPerformance,"Grid_split");
   }
-
-  int nvec = nvector; // Counts down to 1 as we collapse dims
-  Coordinate ldims = full_grid->_ldimensions;
-
-  for(int d=ndim-1;d>=0;d--){
-
-    if ( ratio[d] != 1 ) {
-
-      full_grid ->AllToAll(d,alldata,tmpdata);
-      if ( split_grid->_processors[d] > 1 ) {
-	alldata=tmpdata;
-	split_grid->AllToAll(d,alldata,tmpdata);
-      }
-
-      auto rdims = ldims; 
-      auto     M = ratio[d];
-      auto rsites= lsites*M;// increases rsites by M
-      nvec      /= M;       // Reduce nvec by subdivision factor
-      rdims[d]  *= M;       // increase local dim by same factor
-
-      int sP =   split_grid->_processors[d];
-      int fP =    full_grid->_processors[d];
-
-      int fvol   = lsites;
-      
-      int chunk  = (nvec*fvol)/sP;          GRID_ASSERT(chunk*sP == nvec*fvol);
-
-      // Loop over reordered data post A2A
-      thread_for(c, chunk, {
-	Coordinate coor(ndim);
-	for(int m=0;m<M;m++){
-	  for(int s=0;s<sP;s++){
-	    
-	    // addressing; use lexico
-	    int lex_r;
-	    uint64_t lex_c        = c+chunk*m+chunk*M*s;
-	    uint64_t lex_fvol_vec = c+chunk*s;
-	    uint64_t lex_fvol     = lex_fvol_vec%fvol;
-	    uint64_t lex_vec      = lex_fvol_vec/fvol;
-
-	    // which node sets an adder to the coordinate
-	    Lexicographic::CoorFromIndex(coor, lex_fvol, ldims);	  
-	    coor[d] += m*ldims[d];
-	    Lexicographic::IndexFromCoor(coor, lex_r, rdims);	  
-	    lex_r += lex_vec * rsites;
-
-	    // LexicoFind coordinate & vector number within split lattice
-	    alldata[lex_r] = tmpdata[lex_c];
-
-	  }
-	}
-      });
-      ldims[d]*= ratio[d];
-      lsites  *= ratio[d];
-
-    }
-  }
-  vectorizeFromLexOrdArray(alldata,split);    
 }
 
+template<class Vobj>
+void Grid_split(std::vector<Lattice<Vobj> > & full,Lattice<Vobj>   & split)
+{
+  std::vector<const Lattice<Vobj> *> full_p(full.size());
+  for(int n=0;n<full.size();n++){
+    full_p[n] = &full[n];
+  }
+  Grid_split(full_p,split);
+}
+
+// The same field to every partition
 template<class Vobj>
 void Grid_split(Lattice<Vobj> &full,Lattice<Vobj>   & split)
 {
   int nvector = full.Grid()->_Nprocessors / split.Grid()->_Nprocessors;
-  std::vector<Lattice<Vobj> > full_v(nvector,full.Grid());
-  for(int n=0;n<nvector;n++){
-    full_v[n] = full;
-  }
-  Grid_split(full_v,split);
+  std::vector<const Lattice<Vobj> *> full_p(nvector,&full);
+  Grid_split(full_p,split);
 }
 
 template<class Vobj>
@@ -1725,119 +1828,57 @@ void Grid_unsplit(std::vector<Lattice<Vobj> > & full,Lattice<Vobj>   & split)
 {
   typedef typename Vobj::scalar_object Sobj;
 
-  int full_vecs   = full.size();
-
-  GRID_ASSERT(full_vecs>=1);
-
-  GridBase * full_grid = full[0].Grid();
+  GRID_ASSERT(full.size()>=1);
+  GridBase *full_grid  = full[0].Grid();
   GridBase *split_grid = split.Grid();
 
-  int       ndim  = full_grid->_ndimension;
-  int  full_nproc = full_grid->_Nprocessors;
-  int split_nproc =split_grid->_Nprocessors;
-
-  ////////////////////////////////
-  // Checkerboard management
-  ////////////////////////////////
   int cb = full[0].Checkerboard();
   split.Checkerboard() = cb;
-
-  //////////////////////////////
-  // Checks
-  //////////////////////////////
-  GRID_ASSERT(full_grid->_ndimension==split_grid->_ndimension);
-  for(int n=0;n<full_vecs;n++){
+  for(int n=0;n<full.size();n++){
+    GRID_ASSERT(full[n].Grid() == full_grid);
     GRID_ASSERT(full[n].Checkerboard() == cb);
-    for(int d=0;d<ndim;d++){
-      GRID_ASSERT(full[n].Grid()->_gdimensions[d]==split.Grid()->_gdimensions[d]);
-      GRID_ASSERT(full[n].Grid()->_fdimensions[d]==split.Grid()->_fdimensions[d]);
-    }
   }
 
-  int   nvector   =full_nproc/split_nproc; 
-  GRID_ASSERT(nvector*split_nproc==full_nproc);
-  GRID_ASSERT(nvector == full_vecs);
+  double t_start = usecond();
+  GridSplitSchedule sched(full_grid,split_grid);
+  GRID_ASSERT(sched.nvector == full.size());
 
-  Coordinate ratio(ndim);
-  for(int d=0;d<ndim;d++){
-    ratio[d] = full_grid->_processors[d]/ split_grid->_processors[d];
+  uint64_t fvol  = full_grid->lSites();
+  uint64_t bytes = fvol*sizeof(Sobj);
+  deviceVector<Sobj> staged(split_grid->lSites());
+  deviceVector<Sobj> block(fvol);
+  Sobj *staged_p = &staged[0];
+  Sobj *block_p  = &block[0];
+  double t_alloc = usecond() - t_start;
+
+  double t0 = usecond();
+  GridSplitUnvectorise(split,staged_p,sched.fldims);
+  double t_pack = usecond() - t0;
+
+  double t_comm   = 0;
+  double t_unpack = 0;
+  uint64_t offrank = 0;
+  for(int k=0;k<sched.nvector;k++){
+    double t1 = usecond();
+    offrank += sched.Exchange((void *)&staged_p[sched.Box[k]*fvol],sched.BoxRank[k],
+                              (void *)block_p,sched.FieldRank[k],bytes);
+    double t2 = usecond();
+    GridSplitVectorise(block_p,full[sched.Field[k]],full_grid->_ldimensions);
+    double t3 = usecond();
+    t_comm   += t2-t1;
+    t_unpack += t3-t2;
   }
+  double t_total = usecond() - t_start;
 
-  uint64_t lsites = full_grid->lSites();
-  uint64_t     sz = lsites * nvector;
-  std::vector<Sobj> tmpdata(sz);
-  std::vector<Sobj> alldata(sz);
-  std::vector<Sobj> scalardata(lsites); 
-
-  unvectorizeToLexOrdArray(alldata,split);    
-
-  /////////////////////////////////////////////////////////////////
-  // Start from split grid and work towards full grid
-  /////////////////////////////////////////////////////////////////
-
-  int nvec = 1;
-  uint64_t rsites        = split_grid->lSites();
-  Coordinate rdims = split_grid->_ldimensions;
-
-  for(int d=0;d<ndim;d++){
-
-    if ( ratio[d] != 1 ) {
-
-      auto     M = ratio[d];
-
-      int sP =   split_grid->_processors[d];
-      int fP =    full_grid->_processors[d];
-      
-      auto ldims = rdims;  ldims[d]  /= M;  // Decrease local dims by same factor
-      auto lsites= rsites/M;                // Decreases rsites by M
-      
-      int fvol   = lsites;
-      int chunk  = (nvec*fvol)/sP;          GRID_ASSERT(chunk*sP == nvec*fvol);
-	
-      {
-	// Loop over reordered data post A2A
-	thread_for(c, chunk,{
-	  Coordinate coor(ndim);
-	  for(int m=0;m<M;m++){
-	    for(int s=0;s<sP;s++){
-
-	      // addressing; use lexico
-	      int lex_r;
-	      uint64_t lex_c = c+chunk*m+chunk*M*s;
-	      uint64_t lex_fvol_vec = c+chunk*s;
-	      uint64_t lex_fvol     = lex_fvol_vec%fvol;
-	      uint64_t lex_vec      = lex_fvol_vec/fvol;
-	      
-	      // which node sets an adder to the coordinate
-	      Lexicographic::CoorFromIndex(coor, lex_fvol, ldims);	  
-	      coor[d] += m*ldims[d];
-	      Lexicographic::IndexFromCoor(coor, lex_r, rdims);	  
-	      lex_r += lex_vec * rsites;
-	      
-	      // LexicoFind coordinate & vector number within split lattice
-	      tmpdata[lex_c] = alldata[lex_r];
-	    }
-	  }
-        });
-      }
-
-      if ( split_grid->_processors[d] > 1 ) {
-	split_grid->AllToAll(d,tmpdata,alldata);
-	tmpdata=alldata;
-      }
-      full_grid ->AllToAll(d,tmpdata,alldata);
-      rdims[d]/= M;
-      rsites  /= M;
-      nvec    *= M;       // Increase nvec by subdivision factor
-    }
-  }
-
-  lsites = full_grid->lSites();
-  for(int v=0;v<nvector;v++){
-    thread_for(site, lsites,{
-      scalardata[site] = alldata[v*lsites+site];
-    });
-    vectorizeFromLexOrdArray(scalardata,full[v]);    
+  std::cout << GridLogPerformance << "Grid_unsplit: " << sched.nvector << " x " << bytes/1.0e6
+            << " MB/rank total " << t_total/1.0e6 << " s : alloc " << t_alloc/1.0e6
+            << " pack " << t_pack/1.0e6
+            << " exchange " << t_comm/1.0e6 << " (" << offrank/1.0e6 << " MB off rank, "
+            << offrank/(t_comm+1.0e-9)/1.0e3 << " GB/s)"
+            << " unpack " << t_unpack/1.0e6 << std::endl;
+  // Staging buffers are still live here, so this is the high-water point of the call
+  if ( GridLogPerformance.isActive() ) {
+    HostMemoryReport(full_grid,GridLogPerformance,"Grid_unsplit");
   }
 }
 

Grid/perfmon/HostMemory.h

@@ -0,0 +1,227 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./lib/perfmon/HostMemory.h
+
+    Copyright (C) 2026
+
+Author: Peter Boyle <paboyle@ph.ed.ac.uk>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+*************************************************************************************/
+/*  END LEGAL */
+#pragma once
+
+#include <sys/resource.h>
+#include <unistd.h>
+#include <cctype>
+#include <fstream>
+#ifdef __APPLE__
+#include <mach/mach.h>
+#endif
+
+NAMESPACE_BEGIN(Grid);
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Host memory as the OOM killer sees it, plus the device allocator footprint.
+// Targets are Apple or Linux; on Apple only the RSS and device fields are filled.
+//
+//  RSS      : this process, now and its high-water mark (getrusage).
+//  Node     : MemAvailable and MemTotal from /proc/meminfo (Linux).
+//  Cgroup   : the nearest enclosing memory cgroup with a finite limit (the one the batch
+//             system enforces), its usage now, its high-water mark and the limit (Linux,
+//             v2 or v1 hierarchy). Cgroup usage includes page cache, tmpfs and kernel memory
+//             charged to the job, none of which appear in any process RSS.
+//  Device   : MemoryManager Lattice footprint on the device and its free-block cache. Device
+//             allocations outside the MemoryManager (comms buffers) are not included.
+//
+// All values in GB; -1 where the source does not exist on this platform or kernel.
+/////////////////////////////////////////////////////////////////////////////////////////////
+struct HostMemoryStatus
+{
+  RealD RSS           = -1;
+  RealD RSSPeak       = -1;
+  RealD NodeAvailable = -1;
+  RealD NodeTotal     = -1;
+  RealD CgroupCurrent = -1;
+  RealD CgroupPeak    = -1;
+  RealD CgroupLimit   = -1;
+  RealD DeviceLattice = -1;
+  RealD DeviceCache   = -1;
+};
+
+#ifndef __APPLE__ // Linux
+
+// One number from a file; -1 if the file is absent or holds no number (e.g. "max")
+inline RealD HostMemoryReadBytes(const std::string &file)
+{
+  std::ifstream f(file);
+  if ( !f.good() ) {
+    return -1;
+  }
+  std::string word;
+  f >> word;
+  if ( word.empty() || !isdigit(word[0]) ) {
+    return -1;
+  }
+  return std::stod(word);
+}
+
+// Value of a "Key: N kB" line in /proc/meminfo, in bytes; -1 if absent
+inline RealD HostMemoryMeminfo(const std::string &key)
+{
+  std::ifstream f("/proc/meminfo");
+  std::string name;
+  RealD value;
+  std::string unit;
+  while ( f >> name >> value >> unit ) {
+    if ( name == key + ":" ) {
+      return value*1024.0;
+    }
+  }
+  return -1;
+}
+
+// Fill the cgroup fields from the nearest enclosing memory cgroup with a finite limit, or
+// from this process's own cgroup (limit left at -1) if no ancestor has one
+inline void HostMemoryCgroup(HostMemoryStatus &s)
+{
+  std::ifstream f("/proc/self/cgroup");
+  std::string line;
+  std::string root;
+  std::string path;
+  std::string current;
+  std::string peak;
+  std::string limit;
+  while ( std::getline(f,line) ) {
+    size_t c1 = line.find(':');
+    size_t c2 = line.find(':',c1+1);
+    if ( c1 == std::string::npos || c2 == std::string::npos ) {
+      continue;
+    }
+    std::string controllers = line.substr(c1+1,c2-c1-1);
+    // v2 unified hierarchy: "0::/path"
+    if ( controllers.empty() && root.empty() ) {
+      root    = "/sys/fs/cgroup";
+      path    = line.substr(c2+1);
+      current = "memory.current";
+      peak    = "memory.peak";
+      limit   = "memory.max";
+    }
+    // v1 memory controller: "N:memory:/path", takes precedence
+    if ( (","+controllers+",").find(",memory,") != std::string::npos ) {
+      root    = "/sys/fs/cgroup/memory";
+      path    = line.substr(c2+1);
+      current = "memory.usage_in_bytes";
+      peak    = "memory.max_usage_in_bytes";
+      limit   = "memory.limit_in_bytes";
+      break;
+    }
+  }
+  if ( root.empty() ) {
+    return;
+  }
+
+  // v1 reports "no limit" as a number near 2^63
+  const RealD unlimited = 1.0e18;
+  std::string dir = path;
+  while ( true ) {
+    std::string base = root + dir + "/";
+    RealD lim = HostMemoryReadBytes(base+limit);
+    if ( lim > 0 && lim < unlimited ) {
+      s.CgroupCurrent = HostMemoryReadBytes(base+current)/1.0e9;
+      s.CgroupPeak    = HostMemoryReadBytes(base+peak)/1.0e9;
+      s.CgroupLimit   = lim/1.0e9;
+      return;
+    }
+    if ( dir.empty() || dir == "/" ) {
+      break;
+    }
+    dir = dir.substr(0,dir.rfind('/'));
+  }
+  std::string base = root + path + "/";
+  s.CgroupCurrent = HostMemoryReadBytes(base+current)/1.0e9;
+  s.CgroupPeak    = HostMemoryReadBytes(base+peak)/1.0e9;
+}
+
+#endif
+
+// This rank only; no communication
+inline HostMemoryStatus HostMemoryQuery(void)
+{
+  HostMemoryStatus s;
+
+  struct rusage ru;
+  getrusage(RUSAGE_SELF,&ru);
+#ifdef __APPLE__
+  s.RSSPeak = ru.ru_maxrss/1.0e9;                // bytes on macOS
+  mach_task_basic_info_data_t info;
+  mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;
+  if ( task_info(mach_task_self(),MACH_TASK_BASIC_INFO,(task_info_t)&info,&count) == KERN_SUCCESS ) {
+    s.RSS = info.resident_size/1.0e9;
+  }
+#else // Linux
+  s.RSSPeak = ru.ru_maxrss*1024.0/1.0e9;         // kilobytes on Linux
+  std::ifstream statm("/proc/self/statm");
+  long pages    = 0;
+  long resident = 0;
+  if ( statm >> pages >> resident ) {
+    s.RSS = resident*(RealD)sysconf(_SC_PAGESIZE)/1.0e9;
+  }
+  RealD avail = HostMemoryMeminfo("MemAvailable");
+  RealD total = HostMemoryMeminfo("MemTotal");
+  if ( avail >= 0 ) {
+    s.NodeAvailable = avail/1.0e9;
+  }
+  if ( total >= 0 ) {
+    s.NodeTotal = total/1.0e9;
+  }
+  HostMemoryCgroup(s);
+#endif
+
+  s.DeviceLattice = MemoryManager::DeviceBytes/1.0e9;
+  s.DeviceCache   = MemoryManager::DeviceCacheBytes()/1.0e9;
+  return s;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Collective over grid. One line on log: maxima over ranks, except node MemAvailable which
+// is the minimum over ranks (the tightest node).
+/////////////////////////////////////////////////////////////////////////////////////////////
+inline void HostMemoryReport(GridBase *grid,GridLogger &log,const std::string &label)
+{
+  HostMemoryStatus s = HostMemoryQuery();
+  RealD minus_available = -s.NodeAvailable;
+  grid->GlobalMax(s.RSS);
+  grid->GlobalMax(s.RSSPeak);
+  grid->GlobalMax(minus_available);
+  grid->GlobalMax(s.NodeTotal);
+  grid->GlobalMax(s.CgroupCurrent);
+  grid->GlobalMax(s.CgroupPeak);
+  grid->GlobalMax(s.CgroupLimit);
+  grid->GlobalMax(s.DeviceLattice);
+  grid->GlobalMax(s.DeviceCache);
+  std::cout << log << "HOSTMEM " << label
+            << " : rank RSS " << s.RSS << " peak " << s.RSSPeak
+            << " ; node available " << -minus_available << " of " << s.NodeTotal
+            << " ; cgroup " << s.CgroupCurrent << " peak " << s.CgroupPeak << " limit " << s.CgroupLimit
+            << " ; device lattice " << s.DeviceLattice << " cache " << s.DeviceCache
+            << " GB" << std::endl;
+}
+
+NAMESPACE_END(Grid);

Grid/qcd/action/ActionParams.h

@@ -53,6 +53,7 @@ struct WilsonImplParams {
   AcceleratorVector<Real,Nd> twist_n_2pi_L;
   AcceleratorVector<Complex,Nd> boundary_phases;
   WilsonImplParams()  {
+    overlapCommsCompute=true;
     dirichlet.resize(0);
     partialDirichlet=0;
     boundary_phases.resize(Nd, 1.0);

Grid/qcd/action/fermion/CayleyFermion5D.h

@@ -81,6 +81,15 @@ public:
     mass_minus=_mass_minus;
     SetCoefficientsInternal(_zolo_hi,_gamma,_b,_c);  // Reset coeffs
   } ;
+
+  // Everything SetMass needs, plus the boost, then rebuild the coefficient arrays in clone
+  void CloneCoefficientsInto(CayleyFermion5D<Impl> &clone)
+  {
+    clone.mass_plus  = mass_plus;
+    clone.mass_minus = mass_minus;
+    clone.qmu        = qmu;
+    clone.SetCoefficientsInternal(_zolo_hi,_gamma,_b,_c);
+  }
   void  P(const FermionField &psi, FermionField &chi);
   void  Pdag(const FermionField &psi, FermionField &chi);
   

Grid/qcd/action/fermion/CompactWilsonCloverFermion.h

@@ -189,6 +189,71 @@ public:
 
   void ImportGauge(const GaugeField& _Umu) override;
 
+  // Install already computed clover fields and inverses, and derive the checkerboarded parts
+  void ImportCloverTerms(const CloverDiagonalField& D,
+                         const CloverTriangleField& T,
+                         const CloverDiagonalField& Dinv,
+                         const CloverTriangleField& Tinv)
+  {
+    Diagonal    = D;
+    Triangle    = T;
+    DiagonalInv = Dinv;
+    TriangleInv = Tinv;
+    pickCheckerboard(Even, DiagonalEven,    Diagonal);
+    pickCheckerboard(Even, TriangleEven,    Triangle);
+    pickCheckerboard(Odd,  DiagonalOdd,     Diagonal);
+    pickCheckerboard(Odd,  TriangleOdd,     Triangle);
+    pickCheckerboard(Even, DiagonalInvEven, DiagonalInv);
+    pickCheckerboard(Even, TriangleInvEven, TriangleInv);
+    pickCheckerboard(Odd,  DiagonalInvOdd,  DiagonalInv);
+    pickCheckerboard(Odd,  TriangleInvOdd,  TriangleInv);
+  }
+
+  // Exact type only: derived operators must not inherit this
+  SplitOperator<FermionField> *SplitClone(const Coordinate& mpi_split) override
+  {
+    if ( typeid(*this) != typeid(CompactWilsonCloverFermion<Impl,CloverHelpers>) ) {
+      return nullptr;
+    }
+    if ( Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field and clover fields are overwritten below.
+    // Boundary masks are rebuilt by the constructor from global coordinates.
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    CompactWilsonCloverFermion<Impl,CloverHelpers> *clone =
+      new CompactWilsonCloverFermion<Impl,CloverHelpers>(Uplaceholder,
+                                                         *split->GaugeGrid,
+                                                         *split->GaugeRBGrid,
+                                                         this->mass,
+                                                         csw_r,
+                                                         csw_t,
+                                                         cF,
+                                                         this->anisotropyCoeff,
+                                                         this->Params);
+    // The stored coefficients are post-scaled; the constructor would scale them again
+    clone->csw_r = csw_r;
+    clone->csw_t = csw_t;
+
+    CloverDiagonalField sD(split->GaugeGrid);
+    CloverTriangleField sT(split->GaugeGrid);
+    CloverDiagonalField sDinv(split->GaugeGrid);
+    CloverTriangleField sTinv(split->GaugeGrid);
+    Grid_split(Diagonal,sD);
+    Grid_split(Triangle,sT);
+    Grid_split(DiagonalInv,sDinv);
+    Grid_split(TriangleInv,sTinv);
+    clone->ImportCloverTerms(sD,sT,sDinv,sTinv);
+
+    this->SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
   /////////////////////////////////////////////
   // Helpers
   /////////////////////////////////////////////

Grid/qcd/action/fermion/DWFSlow.h

@@ -109,10 +109,69 @@ public:
   ////////////////////////
   // Derivative interface
   ////////////////////////
-  // Interface calls an internal routine
-  void DhopDeriv(GaugeField &mat,const FermionField &U,const FermionField &V,int dag)  { GRID_ASSERT(0);};
-  void DhopDerivOE(GaugeField &mat,const FermionField &U,const FermionField &V,int dag){ GRID_ASSERT(0);};
-  void DhopDerivEO(GaugeField &mat,const FermionField &U,const FermionField &V,int dag){ GRID_ASSERT(0);};
+  // Force on every link of every s-slice, in the form WilsonFermion5D uses:
+  //   F5[mu](x,s) = tr_spin[ Btilde(x,s) A(x,s)^dag ]
+  //   Btilde(x,s) = -1/2 U_mu(x,s) (1 -+ gamma_mu) B(x+mu,s), lower sign for dag
+  // Dhop5 carries no gauge links and contributes nothing.
+  void DhopDeriv5D(std::vector<GaugeLinkField> &F5,const FermionField &A,const FermionField &B,int dag)
+  {
+    RealD sgn= 1.0;
+    if (dag ) sgn=-1.0;
+
+    Gamma::Algebra Gmu [] = {
+			 Gamma::Algebra::GammaX,
+			 Gamma::Algebra::GammaY,
+			 Gamma::Algebra::GammaZ,
+			 Gamma::Algebra::GammaT
+    };
+
+    FermionField tmp(B.Grid());
+    FermionField Btilde(B.Grid());
+    for(int mu=0;mu<Nd;mu++){
+      tmp    = Umu[mu] * Cshift(B,mu+1,1);
+      Btilde = tmp - Gamma(Gmu[mu])*tmp*sgn;
+      Btilde = -0.5*Btilde;
+      F5[mu] = TraceIndex<SpinIndex>(outerProduct(Btilde,A));
+    }
+  }
+
+  // Every slice holds the same 4D gauge field, so the 4D force is the sum over slices
+  void DhopDeriv(GaugeField &mat,const FermionField &A,const FermionField &B,int dag)
+  {
+    std::vector<GaugeLinkField> F5(Nd,_grid);
+    DhopDeriv5D(F5,A,B,dag);
+
+    GaugeLinkField slice(_grid4);
+    GaugeLinkField sum(_grid4);
+    for(int mu=0;mu<Nd;mu++){
+      sum = Zero();
+      for(int s=0;s<Ls;s++){
+        ExtractSlice(slice,F5[mu],s,0);
+        sum = sum + slice;
+      }
+      PokeIndex<LorentzIndex>(mat,sum,mu);
+    }
+  }
+
+  // The checkerboard here is 5D, so a field on one checkerboard touches links at both 4D
+  // parities: the even-odd forces are returned on the full 4D gauge grid.
+  int CheckerboardedForce(void) { return 0; };
+
+  // A and B on opposite 5D checkerboards, placed in zeroed full fields: the full-lattice
+  // derivative then picks up exactly the links Meooe uses between them.
+  void DhopDerivCB(GaugeField &mat,const FermionField &A,const FermionField &B,int dag)
+  {
+    GRID_ASSERT(A.Checkerboard() != B.Checkerboard());
+    FermionField Afull(_grid);
+    FermionField Bfull(_grid);
+    Afull = Zero();
+    Bfull = Zero();
+    setCheckerboard(Afull,A);
+    setCheckerboard(Bfull,B);
+    DhopDeriv(mat,Afull,Bfull,dag);
+  }
+  void DhopDerivOE(GaugeField &mat,const FermionField &U,const FermionField &V,int dag){ DhopDerivCB(mat,U,V,dag);};
+  void DhopDerivEO(GaugeField &mat,const FermionField &U,const FermionField &V,int dag){ DhopDerivCB(mat,U,V,dag);};
 
   ///////////////////////////////////////////////////////////////
   // non-hermitian hopping term; half cb or both
@@ -200,6 +259,7 @@ public:
     _grid(&Fgrid),
     _cbgrid(&Hgrid),
     _grid4(_Umu.Grid()),
+    _cbgrid4(nullptr),  // the checkerboard is 5D; there is no 4D red-black grid
     Umu(Nd,&Fgrid),
     UmuEven(Nd,&Hgrid),
     UmuOdd(Nd,&Hgrid),
@@ -221,10 +281,6 @@ public:
       MassField =scalar(-mass);
       one       =scalar(1.0);
       MassField =where(coor==Integer(Ls-1),MassField,one);
-      for(int mu=0;mu<Nd;mu++){
-	pickCheckerboard(Even,UmuEven[mu],Umu[mu]);
-	pickCheckerboard(Odd ,UmuOdd[mu],Umu[mu]);
-      }
       pickCheckerboard(Even,MassFieldEven,MassField);
       pickCheckerboard(Odd ,MassFieldOdd,MassField);
       
@@ -239,6 +295,8 @@ public:
       for(int s=0;s<this->Ls;s++){
 	InsertSlice(U4,Umu[mu],s,0);
       }
+      pickCheckerboard(Even,UmuEven[mu],Umu[mu]);
+      pickCheckerboard(Odd ,UmuOdd[mu],Umu[mu]);
     }
   }
 

Grid/qcd/action/fermion/DomainWallFermion.h

@@ -132,6 +132,35 @@ public:
     Approx::zolotarev_free(zdata);
   }
 
+  // Exact type only: derived operators must not inherit this
+  virtual SplitOperator<FermionField> *SplitClone(const Coordinate &mpi_split)
+  {
+    if ( typeid(*this) != typeid(DomainWallFermion<Impl>) ) {
+      return nullptr;
+    }
+    if ( this->Dirichlet || Impl::LsVectorised || Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field is overwritten below
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    DomainWallFermion<Impl> *clone = new DomainWallFermion<Impl>(Uplaceholder,
+                                                                 *split->FermionGrid,
+                                                                 *split->FermionRBGrid,
+                                                                 *split->GaugeGrid,
+                                                                 *split->GaugeRBGrid,
+                                                                 this->mass_plus,
+                                                                 this->M5,
+                                                                 this->Params);
+    this->CloneCoefficientsInto(*clone);
+    this->SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
 };
 
 NAMESPACE_END(Grid);

Grid/qcd/action/fermion/FermionOperator.h

@@ -59,6 +59,38 @@ public:
   virtual GridBase *GaugeGrid(void)           =0;
   virtual GridBase *GaugeRedBlackGrid(void)   =0;
 
+  ////////////////////////////////////////////////////////////////
+  // Split-communicator copies of this operator's gauge and fermion
+  // grids, returned in a new bundle with no Matrix or Linop yet.
+  // Partitions have MPI layout mpi_split. Collective.
+  ////////////////////////////////////////////////////////////////
+  SplitOperator<FermionField> *MakeSplitGrids(const Coordinate &mpi_split)
+  {
+    GridCartesian *U = dynamic_cast<GridCartesian *>(GaugeGrid());
+    GRID_ASSERT(U != nullptr);
+    GRID_ASSERT(mpi_split.size() == Nd);
+
+    double t0 = usecond();
+
+    SplitOperator<FermionField> *split = new SplitOperator<FermionField>();
+
+    split->GaugeGrid   = new GridCartesian(U->FullDimensions(),U->_simd_layout,mpi_split,*U);
+    split->GaugeRBGrid = SpaceTimeGrid::makeFourDimRedBlackGrid(split->GaugeGrid);
+    split->Partition   = GridSplitVectorIndex(U,split->GaugeGrid);
+    split->Partitions  = U->ProcessorCount()/split->GaugeGrid->ProcessorCount();
+
+    if ( FermionGrid()->Nd() == Nd+1 ) {
+      int Ls = FermionGrid()->_fdimensions[0];
+      split->FermionGrid   = SpaceTimeGrid::makeFiveDimGrid(Ls,split->GaugeGrid);
+      split->FermionRBGrid = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,split->GaugeGrid);
+    } else {
+      split->FermionGrid   = split->GaugeGrid;
+      split->FermionRBGrid = split->GaugeRBGrid;
+    }
+    std::cout << GridLogPerformance << "MakeSplitGrids: " << (usecond()-t0)/1.0e6 << " s" << std::endl;
+    return split;
+  }
+
   // override multiply
   virtual void  M    (const FermionField &in, FermionField &out)=0;
   virtual void  Mdag (const FermionField &in, FermionField &out)=0;
@@ -88,6 +120,10 @@ public:
   virtual void DhopDerivEO(GaugeField &mat,const FermionField &U,const FermionField &V,int dag)=0;
   virtual void DhopDerivOE(GaugeField &mat,const FermionField &U,const FermionField &V,int dag)=0;
 
+  // 1 if MoeDeriv/MeoDeriv return a force on a single 4D checkerboard (4D red-black),
+  // 0 if they return it on the full 4D gauge grid (a checkerboard that includes s)
+  virtual int CheckerboardedForce(void) { return 1; };
+
   virtual void  Mdiag  (const FermionField &in, FermionField &out) { Mooee(in,out);};   // Same as Mooee applied to both CB's
   virtual void  Mdir   (const FermionField &in, FermionField &out,int dir,int disp)=0;   // case by case Wilson, Clover, Cayley, ContFrac, PartFrac
   virtual void  MdirAll(const FermionField &in, std::vector<FermionField> &out)=0;   // case by case Wilson, Clover, Cayley, ContFrac, PartFrac

Grid/qcd/action/fermion/MobiusFermion.h

@@ -70,6 +70,37 @@ public:
  
   }
 
+  // Exact type only: derived operators (e.g. ScaledShamirFermion) must not inherit this
+  virtual SplitOperator<FermionField> *SplitClone(const Coordinate &mpi_split)
+  {
+    if ( typeid(*this) != typeid(MobiusFermion<Impl>) ) {
+      return nullptr;
+    }
+    if ( this->Dirichlet || Impl::LsVectorised || Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field is overwritten below
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    MobiusFermion<Impl> *clone = new MobiusFermion<Impl>(Uplaceholder,
+                                                         *split->FermionGrid,
+                                                         *split->FermionRBGrid,
+                                                         *split->GaugeGrid,
+                                                         *split->GaugeRBGrid,
+                                                         this->mass_plus,
+                                                         this->M5,
+                                                         this->_b,
+                                                         this->_c,
+                                                         this->Params);
+    this->CloneCoefficientsInto(*clone);
+    this->SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
 };
 
 NAMESPACE_END(Grid);

Grid/qcd/action/fermion/WilsonCloverFermion.h

@@ -91,6 +91,60 @@ public:
   // Derivative parts unpreconditioned pseudofermions
   void MDeriv(GaugeField &force, const FermionField &X, const FermionField &Y, int dag);
 
+  // Install already computed clover term and inverse, and derive the checkerboarded parts
+  void ImportCloverTerms(const CloverField &C, const CloverField &Cinv)
+  {
+    CloverTerm    = C;
+    CloverTermInv = Cinv;
+    pickCheckerboard(Even, CloverTermEven, CloverTerm);
+    pickCheckerboard(Odd,  CloverTermOdd,  CloverTerm);
+    pickCheckerboard(Even, CloverTermDagEven, adj(CloverTerm));
+    pickCheckerboard(Odd,  CloverTermDagOdd,  adj(CloverTerm));
+    pickCheckerboard(Even, CloverTermInvEven, CloverTermInv);
+    pickCheckerboard(Odd,  CloverTermInvOdd,  CloverTermInv);
+    pickCheckerboard(Even, CloverTermInvDagEven, adj(CloverTermInv));
+    pickCheckerboard(Odd,  CloverTermInvDagOdd,  adj(CloverTermInv));
+  }
+
+  // Exact type only: derived operators must not inherit this
+  virtual SplitOperator<FermionField> *SplitClone(const Coordinate &mpi_split)
+  {
+    if ( typeid(*this) != typeid(WilsonCloverFermion<Impl,CloverHelpers>) ) {
+      return nullptr;
+    }
+    if ( Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field and clover terms are overwritten below
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    WilsonCloverFermion<Impl,CloverHelpers> *clone =
+      new WilsonCloverFermion<Impl,CloverHelpers>(Uplaceholder,
+                                                  *split->GaugeGrid,
+                                                  *split->GaugeRBGrid,
+                                                  this->mass,
+                                                  csw_r,
+                                                  csw_t,
+                                                  this->anisotropyCoeff,
+                                                  this->Params);
+    // The stored coefficients are post-scaled; the constructor would scale them again
+    clone->csw_r = csw_r;
+    clone->csw_t = csw_t;
+
+    CloverField sC(split->GaugeGrid);
+    CloverField sCinv(split->GaugeGrid);
+    Grid_split(CloverTerm,sC);
+    Grid_split(CloverTermInv,sCinv);
+    clone->ImportCloverTerms(sC,sCinv);
+
+    this->SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
 public:
   // here fixing the 4 dimensions, make it more general?
 

Grid/qcd/action/fermion/WilsonCompressor.h

@@ -253,6 +253,11 @@ public:
 
     int dag = compress.dag;
     int face_idx=0;
+#ifdef GRID_FUSED_HALO_LOOPS
+    // queue the face gathers and run them together below (not on the first call, which builds the face tables)
+    const bool fused = this->face_table_computed;
+    if ( fused ) FusedHaloLoops::begin();
+#endif
 #define vet_same_node(a,b) \
       { auto tmp = b;  }
     if ( dag ) { 
@@ -274,6 +279,9 @@ public:
       vet_same_node(this->same_node[Zm],this->HaloGatherDir(source,ZpCompress,Zm,face_idx));
       vet_same_node(this->same_node[Tm],this->HaloGatherDir(source,TpCompress,Tm,face_idx));
     }
+#ifdef GRID_FUSED_HALO_LOOPS
+    if ( fused ) FusedHaloLoops::end();
+#endif
     this->face_table_computed=1;
     GRID_ASSERT(this->u_comm_offset==this->_unified_buffer_size);
     accelerator_barrier();

Grid/qcd/action/fermion/WilsonFermion.h

@@ -147,6 +147,48 @@ public:
   // DoubleStore impl dependent
   void ImportGauge(const GaugeField &_Umu);
 
+  // Install an already doubled gauge field (phases, twists, anisotropy and -1/2 applied)
+  void ImportDoubledGauge(const DoubledGaugeField &U)
+  {
+    Umu = U;
+    pickCheckerboard(Even,UmuEven,Umu);
+    pickCheckerboard(Odd ,UmuOdd ,Umu);
+  }
+
+  // Redistribute this operator's doubled gauge field onto clone's split grid. Collective.
+  void SplitDoubledGaugeInto(WilsonFermion<Impl> &clone)
+  {
+    DoubledGaugeField sUmu(clone.GaugeGrid());
+    Grid_split(Umu,sUmu);
+    clone.ImportDoubledGauge(sUmu);
+  }
+
+  // Exact type only: derived operators (e.g. WilsonTMFermion) must not inherit this
+  virtual SplitOperator<FermionField> *SplitClone(const Coordinate &mpi_split)
+  {
+    if ( typeid(*this) != typeid(WilsonFermion<Impl>) ) {
+      return nullptr;
+    }
+    if ( Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field is overwritten below
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    WilsonFermion<Impl> *clone = new WilsonFermion<Impl>(Uplaceholder,
+                                                         *split->GaugeGrid,
+                                                         *split->GaugeRBGrid,
+                                                         mass,
+                                                         this->Params,
+                                                         anisotropyCoeff);
+    SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
   ///////////////////////////////////////////////////////////////
   // Data members require to support the functionality
   ///////////////////////////////////////////////////////////////

Grid/qcd/action/fermion/WilsonFermion5D.h

@@ -181,6 +181,27 @@ public:
   // DoubleStore
   void ImportGauge(const GaugeField &_Umu);
     
+  // Install an already doubled gauge field (phases, twists and -1/2 applied)
+  void ImportDoubledGauge(const DoubledGaugeField &U)
+  {
+    Umu = U;
+    pickCheckerboard(Even,UmuEven,Umu);
+    pickCheckerboard(Odd ,UmuOdd ,Umu);
+  }
+
+  // Redistribute this operator's doubled gauge field onto clone's split grid. Collective.
+  void SplitDoubledGaugeInto(WilsonFermion5D<Impl> &clone)
+  {
+    double t0 = usecond();
+    DoubledGaugeField sUmu(clone.GaugeGrid());
+    Grid_split(Umu,sUmu);
+    double t1 = usecond();
+    clone.ImportDoubledGauge(sUmu);
+    double t2 = usecond();
+    std::cout << GridLogPerformance << "SplitDoubledGaugeInto: Grid_split " << (t1-t0)/1.0e6
+              << " s import " << (t2-t1)/1.0e6 << " s" << std::endl;
+  }
+
   ///////////////////////////////////////////////////////////////
   // Data members require to support the functionality
   ///////////////////////////////////////////////////////////////

Grid/qcd/action/fermion/WilsonKernels.h

@@ -42,6 +42,12 @@ public:
   enum { CommsAndCompute, CommsThenCompute };
   static int Opt;  
   static int Comms;
+#ifdef A64FX
+  // If set, called by thread 0 every few sites of the interior kernel so that
+  // communications started before it progress during the compute.
+  static void (*CommsProgress)(void *);
+  static void *CommsProgressArg;
+#endif
 };
  
 template<class Impl> class WilsonKernels : public FermionOperator<Impl> , public WilsonKernelsStatic { 

Grid/qcd/action/fermion/ZMobiusFermion.h

@@ -67,6 +67,44 @@ public:
     this->SetCoefficientsInternal(1.0,zgamma,b,c);
   }
 
+  // Exact type only: derived operators must not inherit this
+  virtual SplitOperator<FermionField> *SplitClone(const Coordinate &mpi_split)
+  {
+    if ( typeid(*this) != typeid(ZMobiusFermion<Impl>) ) {
+      return nullptr;
+    }
+    if ( this->Dirichlet || Impl::LsVectorised || Impl::isGparity ) {
+      return nullptr;
+    }
+    SplitOperator<FermionField> *split = this->MakeSplitGrids(mpi_split);
+
+    // Placeholder links; the doubled field is overwritten below
+    GaugeField Uplaceholder(split->GaugeGrid);
+    Uplaceholder = Zero();
+
+    // Constructor argument only; CloneCoefficientsInto restores _gamma at full precision
+    std::vector<ComplexD> gamma(this->Ls);
+    for ( int s=0; s<this->Ls; s++ ) {
+      gamma[s] = this->_gamma[s];
+    }
+
+    ZMobiusFermion<Impl> *clone = new ZMobiusFermion<Impl>(Uplaceholder,
+                                                           *split->FermionGrid,
+                                                           *split->FermionRBGrid,
+                                                           *split->GaugeGrid,
+                                                           *split->GaugeRBGrid,
+                                                           this->mass_plus,
+                                                           this->M5,
+                                                           gamma,
+                                                           this->_b,
+                                                           this->_c,
+                                                           this->Params);
+    this->CloneCoefficientsInto(*clone);
+    this->SplitDoubledGaugeInto(*clone);
+    split->Matrix = clone;
+    return split;
+  }
+
 };
 
 NAMESPACE_END(Grid);

Grid/qcd/action/fermion/implementation/WilsonFermion5DImplementation.h

@@ -346,6 +346,11 @@ void WilsonFermion5D<Impl>::DhopInternalOverlappedComms(StencilImpl & st,
   // do the compute interior
   /////////////////////////////
   int Opt = WilsonKernelsStatic::Opt; // Why pass this. Kernels should know
+#ifdef A64FX
+  // MPI does not progress outside MPI calls: let the interior kernel poll
+  WilsonKernelsStatic::CommsProgress    = [](void *s) { ((StencilImpl *)s)->CommunicateProgress(); };
+  WilsonKernelsStatic::CommsProgressArg = (void *)&st;
+#endif
   if (dag == DaggerYes) {
     GRID_TRACE("DhopDagInterior");
     Kernels::DhopDagKernel(Opt,st,U,st.CommBuf(),LLs,U.oSites(),in,out,1,0);
@@ -353,6 +358,10 @@ void WilsonFermion5D<Impl>::DhopInternalOverlappedComms(StencilImpl & st,
     GRID_TRACE("DhopInterior");
     Kernels::DhopKernel   (Opt,st,U,st.CommBuf(),LLs,U.oSites(),in,out,1,0);
   }
+#ifdef A64FX
+  WilsonKernelsStatic::CommsProgress    = nullptr;
+  WilsonKernelsStatic::CommsProgressArg = nullptr;
+#endif
   
   //ifdef GRID_ACCELERATED
 #if 0

Grid/qcd/action/fermion/implementation/WilsonFermionImplementation.h

@@ -436,6 +436,11 @@ void WilsonFermion<Impl>::DhopInternalOverlappedComms(StencilImpl &st,
   // do the compute interior
   /////////////////////////////
   int Opt = WilsonKernelsStatic::Opt;
+#ifdef A64FX
+  // MPI does not progress outside MPI calls: let the interior kernel poll
+  WilsonKernelsStatic::CommsProgress    = [](void *s) { ((StencilImpl *)s)->CommunicateProgress(); };
+  WilsonKernelsStatic::CommsProgressArg = (void *)&st;
+#endif
   if (dag == DaggerYes) {
     GRID_TRACE("DhopDagInterior");
     Kernels::DhopDagKernel(Opt,st,U,st.CommBuf(),1,U.oSites(),in,out,1,0);
@@ -443,6 +448,10 @@ void WilsonFermion<Impl>::DhopInternalOverlappedComms(StencilImpl &st,
     GRID_TRACE("DhopInterior");
     Kernels::DhopKernel(Opt,st,U,st.CommBuf(),1,U.oSites(),in,out,1,0);
   }
+#ifdef A64FX
+  WilsonKernelsStatic::CommsProgress    = nullptr;
+  WilsonKernelsStatic::CommsProgressArg = nullptr;
+#endif
 
   /////////////////////////////
   // Complete comms

Grid/qcd/action/fermion/implementation/WilsonKernelsAsmBodyA64FX.h

@@ -95,7 +95,6 @@ Author:  Nils Meyer  <nils.meyer@ur.de>  Regensburg University
 #ifdef INTERIOR_AND_EXTERIOR
 
 #define ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)			\
-      basep = st.GetPFInfo(nent,plocal); nent++;			\
       if ( local ) {							            \
     LOAD_CHIMU(base);                                       \
     LOAD_TABLE(PERMUTE_DIR);                                \
@@ -106,23 +105,15 @@ Author:  Nils Meyer  <nils.meyer@ur.de>  Regensburg University
       }									                    \
       base = st.GetInfo(ptype,local,perm,NxtDir,ent,plocal); ent++;	\
     MULT_2SPIN_1(Dir);					                    \
-    PREFETCH_CHIMU(base);                                   \
-    PREFETCH_CHIMU_L2(basep);                               \
-    /* PREFETCH_GAUGE_L1(NxtDir); */                        \
     MULT_2SPIN_2;					                        \
-    if (s == 0) {                                           \
-      if ((Dir == 0) || (Dir == 4)) { PREFETCH_GAUGE_L2(Dir); } \
-    }                                                       \
     RECON;								                    \
 
 /*
-NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
-    though I expected that it would improve on performance
+No software prefetch: the hardware prefetcher handles the neighbour and gauge link streams better on A64FX.
 */
 
 #define ASM_LEG_XP(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)	    \
   base = st.GetInfo(ptype,local,perm,Dir,ent,plocal); ent++; \
-  PREFETCH1_CHIMU(base);						            \
   ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)
 
 #define RESULT(base,basep) SAVE_RESULT(base,basep);
@@ -130,12 +121,11 @@ NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
 #endif
 
 ////////////////////////////////////////////////////////////////////////////////
-// Pre comms kernel -- prefetch like normal because it is mostly right
+// Pre comms kernel -- gauge link prefetch into L2 only
 ////////////////////////////////////////////////////////////////////////////////
 #ifdef INTERIOR
 
 #define ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)			\
-      basep = st.GetPFInfo(nent,plocal); nent++;			\
       if ( local ) {							\
   LOAD_CHIMU(base);                                       \
   LOAD_TABLE(PERMUTE_DIR);                                \
@@ -148,12 +138,15 @@ NB: picking PREFETCH_GAUGE_L2(Dir+4); here results in performance penalty
   RECON;								\
       }									\
   base = st.GetInfo(ptype,local,perm,NxtDir,ent,plocal); ent++;	\
-  PREFETCH_CHIMU(base);						\
-  PREFETCH_CHIMU_L2(basep);                               \
+  if (s == 0) {                                           \
+    if ((Dir == 0) || (Dir == 4)) { PREFETCH_GAUGE_L2(Dir); } \
+  }                                                       \
+
+/* No software prefetch of the neighbours (the hardware prefetcher handles them better on A64FX);
+   the gauge links of the next site are prefetched into L2. */
 
 #define ASM_LEG_XP(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)			\
   base = st.GetInfo(ptype,local,perm,Dir,ent,plocal); ent++;		\
-  PREFETCH1_CHIMU(base);						\
   { ZERO_PSI; }								\
   ASM_LEG(Dir,NxtDir,PERMUTE_DIR,PROJ,RECON)
 

Grid/qcd/action/fermion/implementation/WilsonKernelsImplementation.h

@@ -479,6 +479,49 @@ extern "C" {
     int sF = ss*Ls;							\
     WilsonKernels<Impl>::A(st_v,U_v,buf,sF,sU,Ls,1,in_v,out_v);		\
   });
+#ifdef A64FX
+// Interior kernel with communication progress: thread 0 calls the progress hook
+// between chunks of sites; dynamic scheduling lets the other threads pick up the
+// sites thread 0 has no time for while it drives the communications. Chunks of
+// 8 4d sites for 5d operators, 32 sites for 4d operators (measured on Fugaku).
+#define ASM_CALL_PROGRESS(A)						\
+  if ( WilsonKernelsStatic::CommsProgress ) {				\
+    auto progress     = WilsonKernelsStatic::CommsProgress;		\
+    auto progress_arg = WilsonKernelsStatic::CommsProgressArg;		\
+    const uint64_t chunk = std::max<uint64_t>(8,32/Ls);			\
+    const uint64_t nchunk = (Nsite + chunk - 1)/chunk;			\
+    thread_region {							\
+      const bool master = (thread_num() == 0);				\
+      DO_PRAGMA(omp for schedule(dynamic,1) nowait)			\
+      for(uint64_t c=0;c<nchunk;c++){					\
+	const uint64_t hi = std::min((c+1)*chunk,(uint64_t)Nsite);	\
+	for(uint64_t sss=c*chunk;sss<hi;sss++){				\
+	  int sU = sss;							\
+	  int sF = sss*Ls;						\
+	  WilsonKernels<Impl>::A(st_v,U_v,buf,sF,sU,Ls,1,in_v,out_v);	\
+	}								\
+	if ( master ) progress(progress_arg);				\
+      }									\
+    }									\
+  } else {								\
+    ASM_CALL(A);							\
+  }
+
+// Exterior kernel on the surface 4d sites only (the surface list holds Ls
+// consecutive 5d entries per surface 4d site)
+#define ASM_CALL_EXT(A)							\
+  const uint64_t nsurf = st.surface_list.size()/Ls;			\
+  auto surf = &st.surface_list[0];					\
+  thread_for( sss, nsurf, {						\
+    int sU = surf[sss*Ls]/Ls;						\
+    int sF = sU*Ls;							\
+    WilsonKernels<Impl>::A(st_v,U_v,buf,sF,sU,Ls,1,in_v,out_v);		\
+  });
+#else
+#define ASM_CALL_PROGRESS(A) ASM_CALL(A)
+#define ASM_CALL_EXT(A)      ASM_CALL(A)
+#endif
+
 #define ASM_CALL_SLICE(A)						\
   auto grid = in.Grid() ;						\
   int nt = grid->LocalDimensions()[4];					\
@@ -514,7 +557,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st,  DoubledGaugeField
      if (Opt == WilsonKernelsStatic::OptGeneric    ) { KERNEL_CALLNB(GenericDhopSiteInt); return;}
      if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALLNB(HandDhopSiteInt);    return;}
 #ifndef GRID_CUDA
-     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL(AsmDhopSiteInt);    return;}
+     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL_PROGRESS(AsmDhopSiteInt);    return;}
 #endif
    } else if( exterior ) {
      //     // dependent on result of merge
@@ -522,7 +565,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st,  DoubledGaugeField
      if (Opt == WilsonKernelsStatic::OptGeneric    ) { KERNEL_CALL_EXT(GenericDhopSiteExt); return;}
      if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALL_EXT(HandDhopSiteExt);    return;}
 #ifndef GRID_CUDA
-     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL(AsmDhopSiteExt);    return;}
+     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL_EXT(AsmDhopSiteExt);    return;}
 #endif
    }
    GRID_ASSERT(0 && " Kernel optimisation case not covered ");
@@ -560,7 +603,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st,  DoubledGaugeField
      if (Opt == WilsonKernelsStatic::OptGeneric    ) { KERNEL_CALLNB(GenericDhopSiteDagInt); return;}
      if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALLNB(HandDhopSiteDagInt);    return;}
 #ifndef GRID_CUDA
-     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL(AsmDhopSiteDagInt);     return;}
+     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL_PROGRESS(AsmDhopSiteDagInt);     return;}
 #endif
    } else if( exterior ) {
      // Dependent on result of merge
@@ -568,7 +611,7 @@ void WilsonKernels<Impl>::DhopKernel(int Opt,StencilImpl &st,  DoubledGaugeField
      if (Opt == WilsonKernelsStatic::OptGeneric    ) { KERNEL_CALL_EXT(GenericDhopSiteDagExt); return;}
      if (Opt == WilsonKernelsStatic::OptHandUnroll ) { KERNEL_CALL_EXT(HandDhopSiteDagExt);    return;}
 #ifndef GRID_CUDA
-     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL(AsmDhopSiteDagExt);     return;}
+     if (Opt == WilsonKernelsStatic::OptInlineAsm  ) {  ASM_CALL_EXT(AsmDhopSiteDagExt);     return;}
 #endif
    }
    GRID_ASSERT(0 && " Kernel optimisation case not covered ");

Grid/qcd/action/fermion/instantiation/WilsonKernelsInstantiation.cc

@@ -34,6 +34,10 @@ NAMESPACE_BEGIN(Grid);
 
 // Move these
 int WilsonKernelsStatic::Opt   = WilsonKernelsStatic::OptGeneric;
+#ifdef A64FX
+void (*WilsonKernelsStatic::CommsProgress)(void *) = nullptr;
+void *WilsonKernelsStatic::CommsProgressArg = nullptr;
+#endif
 int WilsonKernelsStatic::Comms = WilsonKernelsStatic::CommsAndCompute;
 
 NAMESPACE_END(Grid);

Grid/qcd/action/pseudofermion/EvenOddSchurDifferentiable.h

@@ -64,6 +64,20 @@ public:
     GRID_ASSERT(U.Checkerboard()==Odd);
     GRID_ASSERT(V.Checkerboard()==U.Checkerboard());
 
+    // Forces returned on the full 4D grid (checkerboard including s): same sequence, summed
+    if ( !this->_Mat.CheckerboardedForce() ) {
+      GaugeField ForceO(Force.Grid());
+      GaugeField ForceE(Force.Grid());
+      this->_Mat.Meooe   (V,tmp1);
+      this->_Mat.MooeeInv(tmp1,tmp2);
+      this->_Mat.MoeDeriv(ForceO,U,tmp2,DaggerNo);
+      this->_Mat.MeooeDag   (U,tmp1);
+      this->_Mat.MooeeInvDag(tmp1,tmp2);
+      this->_Mat.MeoDeriv(ForceE,tmp2,V,DaggerNo);
+      Force = -(ForceO+ForceE);
+      return;
+    }
+
     // NOTE Guido: WE DO NOT WANT TO USE THE ucbgrid GRID FOR THE FORCE
     // it is not conformable with the HMC force field
     // Case: Ls vectorised fields
@@ -109,6 +123,20 @@ public:
     GRID_ASSERT(V.Checkerboard()==Odd);
     GRID_ASSERT(V.Checkerboard()==V.Checkerboard());
 
+    // Forces returned on the full 4D grid (checkerboard including s): same sequence, summed
+    if ( !this->_Mat.CheckerboardedForce() ) {
+      GaugeField ForceO(Force.Grid());
+      GaugeField ForceE(Force.Grid());
+      this->_Mat.MeooeDag   (V,tmp1);
+      this->_Mat.MooeeInvDag(tmp1,tmp2);
+      this->_Mat.MoeDeriv(ForceO,U,tmp2,DaggerYes);
+      this->_Mat.Meooe   (U,tmp1);
+      this->_Mat.MooeeInv(tmp1,tmp2);
+      this->_Mat.MeoDeriv(ForceE,tmp2,V,DaggerYes);
+      Force = -(ForceO+ForceE);
+      return;
+    }
+
     // NOTE Guido: WE DO NOT WANT TO USE THE ucbgrid GRID FOR THE FORCE
     // it is not conformable with the HMC force field
     // INHERIT FROM THE Force field instead

Grid/qcd/smearing/GaugeConfiguration.h

@@ -125,6 +125,16 @@ public:
     return SmearedSet[Level];
   }
 
+  /*! @brief Level n of the smearing chain: 0 is the thin links, n is n smearing steps */
+  const GaugeField &get_level(int n) const
+  {
+    GRID_ASSERT(n >= 0 && n <= (int)smearingLevels);
+    if ( n == 0 ) {
+      return *ThinLinks;
+    }
+    return SmearedSet[n-1];
+  }
+
   //====================================================================
   void set_iLambda(GaugeLinkField& iLambda, GaugeLinkField& e_iQ,
                    const GaugeLinkField& iQ, const GaugeLinkField& Sigmap,
@@ -277,6 +287,45 @@ public:
     SigmaTilde=Gimpl::projectForce(SigmaTilde); // Ta
       
   }
+
+  //====================================================================
+  // Force on the thin links from an action that depends on every level of the chain.
+  // LevelForce[n], n = 0..smearingLevels, is the force on level n in the same
+  // "U times dS/dU" form smeared_force takes (zero fields allowed). The chain rule is
+  // applied from the top level down, adding each level's own force on the way:
+  //   Sigma_L = U_L^dag F_L
+  //   Sigma_n = AnalyticSmearedForce(Sigma_{n+1}, U_n) + U_n^dag F_n
+  //   ThinForce = Ta( U_0 Sigma_0 )
+  // With force on the top level only this is smeared_force.
+  // SmearedConfigurationMasked would need its own: its chain rule is level dependent.
+  virtual void smeared_force_levels(const std::vector<GaugeField> &LevelForce, GaugeField &ThinForce)
+  {
+    GRID_ASSERT(LevelForce.size() == smearingLevels+1);
+
+    GridBase *grid = LevelForce[0].Grid();
+    GaugeField Sigma(grid);
+    GaugeLinkField tmp_mu(grid);
+
+    // Sigma = U_L^dag F_L
+    for (int mu = 0; mu < Nd; mu++) {
+      tmp_mu = adj(peekLorentz(get_level(smearingLevels), mu)) * peekLorentz(LevelForce[smearingLevels], mu);
+      pokeLorentz(Sigma, tmp_mu, mu);
+    }
+
+    for (int n = (int)smearingLevels - 1; n >= 0; n--) {
+      Sigma = AnalyticSmearedForce(Sigma, get_level(n));
+      for (int mu = 0; mu < Nd; mu++) {
+        tmp_mu = peekLorentz(Sigma, mu) + adj(peekLorentz(get_level(n), mu)) * peekLorentz(LevelForce[n], mu);
+        pokeLorentz(Sigma, tmp_mu, mu);
+      }
+    }
+
+    for (int mu = 0; mu < Nd; mu++) {
+      tmp_mu = peekLorentz(*ThinLinks, mu) * peekLorentz(Sigma, mu);
+      pokeLorentz(ThinForce, tmp_mu, mu);
+    }
+    ThinForce = Gimpl::projectForce(ThinForce); // Ta
+  }
   //====================================================================
 
   virtual GaugeField& get_SmearedU() { return SmearedSet[smearingLevels - 1]; }

Grid/simd/Fujitsu_A64FX_intrin_double.h

@@ -39,8 +39,13 @@ Author: Nils Meyer <nils.meyer@ur.de>
 #define UNLOCK_GAUGE(A)  
 #define MASK_REGS                      DECLARATIONS_A64FXd  
 #define SAVE_RESULT(A,B)               RESULT_A64FXd(A);  
-#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_A64FXd(Dir)  
-#define MULT_2SPIN_2                   MULT_2SPIN_2_A64FXd  
+#ifdef GRID_A64FX_FCMLA
+#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_FCMLA_A64FXd(Dir)
+#define MULT_2SPIN_2                   MULT_2SPIN_2_FCMLA_A64FXd
+#else
+#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_FMLA_A64FXd(Dir)
+#define MULT_2SPIN_2                   MULT_2SPIN_2_FMLA_A64FXd
+#endif
 #define LOAD_CHI(base)                 LOAD_CHI_A64FXd(base)  
 #define ZERO_PSI                       ZERO_PSI_A64FXd  
 #define ADD_RESULT(base,basep)         LOAD_CHIMU(base); ADD_RESULT_INTERNAL_A64FXd; RESULT_A64FXd(base)  
@@ -273,8 +278,88 @@ Author: Nils Meyer <nils.meyer@ur.de>
     U_11 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-2));  \
     U_21 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(1));  \
 }
+// MULT_2SPIN with FMLA instead of FCMLA (3 uops on A64FX):
+//   UChi = sum_b Re(U_ab) Chi_b + Im(U_ab) (i Chi_b)
+// column by column. Re(U)/Im(U) are duplicated over each complex with trn1/trn2,
+// i Chi_b is formed in place with fcadd once Chi_b has been used (Chi is dead after
+// the multiply). U_00, U_10, U_20 hold column b, U_01 is a temporary.
+// Define GRID_A64FX_FCMLA to use the original FCMLA version.
+#define MULT_2SPIN_1_FMLA_A64FXd(A)  \
+{ \
+    const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
+    U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-6));  \
+    U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-3));  \
+    U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(0));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_10 = svmul_x(pg1, U_01, Chi_10); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_11 = svmul_x(pg1, U_01, Chi_10); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_12 = svmul_x(pg1, U_01, Chi_10); \
+    Chi_00 = svcadd_x(pg1, zero0, Chi_00, 90); \
+    Chi_10 = svcadd_x(pg1, zero0, Chi_10, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_00); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_10); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_00); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_10); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_00); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_10); \
+    U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-5));  \
+    U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-2));  \
+    U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(1));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
+    Chi_01 = svcadd_x(pg1, zero0, Chi_01, 90); \
+    Chi_11 = svcadd_x(pg1, zero0, Chi_11, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
+    U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-4));  \
+    U_10 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-1));  \
+    U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(2));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
+    Chi_02 = svcadd_x(pg1, zero0, Chi_02, 90); \
+    Chi_12 = svcadd_x(pg1, zero0, Chi_12, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
+}
+#define MULT_2SPIN_2_FMLA_A64FXd
 // MULT_2SPIN
-#define MULT_2SPIN_1_A64FXd(A)  \
+#define MULT_2SPIN_1_FCMLA_A64FXd(A)  \
 { \
     const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
     U_00 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(-6));  \
@@ -300,7 +385,7 @@ Author: Nils Meyer <nils.meyer@ur.de>
     U_20 = svld1_vnum(pg1, (float64_t*)(baseU + 2 * 3 * 64), (int64_t)(2));  \
 }
 // MULT_2SPIN_BACKEND
-#define MULT_2SPIN_2_A64FXd  \
+#define MULT_2SPIN_2_FCMLA_A64FXd  \
 { \
     UChi_00 = svcmla_x(pg1, UChi_00, U_01, Chi_01, 0); \
     UChi_10 = svcmla_x(pg1, UChi_10, U_01, Chi_11, 0); \

Grid/simd/Fujitsu_A64FX_intrin_single.h

@@ -39,8 +39,13 @@ Author: Nils Meyer <nils.meyer@ur.de>
 #define UNLOCK_GAUGE(A)  
 #define MASK_REGS                      DECLARATIONS_A64FXf  
 #define SAVE_RESULT(A,B)               RESULT_A64FXf(A);  
-#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_A64FXf(Dir)  
-#define MULT_2SPIN_2                   MULT_2SPIN_2_A64FXf  
+#ifdef GRID_A64FX_FCMLA
+#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_FCMLA_A64FXf(Dir)
+#define MULT_2SPIN_2                   MULT_2SPIN_2_FCMLA_A64FXf
+#else
+#define MULT_2SPIN_1(Dir)              MULT_2SPIN_1_FMLA_A64FXf(Dir)
+#define MULT_2SPIN_2                   MULT_2SPIN_2_FMLA_A64FXf
+#endif
 #define LOAD_CHI(base)                 LOAD_CHI_A64FXf(base)  
 #define ZERO_PSI                       ZERO_PSI_A64FXf  
 #define ADD_RESULT(base,basep)         LOAD_CHIMU(base); ADD_RESULT_INTERNAL_A64FXf; RESULT_A64FXf(base)  
@@ -273,8 +278,88 @@ Author: Nils Meyer <nils.meyer@ur.de>
     U_11 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-2));  \
     U_21 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(1));  \
 }
+// MULT_2SPIN with FMLA instead of FCMLA (3 uops on A64FX):
+//   UChi = sum_b Re(U_ab) Chi_b + Im(U_ab) (i Chi_b)
+// column by column. Re(U)/Im(U) are duplicated over each complex with trn1/trn2,
+// i Chi_b is formed in place with fcadd once Chi_b has been used (Chi is dead after
+// the multiply). U_00, U_10, U_20 hold column b, U_01 is a temporary.
+// Define GRID_A64FX_FCMLA to use the original FCMLA version.
+#define MULT_2SPIN_1_FMLA_A64FXf(A)  \
+{ \
+    const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
+    U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-6));  \
+    U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-3));  \
+    U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(0));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_10 = svmul_x(pg1, U_01, Chi_10); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_11 = svmul_x(pg1, U_01, Chi_10); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmul_x(pg1, U_01, Chi_00); \
+    UChi_12 = svmul_x(pg1, U_01, Chi_10); \
+    Chi_00 = svcadd_x(pg1, zero0, Chi_00, 90); \
+    Chi_10 = svcadd_x(pg1, zero0, Chi_10, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_00); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_10); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_00); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_10); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_00); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_10); \
+    U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-5));  \
+    U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-2));  \
+    U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(1));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
+    Chi_01 = svcadd_x(pg1, zero0, Chi_01, 90); \
+    Chi_11 = svcadd_x(pg1, zero0, Chi_11, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_01); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_11); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_01); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_11); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_01); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_11); \
+    U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-4));  \
+    U_10 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-1));  \
+    U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(2));  \
+    U_01 = svtrn1(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
+    U_01 = svtrn1(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
+    U_01 = svtrn1(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
+    Chi_02 = svcadd_x(pg1, zero0, Chi_02, 90); \
+    Chi_12 = svcadd_x(pg1, zero0, Chi_12, 90); \
+    U_01 = svtrn2(U_00, U_00); \
+    UChi_00 = svmla_x(pg1, UChi_00, U_01, Chi_02); \
+    UChi_10 = svmla_x(pg1, UChi_10, U_01, Chi_12); \
+    U_01 = svtrn2(U_10, U_10); \
+    UChi_01 = svmla_x(pg1, UChi_01, U_01, Chi_02); \
+    UChi_11 = svmla_x(pg1, UChi_11, U_01, Chi_12); \
+    U_01 = svtrn2(U_20, U_20); \
+    UChi_02 = svmla_x(pg1, UChi_02, U_01, Chi_02); \
+    UChi_12 = svmla_x(pg1, UChi_12, U_01, Chi_12); \
+}
+#define MULT_2SPIN_2_FMLA_A64FXf
 // MULT_2SPIN
-#define MULT_2SPIN_1_A64FXf(A)  \
+#define MULT_2SPIN_1_FCMLA_A64FXf(A)  \
 { \
     const auto & ref(U[sU](A)); baseU = (uint64_t)&ref; \
     U_00 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(-6));  \
@@ -300,7 +385,7 @@ Author: Nils Meyer <nils.meyer@ur.de>
     U_20 = svld1_vnum(pg1, (float32_t*)(baseU + 2 * 3 * 64), (int64_t)(2));  \
 }
 // MULT_2SPIN_BACKEND
-#define MULT_2SPIN_2_A64FXf  \
+#define MULT_2SPIN_2_FCMLA_A64FXf  \
 { \
     UChi_00 = svcmla_x(pg1, UChi_00, U_01, Chi_01, 0); \
     UChi_10 = svcmla_x(pg1, UChi_10, U_01, Chi_11, 0); \

Grid/stencil/SimpleCompressor.h

@@ -11,6 +11,43 @@ public:
 };
 
 
+// On A64FX (OpenMP CPU builds) the face gathers of one halo exchange, and the merges, run in one parallel
+// region: each loop is queued, then every thread takes its static share of each loop, with one barrier at
+// the end instead of one parallel region per face. The gathers also prefetch the source sites 4 entries
+// ahead. CPU only: the queued loops run after the views of the source field are closed, which is safe
+// only because a CPU view is a plain pointer to host memory.
+#if defined(A64FX) && defined(_OPENMP) && !defined(GRID_SIMT)
+#define GRID_FUSED_HALO_LOOPS
+struct FusedHaloLoops {
+  // a queued loop over sites [0,n): body(lo,hi) runs sites lo..hi-1
+  typedef std::pair<uint64_t, std::function<void(uint64_t,uint64_t)> > Loop;
+  static std::vector<Loop> &queue(void) { static std::vector<Loop> q; return q; }
+  static bool &deferring(void) { static bool d = false; return d; }
+  // between begin() and end(), the face loops are queued (push) instead of run; end() runs them all
+  static void begin(void) { deferring() = true; }
+  static void end(void)   { deferring() = false; run(); }
+  static void push(uint64_t n,std::function<void(uint64_t,uint64_t)> body) { queue().push_back({n,body}); }
+  static const int prefetch_distance = 4;
+  template<class T> static inline void prefetch(const T *p) {
+    for(int o=0;o<(int)sizeof(T);o+=256) __builtin_prefetch((const char *)p+o,0,3); // 256-byte lines
+  }
+  static void run(void) {
+    auto &q = queue();
+    if ( q.empty() ) return;
+    // one chunk per thread, as in the lattice reductions (Lattice_reduction.h)
+    const int nthread = GridThread::GetThreads();
+    thread_for(thr,nthread,{
+      for(auto &loop : q) {
+	int mywork, myoff;
+	GridThread::GetWork((int)loop.first,thr,mywork,myoff);
+	if ( mywork > 0 ) loop.second(myoff,myoff+mywork);
+      }
+    });
+    q.clear();
+  }
+};
+#endif
+
 // Compressors will inherit buffer management policies
 // Standard comms buffer management
 class FaceGatherSimple
@@ -29,6 +66,20 @@ public:
     std::pair<int,int> *table_v = & table[0];
     
     auto rhs_v = rhs.View(AcceleratorRead);
+#ifdef GRID_FUSED_HALO_LOOPS
+    if ( FusedHaloLoops::deferring() ) {
+      const vobj *rp = &rhs_v[0];
+      const int pf = FusedHaloLoops::prefetch_distance;
+      FusedHaloLoops::push((uint64_t)num, [=](uint64_t lo,uint64_t hi) {
+	for(uint64_t i=lo;i<hi;i++){
+	  if ( i+pf<hi ) FusedHaloLoops::prefetch(&rp[so+table_v[i+pf].second]);
+	  compress.Compress(buffer[off+table_v[i].first],rp[so+table_v[i].second]);
+	}
+      });
+      rhs_v.ViewClose();
+      return;
+    }
+#endif
     accelerator_forNB( i,num, vobj::Nsimd(), {
 	compress.Compress(buffer[off+table_v[i].first],rhs_v[so+table_v[i].second]);
     });
@@ -48,6 +99,23 @@ public:
     auto p1=&pointers[1][0];
     auto tp=&table[0];
     auto rhs_p = &rhs_v[0];
+#ifdef GRID_FUSED_HALO_LOOPS
+    if ( FusedHaloLoops::deferring() ) {
+      const vobj *rp = rhs_p;
+      const int pf = FusedHaloLoops::prefetch_distance;
+      FusedHaloLoops::push((uint64_t)num, [=](uint64_t lo,uint64_t hi) {
+	for(uint64_t j=lo;j<hi;j++){
+	  if ( j+pf<hi ) {
+	    FusedHaloLoops::prefetch(&rp[so+tp[2*(j+pf)].second]);
+	    FusedHaloLoops::prefetch(&rp[so+tp[2*(j+pf)+1].second]);
+	  }
+	  compress.CompressExchange(p0[j],p1[j],rp[so+tp[2*j].second],rp[so+tp[2*j+1].second],type);
+	}
+      });
+      rhs_v.ViewClose();
+      return;
+    }
+#endif
     accelerator_forNB(j, num, vobj::Nsimd(), {
 	compress.CompressExchange(p0[j],p1[j],
 				  rhs_p[so+tp[2*j  ].second],
@@ -73,6 +141,14 @@ public:
     auto vp0= &mm.vpointers[0][0];
     auto vp1= &mm.vpointers[1][0];
     auto type= mm.type;
+#ifdef GRID_FUSED_HALO_LOOPS
+    if ( FusedHaloLoops::deferring() ) {
+      FusedHaloLoops::push((uint64_t)(mm.buffer_size/2), [=](uint64_t lo,uint64_t hi) {
+	for(uint64_t o=lo;o<hi;o++) decompress.Exchange(mp[2*o],mp[2*o+1],vp0[o],vp1[o],type);
+      });
+      return;
+    }
+#endif
     accelerator_forNB(o,mm.buffer_size/2,Merger::Nsimd,{
 	decompress.Exchange(mp[2*o],mp[2*o+1],vp0[o],vp1[o],type);
     });

Grid/stencil/Stencil.h

@@ -630,6 +630,13 @@ public:
     }
   }
 
+#ifdef A64FX
+  // Progress outstanding communications while computing (see CommunicateBegin)
+  void CommunicateProgress(void)
+  {
+    _grid->StencilSendToRecvFromProgress(MpiReqs);
+  }
+#endif
   void CommunicateComplete(std::vector<std::vector<CommsRequest_t> > &reqs)
   {
     //    std::cout << "Communicate Complete "<<std::endl;
@@ -879,6 +886,15 @@ public:
   template<class decompressor>
   void CommsMerge(decompressor decompress,std::vector<Merge> &mm,std::vector<Decompress> &dd)
   {
+#ifdef GRID_FUSED_HALO_LOOPS
+    if ( dd.size()==0 ) { // merges only: run them in one parallel region (SimpleCompressor.h)
+      FusedHaloLoops::begin();
+      for(int i=0;i<mm.size();i++) decompressor::MergeFace(decompress,mm[i]);
+      FusedHaloLoops::end();
+      acceleratorFenceComputeStream();
+      return;
+    }
+#endif
     for(int i=0;i<mm.size();i++){
       decompressor::MergeFace(decompress,mm[i]);
     }

Grid/util/Init.cc

@@ -95,6 +95,8 @@ NAMESPACE_BEGIN(Grid);
 //////////////////////////////////////////////////////
 static Coordinate Grid_default_latt;
 static Coordinate Grid_default_mpi;
+static Coordinate Grid_default_batched_solver_split;
+static bool       Grid_default_batched_solver_split_node = false;
 
 int GridThread::_threads =1;
 int GridThread::_hyperthreads=1;
@@ -108,6 +110,8 @@ char *GridHostname(void)
 }
 const Coordinate &GridDefaultLatt(void)     {return Grid_default_latt;};
 const Coordinate &GridDefaultMpi(void)      {return Grid_default_mpi;};
+const Coordinate &GridDefaultBatchedSolverSplit(void)     {return Grid_default_batched_solver_split;};
+bool              GridDefaultBatchedSolverSplitNode(void) {return Grid_default_batched_solver_split_node;};
 const Coordinate GridDefaultSimd(int dims,int nsimd)
 {
   Coordinate layout(dims);
@@ -363,17 +367,14 @@ void Grid_init(int *argc,char ***argv)
     GlobalSharedMemory::HPEhypercube = enable;
   }
 
+  if( GridCmdOptionExists(*argv,*argv+*argc,"--tofu") ){
+    GlobalSharedMemory::TofuPlacement = 1;
+  }
+
   if( GridCmdOptionExists(*argv,*argv+*argc,"--shm-hugepages") ){
     GlobalSharedMemory::Hugepages = 1;
   }
 
-#if defined(A64FX)
-  if( GridCmdOptionExists(*argv,*argv+*argc,"--comms-overlap") ){
-    std::cout << "Option --comms-overlap currently not supported on QPACE4. Exiting." << std::endl;
-    exit(EXIT_FAILURE);
-  }
-#endif
-
   //////////////////////////////////////////////////////////
   // Memory manager
   //////////////////////////////////////////////////////////
@@ -537,6 +538,7 @@ void Grid_init(int *argc,char ***argv)
     std::cout<<GridLogMessage<<"Geometry:"<<std::endl;
     std::cout<<GridLogMessage<<std::endl;
     std::cout<<GridLogMessage<<"  --mpi n.n.n.n   : default MPI decomposition"<<std::endl;    
+    std::cout<<GridLogMessage<<"  --tofu          : place ranks on the job's Tofu torus (Fujitsu MPI only)"<<std::endl;
     std::cout<<GridLogMessage<<"  --threads n     : default number of OMP threads"<<std::endl;
     std::cout<<GridLogMessage<<"  --grid n.n.n.n  : default Grid size"<<std::endl;
     std::cout<<GridLogMessage<<"  --shm  M        : allocate M megabytes of shared memory for comms"<<std::endl;
@@ -562,6 +564,7 @@ void Grid_init(int *argc,char ***argv)
     std::cout<<GridLogMessage<<"Performance:"<<std::endl;
     std::cout<<GridLogMessage<<std::endl;
     std::cout<<GridLogMessage<<"  --comms-overlap    : Overlap comms with compute "<<std::endl;    
+    std::cout<<GridLogMessage<<"  --batched-solver-split a.b.c.d|node : batched solvers run independent solves on partitions of this MPI layout (node = one per node)"<<std::endl;
     std::cout<<GridLogMessage<<std::endl;
     std::cout<<GridLogMessage<<"  --dslash-generic: Wilson kernel for generic Nc"<<std::endl;    
     std::cout<<GridLogMessage<<"  --dslash-unroll : Wilson kernel for Nc=3"<<std::endl;    
@@ -590,6 +593,9 @@ void Grid_init(int *argc,char ***argv)
   if( GridCmdOptionExists(*argv,*argv+*argc,"--comms-overlap") ){
     WilsonKernelsStatic::Comms = WilsonKernelsStatic::CommsAndCompute;
     StaggeredKernelsStatic::Comms = StaggeredKernelsStatic::CommsAndCompute;
+#ifdef GRID_FJMPI
+    Stencil_persistent = true;
+#endif
   } else {
     WilsonKernelsStatic::Comms = WilsonKernelsStatic::CommsThenCompute;
     StaggeredKernelsStatic::Comms = StaggeredKernelsStatic::CommsThenCompute;
@@ -613,6 +619,16 @@ void Grid_init(int *argc,char ***argv)
 		  Grid_default_latt,
 		  Grid_default_mpi);
 
+  if( GridCmdOptionExists(*argv,*argv+*argc,"--batched-solver-split") ){
+    std::string arg = GridCmdOptionPayload(*argv,*argv+*argc,"--batched-solver-split");
+    if ( arg == "node" ) {
+      Grid_default_batched_solver_split_node = true;
+    } else {
+      GridCmdOptionIntVector(arg,Grid_default_batched_solver_split);
+      GRID_ASSERT(Grid_default_batched_solver_split.size() == Grid_default_mpi.size());
+    }
+  }
+
   if( GridCmdOptionExists(*argv,*argv+*argc,"--decomposition") ){
     std::cout<<GridLogMessage<<"Grid Default Decomposition patterns\n";
     std::cout<<GridLogMessage<<"\tOpenMP threads : "<<GridThread::GetThreads()<<std::endl;

Grid/util/Init.h

@@ -49,6 +49,11 @@ void Grid_unquiesce_nodes(void);
 const Coordinate  GridDefaultSimd(int dims,int nsimd);
 const Coordinate &GridDefaultLatt(void);
 const Coordinate &GridDefaultMpi(void);
+// --batched-solver-split a.b.c.d : MPI layout of each partition for batched solves.
+// --batched-solver-split node    : one partition per node.
+// Not given: empty Coordinate and false, meaning no split.
+const Coordinate &GridDefaultBatchedSolverSplit(void);
+bool              GridDefaultBatchedSolverSplitNode(void);
 const int        &GridThreads(void)  ;
 void              GridSetThreads(int t) ;
 void GridLogTimestamp(int);

configure.ac

@@ -846,6 +846,33 @@ esac
 AM_CONDITIONAL(BUILD_COMMS_MPI3,  [ test "${comms_type}X" == "mpi3X" ] )
 AM_CONDITIONAL(BUILD_COMMS_NONE,  [ test "${comms_type}X" == "noneX" ])
 
+############### Fujitsu MPI extensions (persistent halo requests, Tofu topology)
+AC_ARG_ENABLE([fjmpi],[AS_HELP_STRING([--enable-fjmpi=auto|yes|no],[Use the Fujitsu MPI extensions (default: auto, if mpi-ext.h declares them)])],[ac_FJMPI=${enable_fjmpi}],[ac_FJMPI=auto])
+fjmpi=no
+case ${ac_FJMPI} in
+     yes|no|auto) ;;
+     *) AC_MSG_ERROR([${ac_FJMPI} unsupported --enable-fjmpi option]);;
+esac
+if test "${comms_type}X" == "mpi3X" && test "${ac_FJMPI}X" != "noX"; then
+  ac_save_CXXFLAGS=$CXXFLAGS
+  CXXFLAGS="$MPI_CXXFLAGS $CXXFLAGS"
+  AC_MSG_CHECKING([for the Fujitsu MPI extensions])
+  AC_COMPILE_IFELSE([AC_LANG_PROGRAM([[#include <mpi.h>
+#include <mpi-ext.h>]],
+    [[MPI_Request r; int c[3];
+      FJMPI_Prequest_start(&r);
+      FJMPI_Topology_get_coords(MPI_COMM_WORLD,0,FJMPI_LOGICAL,3,c);]])],
+    [fjmpi=yes],[fjmpi=no])
+  AC_MSG_RESULT([${fjmpi}])
+  CXXFLAGS=$ac_save_CXXFLAGS
+fi
+if test "${ac_FJMPI}X" == "yesX" && test "${fjmpi}X" != "yesX"; then
+  AC_MSG_ERROR([--enable-fjmpi=yes but the Fujitsu MPI extensions were not found (needs --enable-comms=mpi and Fujitsu MPI)])
+fi
+if test "${fjmpi}X" == "yesX"; then
+  AC_DEFINE([GRID_FJMPI],[1],[Fujitsu MPI extensions])
+fi
+
 ############### RNG selection
 AC_ARG_ENABLE([rng],[AS_HELP_STRING([--enable-rng=ranlux48|mt19937|sitmo],[\
 	            Select Random Number Generator to be used])],\
@@ -955,6 +982,7 @@ Threading                   : ${ac_openmp}
 Acceleration                : ${ac_ACCELERATOR}
 Unified virtual memory      : ${ac_UNIFIED}
 Communications type         : ${comms_type}
+Fujitsu MPI extensions      : ${fjmpi}
 Shared memory allocator     : ${ac_SHM}
 Shared memory mmap path     : ${ac_SHMPATH}
 Default precision           : ${ac_PRECISION}

systems/Frontier/split_batched_mobius.job

@@ -0,0 +1,119 @@
+#!/bin/bash -l
+#SBATCH --job-name=split-batched-mobius
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=8
+#SBATCH --cpus-per-task=7
+#SBATCH --gpus-per-node=8
+#SBATCH --time=1:30:00
+#SBATCH --account=phy157_dwf
+#SBATCH --gpu-bind=none
+#SBATCH --exclusive
+#SBATCH --mem=0
+#SBATCH -S 0
+##SBATCH -q debug
+
+##############################################################################
+# MixedPrecisionConjugateGradientBatched with split inner solves, at production
+# size: Mobius 64^3x128, Ls=12, mass 0.026, b=1.5 c=0.5 M5=1.8, antiperiodic
+# in time, Hadrons default SchurDiagMooeeOperator.
+#
+# 4 nodes = 32 GCDs, --mpi 2.2.2.4 (local 32^4).  --batched-solver-split node
+# makes each node one partition (P=4), so Nbatch=4 is one round with no
+# zero padding.  The driver solves the same batch twice from zero guesses:
+#   UNSPLIT  inner CG on all 32 GCDs, one rhs after another (existing path)
+#   SPLIT    four independent inner CGs, one per node
+# and prints per-rhs iterations, true residuals and wall clock for each.
+#
+# Build (in the Frontier build tree, after ./scripts/filelist in the source
+# root so the new test is in tests/solver/Make.inc):
+#   make -C tests/solver Test_split_mobius_batched
+#
+# Memory, per GCD (evictable Lattice fields, so bounded by --device-mem):
+#   double RB 5d field 1.21 GB, float 0.60 GB
+#   driver src+sol 9.7, solver batch copies 9.7, split partition fields and
+#   inner CG on the 4x larger split local volume ~20, gauge copies ~5
+#   => ~45 GB, hence --device-mem 40000 as in pvdagm_mixed_precision.job.
+# Host memory: the device is an inclusive cache, so host RSS per rank is the whole
+# ~45 GB working set (plus transient Grid_split staging): ~400 GB of a node's 512 GB.
+# The driver prints MEMORY lines (host RSS, peak, allocator caches) at each phase.
+# Grid_split stages 2 x 4.8 GB per rank on the host per transfer.
+##############################################################################
+
+LUSTRE=/lustre/orion/phy157/proj-shared/phy157_dwf/paboyle
+RUNDIR=$LUSTRE/runs/${SLURM_JOB_NAME}_${SLURM_JOB_ID}
+mkdir -p $RUNDIR && cd $RUNDIR && echo "RUNDIR $RUNDIR"
+
+cat << EOF > select_gpu
+#!/bin/bash
+export GPU_MAP=(0 1 2 3 7 6 5 4)
+export NUMA_MAP=(3 3 1 1 2 2 0 0)
+export GPU=\${GPU_MAP[\$SLURM_LOCALID]}
+export NUMA=\${NUMA_MAP[\$SLURM_LOCALID]}
+export HIP_VISIBLE_DEVICES=\$GPU
+unset ROCR_VISIBLE_DEVICES
+if [ \$SLURM_PROCID = "0" ]; then echo \$*; fi
+exec numactl -m \$NUMA -N \$NUMA \$*
+EOF
+chmod +x ./select_gpu
+
+# Source tree and build of THIS branch (feature/splitGridCGBatch) on Frontier
+root=$LUSTRE/SplitCG/Grid/systems/Frontier
+source $root/sourceme-rocm7.2.sh
+
+export OMP_NUM_THREADS=7
+ulimit -c 0                             # no 22 GB GPU core dumps
+export FI_MR_CACHE_MONITOR=kdreg2       # site default; device-buffer MPI on Slingshot
+export MPICH_GPU_SUPPORT_ENABLED=1
+export MPICH_SMP_SINGLE_COPY_MODE=CMA
+export MPICH_OFI_NIC_POLICY=GPU
+# No GRID_ALLOC_NCACHE_LARGE here. With --enable-unified=no every Lattice lives in host
+# memory (the device is an inclusive cache), so a deep large-allocation cache keeps freed
+# host blocks on top of the full working set: measured on the laptop at 64 entries it
+# held about as much again as the working set, enough to OOM a 4-node run of this job.
+module load libfabric
+
+BIN=$root/tests/solver/Test_split_mobius_batched
+OPTS="--accelerator-threads 8 --shm 4096 --shm-mpi 1 --device-mem 40000 --comms-overlap"
+vol=64.64.64.128
+MPI_GEOM=2.2.2.4
+
+# A thermalised 64^3x128 NERSC configuration.  Leave empty for a hot start:
+# fine for plumbing, but mass 0.026 on a hot field is not a physical
+# iteration count and may run long.
+CONFIG=
+
+PHYS="--Ls 12 --mass 0.026 --M5 1.8 --b 1.5 --c 0.5 --nbatch 4 --tol 1e-8"
+if [ -n "$CONFIG" ]; then PHYS="$PHYS --config $CONFIG"; fi
+
+export GRID_STDOUT_ROOT=$RUNDIR/split
+srun -N4 -n32 --kill-on-bad-exit=1 ./select_gpu $BIN --mpi $MPI_GEOM --grid $vol $OPTS \
+     $PHYS --batched-solver-split node \
+     --debug-stdout --log Error,Warning,Message,Performance \
+     > log.split 2>&1
+echo "exit $?"
+
+##############################################################################
+# Readout.  Device OOM appears only on stderr as "hipMalloc failed".
+##############################################################################
+f=$(grep -l "Memory access fault\|NO_TRANSLATION\|hipMalloc failed\|out of memory\|illegal memory access\|GRID_ASSERT" $GRID_STDOUT_ROOT/*/Grid.stderr.* 2>/dev/null | head -1)
+if [ -n "$f" ]; then
+  echo "FAULT in $f"; tail -20 $f | cut -c1-160
+fi
+grep -h "MEMORY\|BatchedSolverSplit\|Inner CG iterations\|Total time\|true residual\|SUMMARY" \
+     $GRID_STDOUT_ROOT/0/Grid.stdout.0 2>/dev/null | cut -c1-200
+
+##############################################################################
+# What to read.
+#
+# BatchedSolverSplit line: the partition layout must equal the node layout
+# (ShmGrid), e.g. [1 2 2 2] of [2 2 2 4] : 4 partitions.  A "straddle node
+# boundaries" warning means the rank reordering did not give node-aligned
+# blocks; the result is still correct but the inner CG talks off node.
+#
+# UNSPLIT vs SPLIT: per-rhs inner iterations should agree closely (the same
+# solves, different reduction order); the SUMMARY lines give the wall clock.
+# The gain is the strong-scaling loss of the 32-GCD inner CG against four
+# 8-GCD CGs with node-local halos.  "Split setup and transfer" in the Total
+# time line prices the clone (grids, gauge split) and the Grid_split /
+# Grid_unsplit traffic, which is host-staged.
+##############################################################################

tests/forces/ForceTest.h

@@ -0,0 +1,138 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./tests/forces/ForceTest.h
+
+    Copyright (C) 2022
+
+Author: Peter Boyle <pboyle@bnl.gov>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+    *************************************************************************************/
+    /*  END LEGAL */
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Finite-difference check of an action's force: refresh, S1 = S(U), step the links by eps
+// along a random momentum P, take the force at the midpoint, step again, S2 = S(U'').
+// The force predicts S2 - S1 to O(eps^3).
+//
+// Runs on a copy of U, so the caller's field is unchanged. Returns S2 - S1 - dSpred for
+// the caller to assert on; halving eps should reduce it by about 8.
+/////////////////////////////////////////////////////////////////////////////////////////////
+template<class Gimpl>
+RealD ForceTest(Action<LatticeGaugeField> &action,LatticeGaugeField & Uin,MomentumFilterBase<LatticeGaugeField> &Filter,RealD eps=0.005)
+{
+  GridBase *UGrid = Uin.Grid();
+
+  std::vector<int> seeds({1,2,3,5});
+  GridSerialRNG            sRNG;         sRNG.SeedFixedIntegers(seeds);
+  GridParallelRNG          RNG4(UGrid);  RNG4.SeedFixedIntegers(seeds);
+
+  LatticeColourMatrix Pmu(UGrid);
+  LatticeGaugeField P(UGrid);
+  LatticeGaugeField UdSdU(UGrid);
+  LatticeGaugeField U(UGrid);
+  U = Uin;
+
+  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
+  std::cout << GridLogMessage << " Force test for "<<action.action_name()<<" eps "<<eps<<std::endl;
+  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
+
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  std::cout << GridLogMessage << " Refresh "<<action.action_name()<<std::endl;
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+
+  Gimpl::generate_momenta(P,sRNG,RNG4);
+  Filter.applyFilter(P);
+
+  action.refresh(U,sRNG,RNG4);
+
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+
+  RealD S1 = action.S(U);
+
+  Gimpl::update_field(P,U,eps);
+
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  std::cout << GridLogMessage << " Derivative "<<action.action_name()<<std::endl;
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  action.deriv(U,UdSdU);
+  UdSdU = Ta(UdSdU);
+  Filter.applyFilter(UdSdU);
+
+  DumpSliceNorm("Force",UdSdU,Nd-1);
+
+  Gimpl::update_field(P,U,eps);
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
+  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+
+  RealD S2 = action.S(U);
+
+  // Use the derivative
+  LatticeComplex dS(UGrid); dS = Zero();
+  for(int mu=0;mu<Nd;mu++){
+    auto UdSdUmu = PeekIndex<LorentzIndex>(UdSdU,mu);
+    Pmu= PeekIndex<LorentzIndex>(P,mu);
+    dS = dS - trace(Pmu*UdSdUmu)*eps*2.0*2.0;
+  }
+  ComplexD dSpred    = sum(dS);
+  RealD diff =  S2-S1-dSpred.real();
+
+  std::cout<< GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
+  std::cout<< GridLogMessage << "S1 : "<< S1    <<std::endl;
+  std::cout<< GridLogMessage << "S2 : "<< S2    <<std::endl;
+  std::cout<< GridLogMessage << "dS : "<< S2-S1 <<std::endl;
+  std::cout<< GridLogMessage << "dSpred : "<< dSpred.real() <<std::endl;
+  std::cout<< GridLogMessage << "diff : "<< diff<<std::endl;
+  std::cout<< GridLogMessage << "*********************************************************"<<std::endl;
+  std::cout<< GridLogMessage << "Done" <<std::endl;
+  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
+  return diff;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// ForceTest at eps0, eps0/2, ... (neps values). A correct force leaves an O(eps^3)
+// discrepancy, so successive ratios approach 8 as eps shrinks; an error in the force adds
+// an O(eps) part and drives them towards 2. A strongly curved action needs a smaller eps0
+// before the ratios settle. Returns the ratio for the smallest pair.
+/////////////////////////////////////////////////////////////////////////////////////////////
+template<class Gimpl>
+RealD ForceTestScaling(Action<LatticeGaugeField> &action,LatticeGaugeField &U,MomentumFilterBase<LatticeGaugeField> &Filter,const std::string &name,
+                       RealD eps0=0.01,int neps=4)
+{
+  std::vector<RealD> eps(neps);
+  std::vector<RealD> d(neps);
+  for(int i=0;i<neps;i++){
+    eps[i] = eps0/(1<<i);
+    d[i]   = ForceTest<Gimpl>(action,U,Filter,eps[i]);
+  }
+  for(int i=0;i<neps;i++){
+    std::cout << GridLogMessage << name << ": eps " << eps[i] << " discrepancy " << d[i];
+    if ( i > 0 ) {
+      std::cout << "  ratio to previous " << d[i-1]/d[i];
+    }
+    std::cout << std::endl;
+  }
+  RealD scale = d[neps-2]/d[neps-1];
+  std::cout << GridLogMessage << name << ": smallest-eps ratio " << scale << " (expect ~8)" << std::endl;
+  return scale;
+}

tests/forces/Test_bdy.cc

@@ -99,91 +99,12 @@ Author: Peter Boyle <pboyle@bnl.gov>
 using namespace std;
 using namespace Grid;
 
+#include "ForceTest.h"
+
 typedef MobiusFermionD FermionAction;
 typedef WilsonImplD FimplD;
 typedef WilsonImplD FermionImplPolicy;
 
-template<class Gimpl>
-void ForceTest(Action<LatticeGaugeField> &action,LatticeGaugeField & U,MomentumFilterBase<LatticeGaugeField> &Filter)
-{
-  GridBase *UGrid = U.Grid();
-
-  std::vector<int> seeds({1,2,3,5});
-  GridSerialRNG            sRNG;         sRNG.SeedFixedIntegers(seeds);
-  GridParallelRNG          RNG4(UGrid);  RNG4.SeedFixedIntegers(seeds);
-
-  LatticeColourMatrix Pmu(UGrid); 
-  LatticeGaugeField P(UGrid); 
-  LatticeGaugeField UdSdU(UGrid); 
-
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  std::cout << GridLogMessage << " Force test for "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  
-  RealD eps=0.005;
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Refresh "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  
-  Gimpl::generate_momenta(P,sRNG,RNG4);
-  Filter.applyFilter(P);
-
-#if  0
-  FieldMetaData header;
-  std::string file("./ckpoint_lat.2000");
-  NerscIO::readConfiguration(U,header,file);
-#else
-  U = 1.0;
-#endif
-  action.refresh(U,sRNG,RNG4);
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-
-  RealD S1 = action.S(U);
-
-  Gimpl::update_field(P,U,eps);
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Derivative "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  action.deriv(U,UdSdU);
-  UdSdU = Ta(UdSdU);
-  Filter.applyFilter(UdSdU);
-
-  DumpSliceNorm("Force",UdSdU,Nd-1);
-  
-  Gimpl::update_field(P,U,eps);
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  
-  RealD S2 = action.S(U);
-
-  // Use the derivative
-  LatticeComplex dS(UGrid); dS = Zero();
-  for(int mu=0;mu<Nd;mu++){
-    auto UdSdUmu = PeekIndex<LorentzIndex>(UdSdU,mu);
-    Pmu= PeekIndex<LorentzIndex>(P,mu);
-    dS = dS - trace(Pmu*UdSdUmu)*eps*2.0*2.0;
-  }
-  ComplexD dSpred    = sum(dS);
-  RealD diff =  S2-S1-dSpred.real();
-
-  std::cout<< GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout<< GridLogMessage << "S1 : "<< S1    <<std::endl;
-  std::cout<< GridLogMessage << "S2 : "<< S2    <<std::endl;
-  std::cout<< GridLogMessage << "dS : "<< S2-S1 <<std::endl;
-  std::cout<< GridLogMessage << "dSpred : "<< dSpred.real() <<std::endl;
-  std::cout<< GridLogMessage << "diff : "<< diff<<std::endl;
-  std::cout<< GridLogMessage << "*********************************************************"<<std::endl;
-  //  GRID_ASSERT(diff<1.0);
-  std::cout<< GridLogMessage << "Done" <<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-}
-
 int main (int argc, char ** argv)
 {
   Grid_init(&argc,&argv);
@@ -230,6 +151,7 @@ int main (int argc, char ** argv)
 
   ///////////////////// Gauge Field and Gauge Forces ////////////////////////////
   LatticeGaugeField U(UGrid);
+  U = 1.0;
 
   RealD beta=6.0;
   WilsonGaugeActionR PlaqAction(beta);

tests/forces/Test_double_ratio.cc

@@ -303,83 +303,7 @@ typedef MobiusFermionD FermionAction;
 typedef WilsonImplD FimplD;
 typedef WilsonImplD FermionImplPolicy;
 
-template<class Gimpl>
-void ForceTest(Action<LatticeGaugeField> &action,LatticeGaugeField & U,MomentumFilterBase<LatticeGaugeField> &Filter)
-{
-  GridBase *UGrid = U.Grid();
-
-  std::vector<int> seeds({1,2,3,5});
-  GridSerialRNG            sRNG;         sRNG.SeedFixedIntegers(seeds);
-  GridParallelRNG          RNG4(UGrid);  RNG4.SeedFixedIntegers(seeds);
-
-  LatticeColourMatrix Pmu(UGrid); 
-  LatticeGaugeField P(UGrid); 
-  LatticeGaugeField UdSdU(UGrid); 
-
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  std::cout << GridLogMessage << " Force test for "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  
-  RealD eps=0.005;
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Refresh "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  
-  Gimpl::generate_momenta(P,sRNG,RNG4);
-  Filter.applyFilter(P);
-
-  FieldMetaData header;
-  std::string file("./ckpoint_lat.2000");
-  NerscIO::readConfiguration(U,header,file);
- 
-  action.refresh(U,sRNG,RNG4);
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-
-  RealD S1 = action.S(U);
-
-  Gimpl::update_field(P,U,eps);
-
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Derivative "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  action.deriv(U,UdSdU);
-  UdSdU = Ta(UdSdU);
-  Filter.applyFilter(UdSdU);
-
-  DumpSliceNorm("Force",UdSdU,Nd-1);
-  
-  Gimpl::update_field(P,U,eps);
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout << GridLogMessage << " Action "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  
-  RealD S2 = action.S(U);
-
-  // Use the derivative
-  LatticeComplex dS(UGrid); dS = Zero();
-  for(int mu=0;mu<Nd;mu++){
-    auto UdSdUmu = PeekIndex<LorentzIndex>(UdSdU,mu);
-    Pmu= PeekIndex<LorentzIndex>(P,mu);
-    dS = dS - trace(Pmu*UdSdUmu)*eps*2.0*2.0;
-  }
-  ComplexD dSpred    = sum(dS);
-  RealD diff =  S2-S1-dSpred.real();
-
-  std::cout<< GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout<< GridLogMessage << "S1 : "<< S1    <<std::endl;
-  std::cout<< GridLogMessage << "S2 : "<< S2    <<std::endl;
-  std::cout<< GridLogMessage << "dS : "<< S2-S1 <<std::endl;
-  std::cout<< GridLogMessage << "dSpred : "<< dSpred.real() <<std::endl;
-  std::cout<< GridLogMessage << "diff : "<< diff<<std::endl;
-  std::cout<< GridLogMessage << "*********************************************************"<<std::endl;
-  //  GRID_ASSERT(diff<1.0);
-  std::cout<< GridLogMessage << "Done" <<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-}
+#include "ForceTest.h"
 
 int main (int argc, char ** argv)
 {
@@ -427,6 +351,9 @@ int main (int argc, char ** argv)
 
   ///////////////////// Gauge Field and Gauge Forces ////////////////////////////
   LatticeGaugeField U(UGrid);
+  FieldMetaData header;
+  std::string file("./ckpoint_lat.2000");
+  NerscIO::readConfiguration(U,header,file);
 
   RealD beta=6.0;
   WilsonGaugeActionR PlaqAction(beta);

tests/forces/Test_dwf_ratio_4dpf_force.cc

@@ -35,7 +35,7 @@ Author: Peter Boyle <pboyle@bnl.gov>
 //     composition identity [P M^-1 V Pdag][P V^-1 M Pdag] = 1 for the chosen
 //     (P,Pdag) wall pair.  NO PREDICTION is made about which convention
 //     passes -- that is what the test decides.
-//  T2 ForceTest (idiom from Test_double_ratio.cc) : midpoint-derivative
+//  T2 ForceTest (tests/forces/ForceTest.h) : midpoint-derivative
 //     check of deriv against S.  Should PASS for BOTH conventions (S and
 //     deriv use the same literal-adjoint pair by construction).
 //  T3 Trivial-ratio control (V == M) : T1 with NumOp = DenOp.  The solve
@@ -124,69 +124,7 @@ RealD HeatbathIdentityTest(TwoFlavourRatio4DPseudoFermionAction<Impl> &action,
   return defect;
 }
 
-////////////////////////////////////////////////////////////////////
-// T2 : ForceTest idiom from Test_double_ratio.cc (midpoint derivative)
-////////////////////////////////////////////////////////////////////
-template<class Gimpl>
-void ForceTest(Action<LatticeGaugeField> &action,LatticeGaugeField & U,MomentumFilterBase<LatticeGaugeField> &Filter)
-{
-  GridBase *UGrid = U.Grid();
-
-  std::vector<int> seeds({1,2,3,5});
-  GridSerialRNG            sRNG;         sRNG.SeedFixedIntegers(seeds);
-  GridParallelRNG          RNG4(UGrid);  RNG4.SeedFixedIntegers(seeds);
-
-  LatticeColourMatrix Pmu(UGrid);
-  LatticeGaugeField P(UGrid);
-  LatticeGaugeField UdSdU(UGrid);
-
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  std::cout << GridLogMessage << " Force test for "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-
-  RealD eps=0.005;
-
-  Gimpl::generate_momenta(P,sRNG,RNG4);
-  Filter.applyFilter(P);
-
-  action.refresh(U,sRNG,RNG4);
-
-  RealD S1 = action.S(U);
-
-  Gimpl::update_field(P,U,eps);
-
-  action.deriv(U,UdSdU);
-  UdSdU = Ta(UdSdU);
-  Filter.applyFilter(UdSdU);
-
-  DumpSliceNorm("Force",UdSdU,Nd-1);
-
-  Gimpl::update_field(P,U,eps);
-
-  RealD S2 = action.S(U);
-
-  // Use the derivative
-  LatticeComplex dS(UGrid); dS = Zero();
-  for(int mu=0;mu<Nd;mu++){
-    auto UdSdUmu = PeekIndex<LorentzIndex>(UdSdU,mu);
-    Pmu= PeekIndex<LorentzIndex>(P,mu);
-    dS = dS - trace(Pmu*UdSdUmu)*eps*2.0*2.0;
-  }
-  ComplexD dSpred    = sum(dS);
-  RealD diff =  S2-S1-dSpred.real();
-
-  std::cout<< GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout<< GridLogMessage << "S1 : "<< S1    <<std::endl;
-  std::cout<< GridLogMessage << "S2 : "<< S2    <<std::endl;
-  std::cout<< GridLogMessage << "dS : "<< S2-S1 <<std::endl;
-  std::cout<< GridLogMessage << "dSpred : "<< dSpred.real() <<std::endl;
-  std::cout<< GridLogMessage << "diff : "<< diff<<std::endl;
-  std::cout<< GridLogMessage << "diff/dS : "<< diff/(S2-S1)<<std::endl;
-  std::cout<< GridLogMessage << "*********************************************************"<<std::endl;
-  //  GRID_ASSERT(diff<1.0);
-  std::cout<< GridLogMessage << "Done" <<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-}
+#include "ForceTest.h"
 
 int main (int argc, char ** argv)
 {

tests/forces/Test_dwf_ratio_leftprec.cc

@@ -38,7 +38,7 @@ Author: Peter Boyle <pboyle@bnl.gov>
 //             BOTH classes (E0a also validates the twin-eta capture).
 //   E1      : S_classic == S_leftprec       (relative, ~1e-8)
 //   E2      : deriv_classic == deriv_leftprec (pointwise field norm, ~1e-8)
-//   F1      : ForceTest (Test_double_ratio.cc idiom) on the LeftPrec class.
+//   F1      : ForceTest (tests/forces/ForceTest.h) on the LeftPrec class.
 //
 // All asserts are hard: this is the regression gate for the new class.
 // Run small, e.g.:  ./Test_dwf_ratio_leftprec --grid 8.8.8.8
@@ -138,67 +138,7 @@ public:
   }
 };
 
-////////////////////////////////////////////////////////////////////
-// ForceTest idiom from Test_double_ratio.cc (midpoint derivative)
-////////////////////////////////////////////////////////////////////
-template<class Gimpl>
-void ForceTest(Action<LatticeGaugeField> &action,LatticeGaugeField & U,MomentumFilterBase<LatticeGaugeField> &Filter)
-{
-  GridBase *UGrid = U.Grid();
-
-  std::vector<int> seeds({1,2,3,5});
-  GridSerialRNG            sRNG;         sRNG.SeedFixedIntegers(seeds);
-  GridParallelRNG          RNG4(UGrid);  RNG4.SeedFixedIntegers(seeds);
-
-  LatticeColourMatrix Pmu(UGrid);
-  LatticeGaugeField P(UGrid);
-  LatticeGaugeField UdSdU(UGrid);
-
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-  std::cout << GridLogMessage << " Force test for "<<action.action_name()<<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-
-  RealD eps=0.005;
-
-  Gimpl::generate_momenta(P,sRNG,RNG4);
-  Filter.applyFilter(P);
-
-  action.refresh(U,sRNG,RNG4);
-
-  RealD S1 = action.S(U);
-
-  Gimpl::update_field(P,U,eps);
-
-  action.deriv(U,UdSdU);
-  UdSdU = Ta(UdSdU);
-  Filter.applyFilter(UdSdU);
-
-  DumpSliceNorm("Force",UdSdU,Nd-1);
-
-  Gimpl::update_field(P,U,eps);
-
-  RealD S2 = action.S(U);
-
-  LatticeComplex dS(UGrid); dS = Zero();
-  for(int mu=0;mu<Nd;mu++){
-    auto UdSdUmu = PeekIndex<LorentzIndex>(UdSdU,mu);
-    Pmu= PeekIndex<LorentzIndex>(P,mu);
-    dS = dS - trace(Pmu*UdSdUmu)*eps*2.0*2.0;
-  }
-  ComplexD dSpred    = sum(dS);
-  RealD diff =  S2-S1-dSpred.real();
-
-  std::cout<< GridLogMessage << "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++"<<std::endl;
-  std::cout<< GridLogMessage << "S1 : "<< S1    <<std::endl;
-  std::cout<< GridLogMessage << "S2 : "<< S2    <<std::endl;
-  std::cout<< GridLogMessage << "dS : "<< S2-S1 <<std::endl;
-  std::cout<< GridLogMessage << "dSpred : "<< dSpred.real() <<std::endl;
-  std::cout<< GridLogMessage << "diff : "<< diff<<std::endl;
-  std::cout<< GridLogMessage << "diff/dS : "<< diff/(S2-S1)<<std::endl;
-  std::cout<< GridLogMessage << "*********************************************************"<<std::endl;
-  std::cout<< GridLogMessage << "Done" <<std::endl;
-  std::cout << GridLogMessage << "*********************************************************"<<std::endl;
-}
+#include "ForceTest.h"
 
 int main (int argc, char ** argv)
 {

tests/solver/Test_dwf_mrhs_cg.cc

@@ -56,14 +56,12 @@ int main (int argc, char ** argv)
   /////////////////////////////////////////////
   // Split into 1^4 mpi communicators
   /////////////////////////////////////////////
-  for(int i=0;i<argc;i++){
-    if(std::string(argv[i]) == "--split"){
-      for(int k=0;k<mpi_layout.size();k++){
-	std::stringstream ss; 
-	ss << argv[i+1+k]; 
-	ss >> mpi_split[k];
-      }
-      break;
+  if ( GridCmdOptionExists(argv,argv+argc,"--split") ) {
+    std::vector<int> split;
+    GridCmdOptionIntVector(GridCmdOptionPayload(argv,argv+argc,"--split"),split);
+    GRID_ASSERT(split.size() == mpi_layout.size());
+    for(int k=0;k<mpi_layout.size();k++){
+      mpi_split[k] = split[k];
     }
   }
 

tests/solver/Test_dwf_mrhs_cg_mpi.cc

@@ -59,14 +59,12 @@ int main (int argc, char ** argv)
   // Split into 1^4 mpi communicators
   /////////////////////////////////////////////
 
-  for(int i=0;i<argc;i++){
-    if(std::string(argv[i]) == "--split"){
-      for(int k=0;k<mpi_layout.size();k++){
-	std::stringstream ss; 
-	ss << argv[i+1+k]; 
-	ss >> mpi_split[k];
-      }
-      break;
+  if ( GridCmdOptionExists(argv,argv+argc,"--split") ) {
+    std::vector<int> split;
+    GridCmdOptionIntVector(GridCmdOptionPayload(argv,argv+argc,"--split"),split);
+    GRID_ASSERT(split.size() == mpi_layout.size());
+    for(int k=0;k<mpi_layout.size();k++){
+      mpi_split[k] = split[k];
     }
   }
 

tests/solver/Test_mobius_bcg.cc

@@ -63,14 +63,12 @@ int main (int argc, char ** argv)
   // Split into 1^4 mpi communicators
   /////////////////////////////////////////////
 
-  for(int i=0;i<argc;i++){
-    if(std::string(argv[i]) == "--split"){
-      for(int k=0;k<mpi_layout.size();k++){
-	std::stringstream ss; 
-	ss << argv[i+1+k]; 
-	ss >> mpi_split[k];
-      }
-      break;
+  if ( GridCmdOptionExists(argv,argv+argc,"--split") ) {
+    std::vector<int> split;
+    GridCmdOptionIntVector(GridCmdOptionPayload(argv,argv+argc,"--split"),split);
+    GRID_ASSERT(split.size() == mpi_layout.size());
+    for(int k=0;k<mpi_layout.size();k++){
+      mpi_split[k] = split[k];
     }
   }
 

tests/solver/Test_split_grid.cc

@@ -56,14 +56,12 @@ int main (int argc, char ** argv)
   // Split into 1^4 mpi communicators
   /////////////////////////////////////////////
 
-  for(int i=0;i<argc;i++){
-    if(std::string(argv[i]) == "--split"){
-      for(int k=0;k<mpi_layout.size();k++){
-	std::stringstream ss; 
-	ss << argv[i+1+k]; 
-	ss >> mpi_split[k];
-      }
-      break;
+  if ( GridCmdOptionExists(argv,argv+argc,"--split") ) {
+    std::vector<int> split;
+    GridCmdOptionIntVector(GridCmdOptionPayload(argv,argv+argc,"--split"),split);
+    GRID_ASSERT(split.size() == mpi_layout.size());
+    for(int k=0;k<mpi_layout.size();k++){
+      mpi_split[k] = split[k];
     }
   }
 

tests/solver/Test_split_mixedprec_batched.cc

@@ -0,0 +1,355 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./tests/solver/Test_split_mixedprec_batched.cc
+
+    Copyright (C) 2026
+
+Author: Peter Boyle <paboyle@ph.ed.ac.uk>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+*************************************************************************************/
+/*  END LEGAL */
+/////////////////////////////////////////////////////////////////////////////////////////////
+// MixedPrecisionConjugateGradientBatched with inner solves on split-communicator
+// partitions (--batched-solver-split), against the same solver unsplit.
+//
+//   mpirun -n 2 ./Test_split_mixedprec_batched --grid 8.8.8.8 --mpi 1.1.1.2 --batched-solver-split 1.1.1.1
+//
+// With no --batched-solver-split the partitions are single ranks (1.1.1.1).
+// NBATCH is deliberately not a multiple of the partition count, so the last group is
+// zero-padded.
+/////////////////////////////////////////////////////////////////////////////////////////////
+#include <Grid/Grid.h>
+
+using namespace std;
+using namespace Grid;
+
+const int   NBATCH    = 3;
+const RealD TOLERANCE = 1.0e-8;
+
+template<class Field>
+RealD RelativeDifference(const Field &a,const Field &b)
+{
+  Field diff(a.Grid());
+  diff = a - b;
+  return std::sqrt(norm2(diff)/norm2(b));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// The solver assumes vector p of a group lands in the partition whose
+// GridSplitVectorIndex is p: fill vector p with the constant p+1, split, and check
+// every partition sees its own index.
+/////////////////////////////////////////////////////////////////////////////////////////////
+void CheckPartitionOrder(GridCartesian *UGrid,const Coordinate &layout)
+{
+  GridCartesian SGrid(UGrid->FullDimensions(),UGrid->_simd_layout,layout,*UGrid);
+  int P     = UGrid->ProcessorCount()/SGrid.ProcessorCount();
+  int index = GridSplitVectorIndex(UGrid,&SGrid);
+
+  std::vector<LatticeComplexF> full(P,UGrid);
+  for(int p=0;p<P;p++){
+    full[p] = ComplexF(p+1,0.0);
+  }
+  LatticeComplexF split(&SGrid);
+  Grid_split(full,split);
+
+  LatticeComplexF expect(&SGrid);
+  expect = ComplexF(index+1,0.0);
+  RealD err = norm2(split - expect);
+
+  std::cout << GridLogMessage << "Partition order: vector index " << index << " error " << err << std::endl;
+  GRID_ASSERT(err == 0.0);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+// For one pair of double/float linear operators:
+//  (1) red-black-aware split/unsplit round trip is exact;
+//  (2) the split clone's HermOp agrees with the full-grid HermOp;
+//  (3) the batched solver with split inner solves agrees with the unsplit solver;
+//  (4) the split solution has true residual at the requested tolerance.
+/////////////////////////////////////////////////////////////////////////////////////////////
+template<class FieldD,class FieldF>
+void CheckSplitBatched(const std::string &name,
+                       LinearOperatorBase<FieldD> &Linop_d,
+                       LinearOperatorBase<FieldF> &Linop_f,
+                       GridBase *grid_f,
+                       std::vector<FieldD> &src,
+                       const Coordinate &layout)
+{
+  std::cout << GridLogMessage << "==================================================" << std::endl;
+  std::cout << GridLogMessage << name << std::endl;
+  std::cout << GridLogMessage << "==================================================" << std::endl;
+
+  int cb = src[0].Checkerboard();
+
+  SplitOperator<FieldF> *split = Linop_f.SplitClone(layout);
+  GRID_ASSERT(split != nullptr);
+  int P = split->Partitions;
+
+  std::vector<FieldF> full(P,grid_f);
+  std::vector<FieldF> back(P,grid_f);
+  std::vector<FieldF> Mfull(P,grid_f);
+  for(int p=0;p<P;p++){
+    full[p].Checkerboard()  = cb;
+    back[p].Checkerboard()  = cb;
+    Mfull[p].Checkerboard() = cb;
+    precisionChange(full[p],src[p%src.size()]);
+  }
+
+  // (1) round trip
+  FieldF s_in(split->FieldGrid);
+  FieldF s_out(split->FieldGrid);
+  Grid_split(full,s_in);
+  Grid_unsplit(back,s_in);
+  for(int p=0;p<P;p++){
+    RealD err = norm2(back[p] - full[p]);
+    std::cout << GridLogMessage << name << ": split/unsplit round trip " << p << " error " << err << std::endl;
+    GRID_ASSERT(err == 0.0);
+  }
+
+  // (2) operator
+  split->Linop->HermOp(s_in,s_out);
+  Grid_unsplit(back,s_out);
+  for(int p=0;p<P;p++){
+    Linop_f.HermOp(full[p],Mfull[p]);
+    RealD err = RelativeDifference(back[p],Mfull[p]);
+    std::cout << GridLogMessage << name << ": split HermOp vs full HermOp " << p << " relative difference " << err << std::endl;
+    GRID_ASSERT(err < 1.0e-6);
+  }
+  delete split;
+
+  // (3) solver, unsplit then split
+  int NBatch = src.size();
+  std::vector<FieldD> sol_ref(NBatch,src[0].Grid());
+  std::vector<FieldD> sol_split(NBatch,src[0].Grid());
+  for(int i=0;i<NBatch;i++){
+    sol_ref[i].Checkerboard()   = cb;
+    sol_split[i].Checkerboard() = cb;
+    sol_ref[i]   = Zero();
+    sol_split[i] = Zero();
+  }
+
+  {
+    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;
+  for(int i=0;i<NBatch;i++){
+    RealD diff = RelativeDifference(sol_split[i],sol_ref[i]);
+    Linop_d.HermOp(sol_split[i],Msol);
+    RealD resid = RelativeDifference(Msol,src[i]);
+    std::cout << GridLogMessage << name << ": rhs " << i
+              << " split vs unsplit " << diff
+              << " true residual " << resid << std::endl;
+    GRID_ASSERT(diff  < 1.0e-6);
+    GRID_ASSERT(resid < 10.0*TOLERANCE);
+  }
+}
+
+int main (int argc, char ** argv)
+{
+  Grid_init(&argc,&argv);
+
+  const int Ls = 8;
+
+  Coordinate layout = GridDefaultBatchedSolverSplit();
+  if ( layout.size() == 0 ) {
+    layout = Coordinate(Nd,1);
+  }
+
+  GridCartesian         *UGrid_d   = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(), GridDefaultSimd(Nd,vComplexD::Nsimd()), GridDefaultMpi());
+  GridRedBlackCartesian *UrbGrid_d = SpaceTimeGrid::makeFourDimRedBlackGrid(UGrid_d);
+  GridCartesian         *FGrid_d   = SpaceTimeGrid::makeFiveDimGrid(Ls,UGrid_d);
+  GridRedBlackCartesian *FrbGrid_d = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,UGrid_d);
+
+  GridCartesian         *UGrid_f   = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(), GridDefaultSimd(Nd,vComplexF::Nsimd()), GridDefaultMpi());
+  GridRedBlackCartesian *UrbGrid_f = SpaceTimeGrid::makeFourDimRedBlackGrid(UGrid_f);
+  GridCartesian         *FGrid_f   = SpaceTimeGrid::makeFiveDimGrid(Ls,UGrid_f);
+  GridRedBlackCartesian *FrbGrid_f = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,UGrid_f);
+
+  CheckPartitionOrder(UGrid_f,layout);
+
+  std::vector<int> seeds4({1,2,3,4});
+  std::vector<int> seeds5({5,6,7,8});
+  GridParallelRNG RNG4(UGrid_d);
+  GridParallelRNG RNG5(FGrid_d);
+  RNG4.SeedFixedIntegers(seeds4);
+  RNG5.SeedFixedIntegers(seeds5);
+
+  LatticeGaugeFieldD Umu_d(UGrid_d);
+  LatticeGaugeFieldF Umu_f(UGrid_f);
+  SU<Nc>::HotConfiguration(RNG4,Umu_d);
+  precisionChange(Umu_f,Umu_d);
+
+  // Antiperiodic in time for some actions, so boundary phases must travel with the links
+  WilsonImplParams antiperiodic;
+  antiperiodic.boundary_phases[Nd-1] = -1.0;
+
+  // Sources: odd checkerboard (Schur) and full lattice (MdagM), 4d and 5d
+  std::vector<LatticeFermionD> src4_o(NBATCH,UrbGrid_d);
+  std::vector<LatticeFermionD> src5_o(NBATCH,FrbGrid_d);
+  std::vector<LatticeFermionD> src5(NBATCH,FGrid_d);
+  LatticeFermionD tmp4(UGrid_d);
+  LatticeFermionD tmp5(FGrid_d);
+  for(int i=0;i<NBATCH;i++){
+    random(RNG4,tmp4);
+    random(RNG5,tmp5);
+    pickCheckerboard(Odd,src4_o[i],tmp4);
+    pickCheckerboard(Odd,src5_o[i],tmp5);
+    src5[i] = tmp5;
+  }
+
+  //////////////////////////////////////////
+  // Wilson
+  //////////////////////////////////////////
+  {
+    RealD mass = 0.1;
+    WilsonFermionD Dd(Umu_d,*UGrid_d,*UrbGrid_d,mass);
+    WilsonFermionF Df(Umu_f,*UGrid_f,*UrbGrid_f,mass);
+    SchurDiagMooeeOperator<WilsonFermionD,LatticeFermionD> Ld(Dd);
+    SchurDiagMooeeOperator<WilsonFermionF,LatticeFermionF> Lf(Df);
+    CheckSplitBatched("Wilson SchurDiagMooee",Ld,Lf,UrbGrid_f,src4_o,layout);
+  }
+
+  //////////////////////////////////////////
+  // Wilson clover
+  //////////////////////////////////////////
+  {
+    RealD mass  = 0.1;
+    RealD csw_r = 1.0;
+    RealD csw_t = 1.0;
+    WilsonCloverFermionD Dd(Umu_d,*UGrid_d,*UrbGrid_d,mass,csw_r,csw_t);
+    WilsonCloverFermionF Df(Umu_f,*UGrid_f,*UrbGrid_f,mass,csw_r,csw_t);
+    SchurDiagMooeeOperator<WilsonCloverFermionD,LatticeFermionD> Ld(Dd);
+    SchurDiagMooeeOperator<WilsonCloverFermionF,LatticeFermionF> Lf(Df);
+    CheckSplitBatched("WilsonClover SchurDiagMooee",Ld,Lf,UrbGrid_f,src4_o,layout);
+  }
+
+  //////////////////////////////////////////
+  // Compact Wilson clover, antiperiodic
+  //////////////////////////////////////////
+  {
+    RealD mass  = 0.1;
+    RealD csw_r = 1.0;
+    RealD csw_t = 1.0;
+    RealD cF    = 1.0;
+    WilsonAnisotropyCoefficients anis;
+    CompactWilsonCloverFermionD Dd(Umu_d,*UGrid_d,*UrbGrid_d,mass,csw_r,csw_t,cF,anis,antiperiodic);
+    CompactWilsonCloverFermionF Df(Umu_f,*UGrid_f,*UrbGrid_f,mass,csw_r,csw_t,cF,anis,antiperiodic);
+    SchurDiagMooeeOperator<CompactWilsonCloverFermionD,LatticeFermionD> Ld(Dd);
+    SchurDiagMooeeOperator<CompactWilsonCloverFermionF,LatticeFermionF> Lf(Df);
+    CheckSplitBatched("CompactWilsonClover SchurDiagMooee antiperiodic",Ld,Lf,UrbGrid_f,src4_o,layout);
+  }
+
+  //////////////////////////////////////////
+  // Domain wall
+  //////////////////////////////////////////
+  {
+    RealD mass = 0.1;
+    RealD M5   = 1.8;
+    DomainWallFermionD Dd(Umu_d,*FGrid_d,*FrbGrid_d,*UGrid_d,*UrbGrid_d,mass,M5);
+    DomainWallFermionF Df(Umu_f,*FGrid_f,*FrbGrid_f,*UGrid_f,*UrbGrid_f,mass,M5);
+    SchurDiagMooeeOperator<DomainWallFermionD,LatticeFermionD> Ld(Dd);
+    SchurDiagMooeeOperator<DomainWallFermionF,LatticeFermionF> Lf(Df);
+    CheckSplitBatched("DomainWall SchurDiagMooee",Ld,Lf,FrbGrid_f,src5_o,layout);
+  }
+
+  //////////////////////////////////////////
+  // Mobius, antiperiodic, with unequal
+  // masses; all wrapper kinds
+  //////////////////////////////////////////
+  {
+    RealD mass = 0.1;
+    RealD M5   = 1.8;
+    RealD b    = 1.5;
+    RealD c    = 0.5;
+    MobiusFermionD Dd(Umu_d,*FGrid_d,*FrbGrid_d,*UGrid_d,*UrbGrid_d,mass,M5,b,c,antiperiodic);
+    MobiusFermionF Df(Umu_f,*FGrid_f,*FrbGrid_f,*UGrid_f,*UrbGrid_f,mass,M5,b,c,antiperiodic);
+    Dd.SetMass(0.1,0.12);
+    Df.SetMass(0.1,0.12);
+    {
+      SchurDiagMooeeOperator<MobiusFermionD,LatticeFermionD> Ld(Dd);
+      SchurDiagMooeeOperator<MobiusFermionF,LatticeFermionF> Lf(Df);
+      CheckSplitBatched("Mobius SchurDiagMooee antiperiodic",Ld,Lf,FrbGrid_f,src5_o,layout);
+    }
+    {
+      SchurDiagOneOperator<MobiusFermionD,LatticeFermionD> Ld(Dd);
+      SchurDiagOneOperator<MobiusFermionF,LatticeFermionF> Lf(Df);
+      CheckSplitBatched("Mobius SchurDiagOne antiperiodic",Ld,Lf,FrbGrid_f,src5_o,layout);
+    }
+    {
+      SchurDiagTwoOperator<MobiusFermionD,LatticeFermionD> Ld(Dd);
+      SchurDiagTwoOperator<MobiusFermionF,LatticeFermionF> Lf(Df);
+      CheckSplitBatched("Mobius SchurDiagTwo antiperiodic",Ld,Lf,FrbGrid_f,src5_o,layout);
+    }
+    {
+      MdagMLinearOperator<MobiusFermionD,LatticeFermionD> Ld(Dd);
+      MdagMLinearOperator<MobiusFermionF,LatticeFermionF> Lf(Df);
+      CheckSplitBatched("Mobius MdagM antiperiodic",Ld,Lf,FGrid_f,src5,layout);
+    }
+  }
+
+  //////////////////////////////////////////
+  // ZMobius
+  //////////////////////////////////////////
+  {
+    RealD mass = 0.1;
+    RealD M5   = 1.8;
+    RealD b    = 1.0;
+    RealD c    = 0.0;
+    std::vector<ComplexD> gamma(Ls);
+    for(int s=0;s<Ls;s++){
+      gamma[s] = ComplexD(1.0+0.05*s, (s%2) ? 0.02 : -0.02);
+    }
+    ZMobiusFermionD Dd(Umu_d,*FGrid_d,*FrbGrid_d,*UGrid_d,*UrbGrid_d,mass,M5,gamma,b,c);
+    ZMobiusFermionF Df(Umu_f,*FGrid_f,*FrbGrid_f,*UGrid_f,*UrbGrid_f,mass,M5,gamma,b,c);
+    SchurDiagMooeeOperator<ZMobiusFermionD,LatticeFermionD> Ld(Dd);
+    SchurDiagMooeeOperator<ZMobiusFermionF,LatticeFermionF> Lf(Df);
+    CheckSplitBatched("ZMobius SchurDiagMooee",Ld,Lf,FrbGrid_f,src5_o,layout);
+  }
+
+  //////////////////////////////////////////
+  // A derived operator must not inherit
+  // its parent's SplitClone
+  //////////////////////////////////////////
+  {
+    RealD mass = 0.1;
+    RealD mu   = 0.1;
+    WilsonTMFermionF Df(Umu_f,*UGrid_f,*UrbGrid_f,mass,mu);
+    SchurDiagMooeeOperator<WilsonTMFermionF,LatticeFermionF> Lf(Df);
+    SplitOperator<LatticeFermionF> *split = Lf.SplitClone(layout);
+    std::cout << GridLogMessage << "WilsonTM SplitClone refused: " << (split == nullptr) << std::endl;
+    GRID_ASSERT(split == nullptr);
+  }
+
+  std::cout << GridLogMessage << "Test_split_mixedprec_batched: all checks passed" << std::endl;
+
+  Grid_finalize();
+}

tests/solver/Test_split_mobius_batched.cc

@@ -0,0 +1,236 @@
+/*************************************************************************************
+
+    Grid physics library, www.github.com/paboyle/Grid
+
+    Source file: ./tests/solver/Test_split_mobius_batched.cc
+
+    Copyright (C) 2026
+
+Author: Peter Boyle <paboyle@ph.ed.ac.uk>
+
+    This program is free software; you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation; either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License along
+    with this program; if not, write to the Free Software Foundation, Inc.,
+    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+    See the full license in the file "LICENSE" in the top level distribution directory
+*************************************************************************************/
+/*  END LEGAL */
+/////////////////////////////////////////////////////////////////////////////////////////////
+// Production-size timing of MixedPrecisionConjugateGradientBatched for Mobius with the
+// Hadrons default SchurDiagMooeeOperator: the same batch solved without and then with
+// split inner solves (--batched-solver-split), reporting wall clock, per-rhs iterations
+// and true residuals for each.
+//
+//   --Ls 12 --mass 0.026 --M5 1.8 --b 1.5 --c 0.5 --nbatch 4 --tol 1e-8
+//   --config <NERSC file>    (omit for a hot configuration: timing only, not physics)
+//   --nounsplit              (skip the reference unsplit solve)
+//   --repeat N               (split solve N times; host RSS must not grow between them)
+//
+// MEMORY lines report host RSS (current and peak) and allocator cache sizes, maximum over
+// ranks, at each phase: with --enable-unified=no every Lattice lives in host memory.
+//
+// Only one solution vector is kept, so the driver's own footprint is two batches of
+// double red-black 5d fields; the solver adds about as much again.
+/////////////////////////////////////////////////////////////////////////////////////////////
+#include <Grid/Grid.h>
+#include <sys/resource.h>
+#ifdef __APPLE__
+#include <mach/mach.h>
+#endif
+
+using namespace std;
+using namespace Grid;
+
+// Host memory of this process in GB: current resident set and its high-water mark
+void HostRSS(RealD &current,RealD &peak)
+{
+  struct rusage ru;
+  getrusage(RUSAGE_SELF,&ru);
+#ifdef __APPLE__
+  peak = ru.ru_maxrss/1.0e9;              // bytes on macOS
+  mach_task_basic_info_data_t info;
+  mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;
+  task_info(mach_task_self(),MACH_TASK_BASIC_INFO,(task_info_t)&info,&count);
+  current = info.resident_size/1.0e9;
+#else
+  peak = ru.ru_maxrss*1024.0/1.0e9;       // kilobytes on Linux
+  long pages = 0;
+  long resident = 0;
+  FILE *f = fopen("/proc/self/statm","r");
+  if ( f ) {
+    if ( fscanf(f,"%ld %ld",&pages,&resident) != 2 ) {
+      resident = 0;
+    }
+    fclose(f);
+  }
+  current = resident*(RealD)sysconf(_SC_PAGESIZE)/1.0e9;
+#endif
+}
+
+// Largest values over ranks: host RSS now and at peak, and the allocator caches
+void ReportMemory(GridBase *grid,const std::string &phase)
+{
+  RealD rss;
+  RealD peak;
+  HostRSS(rss,peak);
+  RealD hostcache = MemoryManager::HostCacheBytes()/1.0e9;
+  RealD devcache  = MemoryManager::DeviceCacheBytes()/1.0e9;
+  grid->GlobalMax(rss);
+  grid->GlobalMax(peak);
+  grid->GlobalMax(hostcache);
+  grid->GlobalMax(devcache);
+  std::cout << GridLogMessage << "MEMORY " << phase
+            << " : host RSS " << rss << " GB, peak " << peak
+            << " GB; allocator cache host " << hostcache << " GB, device " << devcache
+            << " GB (max over ranks)" << std::endl;
+  HostMemoryReport(grid,GridLogMessage,phase);
+}
+
+typedef LatticeFermionD FieldD;
+typedef LatticeFermionF FieldF;
+
+template<class T>
+T CmdOption(int argc,char **argv,const std::string &name,T def)
+{
+  T val = def;
+  if ( GridCmdOptionExists(argv,argv+argc,name) ) {
+    std::stringstream ss(GridCmdOptionPayload(argv,argv+argc,name));
+    ss >> val;
+  }
+  return val;
+}
+
+void SolveAndReport(const std::string &label,
+                    MixedPrecisionConjugateGradientBatched<FieldD,FieldF> &mCG,
+                    LinearOperatorBase<FieldD> &Linop_d,
+                    std::vector<FieldD> &src,
+                    std::vector<FieldD> &sol)
+{
+  int nbatch = src.size();
+  for(int i=0;i<nbatch;i++){
+    sol[i].Checkerboard() = src[i].Checkerboard();
+    sol[i] = Zero();
+  }
+
+  std::cout << GridLogMessage << "==================================================" << std::endl;
+  std::cout << GridLogMessage << label << " batched solve, nbatch " << nbatch << std::endl;
+  std::cout << GridLogMessage << "==================================================" << std::endl;
+
+  RealD t0 = usecond();
+  mCG(src,sol);
+  RealD t1 = usecond();
+
+  FieldD Msol(src[0].Grid());
+  Msol.Checkerboard() = src[0].Checkerboard();
+  RealD worst = 0.0;
+  for(int i=0;i<nbatch;i++){
+    Linop_d.HermOp(sol[i],Msol);
+    Msol = Msol - src[i];
+    RealD resid = std::sqrt(norm2(Msol)/norm2(src[i]));
+    worst = std::max(worst,resid);
+    std::cout << GridLogMessage << label << ": rhs " << i << " true residual " << resid << std::endl;
+  }
+  std::cout << GridLogMessage << label << ": SUMMARY wall clock " << (t1-t0)/1.0e6
+            << " s for " << nbatch << " rhs, " << (t1-t0)/1.0e6/nbatch
+            << " s/rhs, worst true residual " << worst << std::endl;
+}
+
+int main (int argc, char ** argv)
+{
+  Grid_init(&argc,&argv);
+
+  int         Ls        = CmdOption<int>        (argc,argv,"--Ls",12);
+  RealD       mass      = CmdOption<RealD>      (argc,argv,"--mass",0.026);
+  RealD       M5        = CmdOption<RealD>      (argc,argv,"--M5",1.8);
+  RealD       b         = CmdOption<RealD>      (argc,argv,"--b",1.5);
+  RealD       c         = CmdOption<RealD>      (argc,argv,"--c",0.5);
+  int         nbatch    = CmdOption<int>        (argc,argv,"--nbatch",4);
+  RealD       tol       = CmdOption<RealD>      (argc,argv,"--tol",1.0e-8);
+  std::string config    = CmdOption<std::string>(argc,argv,"--config",std::string(""));
+  bool        unsplit   = !GridCmdOptionExists(argv,argv+argc,"--nounsplit");
+  int         repeat    = CmdOption<int>        (argc,argv,"--repeat",1);
+
+  std::cout << GridLogMessage << "Mobius Ls " << Ls << " mass " << mass << " M5 " << M5
+            << " b " << b << " c " << c << " nbatch " << nbatch << " tol " << tol << std::endl;
+
+  GridCartesian         *UGrid_d   = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(), GridDefaultSimd(Nd,vComplexD::Nsimd()), GridDefaultMpi());
+  GridRedBlackCartesian *UrbGrid_d = SpaceTimeGrid::makeFourDimRedBlackGrid(UGrid_d);
+  GridCartesian         *FGrid_d   = SpaceTimeGrid::makeFiveDimGrid(Ls,UGrid_d);
+  GridRedBlackCartesian *FrbGrid_d = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,UGrid_d);
+
+  GridCartesian         *UGrid_f   = SpaceTimeGrid::makeFourDimGrid(GridDefaultLatt(), GridDefaultSimd(Nd,vComplexF::Nsimd()), GridDefaultMpi());
+  GridRedBlackCartesian *UrbGrid_f = SpaceTimeGrid::makeFourDimRedBlackGrid(UGrid_f);
+  GridCartesian         *FGrid_f   = SpaceTimeGrid::makeFiveDimGrid(Ls,UGrid_f);
+  GridRedBlackCartesian *FrbGrid_f = SpaceTimeGrid::makeFiveDimRedBlackGrid(Ls,UGrid_f);
+
+  std::vector<int> seeds4({1,2,3,4});
+  std::vector<int> seeds5({5,6,7,8});
+  GridParallelRNG RNG4(UGrid_d);
+  GridParallelRNG RNG5(FGrid_d);
+  RNG4.SeedFixedIntegers(seeds4);
+  RNG5.SeedFixedIntegers(seeds5);
+
+  LatticeGaugeFieldD Umu_d(UGrid_d);
+  LatticeGaugeFieldF Umu_f(UGrid_f);
+  if ( config.size() ) {
+    FieldMetaData header;
+    NerscIO::readConfiguration(Umu_d,header,config);
+  } else {
+    std::cout << GridLogMessage << "No --config: hot configuration, timing only" << std::endl;
+    SU<Nc>::HotConfiguration(RNG4,Umu_d);
+  }
+  precisionChange(Umu_f,Umu_d);
+  ReportMemory(UGrid_d,"gauge field ready");
+
+  // Antiperiodic in time, as in production
+  WilsonImplParams params;
+  params.boundary_phases[Nd-1] = -1.0;
+
+  MobiusFermionD Dd(Umu_d,*FGrid_d,*FrbGrid_d,*UGrid_d,*UrbGrid_d,mass,M5,b,c,params);
+  MobiusFermionF Df(Umu_f,*FGrid_f,*FrbGrid_f,*UGrid_f,*UrbGrid_f,mass,M5,b,c,params);
+  SchurDiagMooeeOperator<MobiusFermionD,FieldD> Linop_d(Dd);
+  SchurDiagMooeeOperator<MobiusFermionF,FieldF> Linop_f(Df);
+
+  std::vector<FieldD> src(nbatch,FrbGrid_d);
+  std::vector<FieldD> sol(nbatch,FrbGrid_d);
+  {
+    FieldD tmp(FGrid_d);
+    for(int i=0;i<nbatch;i++){
+      random(RNG5,tmp);
+      pickCheckerboard(Odd,src[i],tmp);
+    }
+  }
+
+  ReportMemory(UGrid_d,"operators and sources ready");
+
+  if ( unsplit ) {
+    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");
+  }
+
+  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();
+}