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/SplitOperator.h

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

Grid/algorithms/iterative/ConjugateGradientMixedPrecBatched.h

@@ -53,9 +53,17 @@ public:
   
   // Inner solves on independent partitions of the communicator (--batched-solver-split).
   // BatchedSplit is the partition MPI layout, empty for no split; BatchedSplitNode selects
-  // one partition per node. Default from the command line; callers may override.
-  Coordinate BatchedSplit;
-  bool       BatchedSplitNode;
+  // one partition per node. Fixed at construction, default from the command line.
+  const Coordinate BatchedSplit;
+  const bool       BatchedSplitNode;
+
+  // Copy of Linop_f on the partitions, made once at construction and owned; nullptr when
+  // there is no split or the operator cannot be split. It holds the gauge field as it was at
+  // construction: a solver must not outlive a change of the operator's gauge field. The
+  // outer defect correction uses Linop_d, so a stale copy would cost convergence, not
+  // correctness.
+  SplitOperator<FieldF> *split = nullptr;
+  int                    Partitions = 1;
 
   MixedPrecisionConjugateGradientBatched(RealD tol, 
           Integer maxinnerit, 
@@ -64,11 +72,37 @@ public:
           GridBase* _sp_grid, 
           LinearOperatorBase<FieldF> &_Linop_f, 
           LinearOperatorBase<FieldD> &_Linop_d,
-          bool _updateResidual=true) :
+          bool _updateResidual=true,
+          const Coordinate &_split=GridDefaultBatchedSolverSplit(),
+          bool _split_node=GridDefaultBatchedSolverSplitNode()) :
     Linop_f(_Linop_f), Linop_d(_Linop_d),
     Tolerance(tol), InnerTolerance(tol), MaxInnerIterations(maxinnerit), MaxOuterIterations(maxouterit), MaxPatchupIterations(maxpatchit), SinglePrecGrid(_sp_grid),
     OuterLoopNormMult(100.), guesser(NULL), updateResidual(_updateResidual),
-    BatchedSplit(GridDefaultBatchedSolverSplit()), BatchedSplitNode(GridDefaultBatchedSolverSplitNode()) { };
+    BatchedSplit(_split), BatchedSplitNode(_split_node)
+  {
+    Coordinate layout = BatchedSolverSplitLayout(SinglePrecGrid,BatchedSplit,BatchedSplitNode,Partitions);
+    if ( Partitions > 1 ) {
+      double t0 = usecond();
+      split = Linop_f.SplitClone(layout);
+      double t1 = usecond();
+      if ( split == nullptr ) {
+        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: operator cannot be split; serial inner solves" << std::endl;
+      } else {
+        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: split clone " << (t1-t0)/1.0e6 << " s" << std::endl;
+        HostMemoryReport(SinglePrecGrid,GridLogMessage,"MixedPrecisionConjugateGradientBatched: after split clone");
+      }
+    }
+  };
+
+  // Owns split
+  MixedPrecisionConjugateGradientBatched(const MixedPrecisionConjugateGradientBatched &) = delete;
+
+  MixedPrecisionConjugateGradientBatched &operator=(const MixedPrecisionConjugateGradientBatched &) = delete;
+
+  ~MixedPrecisionConjugateGradientBatched(void)
+  {
+    delete split;
+  }
 
   void useGuesser(LinearFunction<FieldF> &g){
     guesser = &g;
@@ -98,6 +132,8 @@ public:
 
     GridStopWatch InnerCGtimer;
     GridStopWatch PrecChangeTimer;
+    GridStopWatch OuterResidualTimer;
+    GridStopWatch PatchupTimer;
     
     int cb = src_d_in[0].Checkerboard();
     
@@ -137,18 +173,10 @@ public:
     ConjugateGradient<FieldF> CG_f(inner_tol, MaxInnerIterations);
     CG_f.ErrorOnNoConverge = false;
     
-    //Optionally clone the single precision operator onto split-communicator partitions
-    GridStopWatch SplitTimer;
-    SplitOperator<FieldF> *split = nullptr;
-    int partitions;
-    Coordinate layout = BatchedSolverSplitLayout(SinglePrecGrid,BatchedSplit,BatchedSplitNode,NBatch,partitions);
-    if ( partitions > 1 ) {
-      SplitTimer.Start();
-      split = Linop_f.SplitClone(layout);
-      SplitTimer.Stop();
-      if ( split == nullptr ) {
-        std::cout << GridLogMessage << "MixedPrecisionConjugateGradientBatched: operator cannot be split; serial inner solves" << std::endl;
-      }
+    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
@@ -156,13 +184,16 @@ public:
     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;
 
@@ -191,7 +222,7 @@ public:
       }
 
       if ( split != nullptr ) {
-        InnerSplitSolves(*split, CG_f, src_f, sol_f, TotalInnerIterations, InnerCGtimer, SplitTimer);
+        InnerSplitSolves(*split, CG_f, src_f, sol_f, TotalInnerIterations, InnerCGtimer, splitTimers);
       }
 
       for (int i=0; i<NBatch; i++) {
@@ -213,17 +244,17 @@ public:
 
     }
 
-    delete split;
-    
     //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();
 
@@ -232,12 +263,36 @@ public:
       std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: solve " << i << " Inner CG iterations " << TotalInnerIterations[i] << " Restarts " << TotalOuterIterations << " Final CG iterations " << TotalFinalStepIterations[i] << std::endl;
     }
     std::cout << GridLogMessage << std::endl;
-    std::cout<<GridLogMessage<<"MixedPrecisionConjugateGradientBatched: Total time " << TotalTimer.Elapsed() << " Precision change " << PrecChangeTimer.Elapsed() << " Inner CG total " << InnerCGtimer.Elapsed() << " Split setup and transfer " << SplitTimer.Elapsed() << 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.
@@ -248,7 +303,7 @@ private:
                         std::vector<FieldF> &sol_f,
                         std::vector<Integer> &TotalInnerIterations,
                         GridStopWatch &InnerCGtimer,
-                        GridStopWatch &SplitTimer)
+                        SplitTimers &timers)
   {
     int NBatch = src_f.size();
     int P      = split.Partitions;
@@ -263,6 +318,7 @@ private:
     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;
@@ -274,22 +330,34 @@ private:
           group_sol[p] = Zero();
         }
       }
+      timers.Staging.Stop();
 
       // The initial guess (e.g. from the guesser) travels with the source
-      SplitTimer.Start();
+      timers.Split.Start();
       Grid_split(group_src,s_src);
       Grid_split(group_sol,s_sol);
-      SplitTimer.Stop();
+      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();
 
-      SplitTimer.Start();
+      timers.Unsplit.Start();
       Grid_unsplit(group_sol,s_sol);
-      SplitTimer.Stop();
+      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;
       }
@@ -297,11 +365,15 @@ private:
         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();
     }
   }
 };

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];
 
@@ -262,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]);
     }
   }
@@ -421,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,
@@ -428,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,
@@ -454,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;
@@ -475,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;
@@ -506,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/fermion/FermionOperator.h

@@ -70,6 +70,8 @@ public:
     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);
@@ -85,6 +87,7 @@ public:
       split->FermionGrid   = split->GaugeGrid;
       split->FermionRBGrid = split->GaugeRBGrid;
     }
+    std::cout << GridLogPerformance << "MakeSplitGrids: " << (usecond()-t0)/1.0e6 << " s" << std::endl;
     return split;
   }
 

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/WilsonFermion5D.h

@@ -192,9 +192,14 @@ public:
   // 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;
   }
 
   ///////////////////////////////////////////////////////////////

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

@@ -367,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
   //////////////////////////////////////////////////////////
@@ -541,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;
@@ -595,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;

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}

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

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

tests/solver/Test_split_mobius_batched.cc

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