step_solve_mpi.hpp

Go to the documentation of this file.

 
#ifndef DEM_STEP_SOLVE_MPI_HPP
#define DEM_STEP_SOLVE_MPI_HPP
 
#ifdef PECLET_DEM_MPI
 
#include <algorithm>
#include <cmath>
#include <Kokkos_Core.hpp>
 
#include "mpi_halo.hpp"  // ParticleHalo
#include "particles.hpp"
#include "solve_driver.hpp"        // demSolveContacts, maxOwnedRadius
#include "solve_driver_force.hpp"  // demStepForce + HertzMindlinLaw + fillWorldRadiiKokkos
#include "step_solve.hpp"          // narrowPhaseGrow
 
namespace peclet::dem {
 
struct MpiSolveHooks {
  static constexpr bool distributed = true;
  ParticleHalo& halo;
  int syncEvery;
  bool forwardRotation;
  int numManifoldsVisible;
  float allMax(float v) const {
    float g = v;
    MPI_Allreduce(&v, &g, 1, MPI_FLOAT, MPI_MAX, halo.comm());
    return g;
  }
  float allMaxAny(float v, bool& any) const {
    const float in[2] = {v, any ? 1.0f : 0.0f};
    float out[2] = {v, in[1]};
    MPI_Allreduce(in, out, 2, MPI_FLOAT, MPI_MAX, halo.comm());
    any = out[1] > 0.0f;
    return out[0];
  }
  bool syncPoint(int it) const { return (it + 1) % syncEvery == 0; }
  int visibleManifolds(int) const { return numManifoldsVisible; }
  void beginSolve(Particles& P) const { halo.markVelocityBaseline(P); }
  void syncVelocities(Particles& P) const { halo.syncVelocities(P, forwardRotation); }
  void publishPositions(Particles& P) const { halo.publishPositions(P, forwardRotation); }
  void syncPositions(Particles& P) const { halo.syncPositions(P, forwardRotation); }
  void syncFrictionCounts(Particles& P) const { halo.syncFrictionCounts(P); }
  void syncContactCounts(Particles& P) const { halo.syncContactCounts(P); }
  // Rank-level M (docs/contact_solve_framework.md §13.3).
  RankK rankK() const {
    return RankK{halo.exchanges(), halo.rankColor(), halo.numRankColors(), syncEvery};
  }
  void openVelocityPhase(Particles& P, bool poisson) const {
    halo.openVelocityPhase(P, forwardRotation, poisson);
  }
  void openPositionCounts(Particles& P, bool velocityMask) const {
    halo.openPositionCounts(P, velocityMask);
  }
  void restoreOrphanBalance(Particles& P) const { halo.restoreOrphanBalance(P); }
};
 
inline constexpr float kDriftSlack = 0.25f;
 
inline void driftVoteLocalKokkos(const Particles& P, const ParticleHalo::BlockBox box, float dt,
                                 float& eMax, float& dMax, float* scaleMax = nullptr) {
  const F3 g = P.gravity;
  const auto pos = P.pos;
  const auto vel = P.vel;
  const auto ext = P.extForce;
  const auto invM = P.invMass;
  const auto sc = P.scale;
  const float lo0 = box.lo.x, lo1 = box.lo.y, lo2 = box.lo.z, hi0 = box.hi.x, hi1 = box.hi.y,
              hi2 = box.hi.z, L0 = box.period.x, L1 = box.period.y, L2 = box.period.z;
  float e = 0.0f, d = 0.0f, smax = 0.0f;
  Kokkos::parallel_reduce(
      "peclet::dem::drift_vote", Kokkos::RangePolicy<CpExec>(0, P.numReal),
      KOKKOS_LAMBDA(int i, float& em, float& dm, float& sm) {
        const float im = invM(i);
        const float ax = g.x + ext(i, 0) * im, ay = g.y + ext(i, 1) * im, az = g.z + ext(i, 2) * im;
        const float v0 = vel(i, 0), v1 = vel(i, 1), v2 = vel(i, 2);
        const float di = Kokkos::sqrt(v0 * v0 + v1 * v1 + v2 * v2) * dt +
                         Kokkos::sqrt(ax * ax + ay * ay + az * az) * dt * dt;
        float d2 = 0.0f;
#define PECLET_DEM_DRIFT_AXIS(X, LO, HI, L)                                          \
  {                                                                                  \
    float gk = Kokkos::fmax(0.0f, Kokkos::fmax((LO) - (X), (X) - (HI)));             \
    if ((L) > 0.0f) {                                                                \
      const float xm = (X) - (L), xp = (X) + (L);                                    \
      gk = Kokkos::fmin(gk, Kokkos::fmax(0.0f, Kokkos::fmax((LO) - xm, xm - (HI)))); \
      gk = Kokkos::fmin(gk, Kokkos::fmax(0.0f, Kokkos::fmax((LO) - xp, xp - (HI)))); \
    }                                                                                \
    d2 += gk * gk;                                                                   \
  }
        PECLET_DEM_DRIFT_AXIS(pos(i, 0), lo0, hi0, L0)
        PECLET_DEM_DRIFT_AXIS(pos(i, 1), lo1, hi1, L1)
        PECLET_DEM_DRIFT_AXIS(pos(i, 2), lo2, hi2, L2)
#undef PECLET_DEM_DRIFT_AXIS
        const float ei = Kokkos::sqrt(d2) + di;
        em = ei > em ? ei : em;
        dm = di > dm ? di : dm;
        sm = sc(i) > sm ? sc(i) : sm;
      },
      Kokkos::Max<float>(e), Kokkos::Max<float>(d), Kokkos::Max<float>(smax));
  eMax = e;
  dMax = d;
  if (scaleMax)
    *scaleMax = smax;
}
 
inline float globalMaxRadius(const Particles& P, MPI_Comm comm) {
  float r = P.numReal > 0 ? maxOwnedRadius(P) : 0.0f, g = r;
  MPI_Allreduce(&r, &g, 1, MPI_FLOAT, MPI_MAX, comm);
  return g > 0.0f ? g : maxOwnedRadius(P);
}
 
inline double xpbdContactReach(float rMax) {
  return 2.1 * static_cast<double>(rMax);
}
 
inline void debugCaptureOwnedContacts(Particles& P, const ParticleHalo& halo, int ncOwned) {
  DebugContactCapture& cap = P.debugCaptured;
  const int no = P.numReal;
  auto hc = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), P.contacts);
  auto hg = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), P.gid);
  auto hp = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), P.posPred);
  auto hr = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), P.rad);
  const std::vector<float> shift = halo.debugGhostShiftsBySlot();
  const double box[3] = {static_cast<double>(P.domain.max.x) - P.domain.min.x,
                         static_cast<double>(P.domain.max.y) - P.domain.min.y,
                         static_cast<double>(P.domain.max.z) - P.domain.min.z};
  auto image = [&](int slot, std::vector<int>& out) {
    for (int d = 0; d < 3; ++d) {
      const float sh = slot < no ? 0.0f : shift[3 * static_cast<std::size_t>(slot - no) + d];
      out.push_back(sh == 0.0f ? 0 : static_cast<int>(std::lround(sh / box[d])));
    }
  };
  cap = DebugContactCapture{};
  for (int i = 0; i < ncOwned; ++i) {
    const ContactC& c = hc(i);
    cap.gidA.push_back(hg(c.bodyA));
    cap.gidB.push_back(c.bodyB >= 0 ? hg(c.bodyB) : -1);
    image(c.bodyA, cap.imageA);
    if (c.bodyB >= 0)
      image(c.bodyB, cap.imageB);
    else
      cap.imageB.insert(cap.imageB.end(), {0, 0, 0});
    cap.dist.push_back(c.dist);
  }
  for (int i = 0; i < no; ++i) {
    cap.gid.push_back(hg(i));
    for (int d = 0; d < 3; ++d)
      cap.posPred.push_back(hp(i, d));
    cap.rad.push_back(hr(i));
  }
}
 
struct MpiDriftVote {
  float rMax;    // global max radius (pre-growth)
  float margin;  // 0.1 R_max: the broad/narrow-phase margin
  float slack;   // S = kDriftSlack R_max
  float dMax;    // global max predicted displacement of the substep
};
inline MpiDriftVote mpiDriftVote(Particles& P, ParticleHalo& halo) {
  float vote[3] = {0.0f, 0.0f, 0.0f}, smax = 0.0f;
  driftVoteLocalKokkos(P, halo.blockBox(), P.dt, vote[1], vote[2], &smax);
  // = maxOwnedRadius(P) on a non-empty rank (same float product), 0 on an empty one.
  vote[0] = P.numReal > 0 ? smax * P.globalScale * P.baseRadius : 0.0f;
  float g[3];
  MPI_Allreduce(vote, g, 3, MPI_FLOAT, MPI_MAX, halo.comm());
  const MpiDriftVote v{g[0], 0.1f * g[0], kDriftSlack * g[0], g[2]};
  int commSize = 1;
  MPI_Comm_size(halo.comm(), &commSize);
  // One rank owns everything: there is no other owner to move a particle to (a migration would be
  // a host round trip that changes nothing).
#if defined(PECLET_DEM_TEST_MUTANT) && PECLET_DEM_TEST_MUTANT == 5
  if (false) {  // G13 mutant 5: the drift vote never migrates
#else
  if (commSize > 1 && g[1] >= v.slack) {
#endif
    halo.migrateToBlocks(P);
    ++P.splitStats.driftMigrations;
  }
  return v;
}
inline double mpiXpbdBand(const Particles& P, ParticleHalo& halo, double rcut,
                          const MpiDriftVote& v) {
  return std::max(rcut, xpbdContactReach(globalMaxRadius(P, halo.comm())) +
                            static_cast<double>(v.slack) +
                            (halo.verletSkin() > 0.0f ? 0.0 : static_cast<double>(v.dMax)));
}
 
inline void demStepMpi(Particles& P, ParticleHalo& halo, double rcut, int syncEvery,
                       bool forwardRotation) {
  CpExec space;
  Kokkos::Profiling::pushRegion("dem::vote");
  const MpiDriftVote vote = mpiDriftVote(P, halo);
  Kokkos::Profiling::popRegion();
  // Broadphase / narrow-phase margin: 0.1 R_max over ALL ranks, so both owners of a cross-face
  // pair report it at the same gap (a rank-local R_max dropped the pair on the rank whose own
  // grains are small while the other kept it).
  const float margin = vote.margin;
 
  if (P.growthFactor != -1.0f && P.growthRate != 0.0f) {
    P.growthFactor *= std::exp(P.growthRate * P.dt);
    if (P.growthFactor > 1.0f)
      P.growthFactor = 1.0f;
  }
  if (P.growthFactor > 0.0f)
    updateGrowthScalesKokkos(P.numReal, P.scale, P.targetScale, P.growthFactor);
  const double band = mpiXpbdBand(P, halo, rcut, vote);
 
  // 1. Predict velocity on the owned set (no ghosts yet -> numParticles == numReal).
  P.numParticles = P.numReal;
  predictVelocityKokkos(P.numReal, P.pos, P.invMass, P.vel, P.quat, P.angVel, P.invInertia,
                        P.posPred, P.quatPred, P.velPred, P.angVelPred, P.deltaPos, P.deltaQuat,
                        P.deltaVel, P.deltaAngVel, P.constraintCounts, P.gravity, P.dt, P.extForce,
                        P.extTorque);
 
  // 2. Gather ghosts (real mass) from owners over the halo: full state -- including gid,
  //    materialId and the warm grounded level -- into the ghost slots; sets
  //    P.numParticles = numReal + numGhost and self-maps realIndices.
  Kokkos::Profiling::pushRegion("dem::gather");
  halo.gather(P, band);
  Kokkos::Profiling::popRegion();
  // The velocity-phase slot map (docs/contact_solve_framework.md §6.2): one velocity slot per body
  // per rank -- a periodic self image maps to the owned slot, other images of one body to its
  // lowest ghost slot. Here, not in gather (the Hertz engine shares gather). Positions stay per
  // slot.
  Kokkos::Profiling::pushRegion("dem::mapSlots");
  halo.mapVelocitySlots(P);
  Kokkos::Profiling::popRegion();
 
  fillWorldRadiiKokkos(P.scale, P.rad, P.globalScale, P.baseRadius, P.numParticles);
 
  // 3. Broad/narrow phase + manifold reduction over owned + ghosts (contactSlot map included:
  // the PGS friction bound and the position-channel Coulomb carry read through it).
  // findCollisionsGrow fences + reads the pair count back to host and guarantees np <= P.pairs
  // extent (growing the buffer on overflow) so the narrowphase never reads P.pairs out of bounds.
  Kokkos::Profiling::pushRegion("dem::broad");
  const int np = findCollisionsGrow(P, margin);
  Kokkos::Profiling::popRegion();
 
  Kokkos::deep_copy(space, P.contactCount, 0);
  Kokkos::deep_copy(space, P.maxOverlap, 0.0f);
  Kokkos::Profiling::pushRegion("dem::narrow");
  const int nc = narrowPhaseGrow(P, np, margin);
  Kokkos::Profiling::popRegion();
 
  // Contact ownership (docs/mpi_momentum_conservation.md §2.1-2.2): owned contacts first, then
  // the manifolds reduced owned-first. The solve runs on [0, ncOwned) / [0, nmOwned); the
  // visible counts stay in P.contactCount / P.manifoldCount for the getters.
  Kokkos::Profiling::pushRegion("dem::ownership");
  const int ncOwned = partitionContactsKokkos(P.contacts, nc, halo.contactOwnership(P));
  Kokkos::Profiling::popRegion();
  if (P.debugCapture)
    debugCaptureOwnedContacts(P, halo, ncOwned);
  int nmOwned = 0;
  reduceContactsToManifoldsKokkos(P.contacts, nc, P.manifolds, P.manifoldCount, P.contactSlot,
                                  ncOwned, &nmOwned);
  const int nmVisible = readInt(P.manifoldCount);
 
  // 4-6. The shared modern velocity + position solve, distributed: rank-local colouring over the
  // owned + ghost body slots (nBodies = numParticles; realIndices are self-mapped, so ghost
  // copies evolve in place between reconciliations), persistent-pair keys from the global ids.
  Kokkos::Profiling::pushRegion("dem::solve");
  demSolveContacts(P, ncOwned, nmOwned, P.numParticles, P.gid,
                   MpiSolveHooks{halo, syncEvery < 1 ? 1 : syncEvery, forwardRotation, nmVisible});
  Kokkos::Profiling::popRegion();
 
  // 7. Commit (owned results kept; the ghost slots, reconciled by the final sync, are dropped and
  // re-gathered next substep).
  finalCommitKokkos(P.numReal, P.pos, P.invMass, P.posPred, P.quat, P.quatPred, P.domain);
 
  if (P.thermostatTau > 0.0f && P.dt > 0.0f)
    applyThermostatKokkos(P.numReal, P.vel, P.invMass, P.angVel, P.invInertia, P.quat,
                          P.thermostatKB, P.thermostatTau, P.thermostatTemp, P.dt);
 
  P.numParticles = P.numReal;  // restore owned-only active count for getters
}
 
struct MpiForceHooks {
  static constexpr bool distributed = true;
  ParticleHalo& halo;
  double band;
  float slack;  // drift slack S = kDriftSlack R_max (§5.5)
  float allMax(float v) const {
    float g = v;
    MPI_Allreduce(&v, &g, 1, MPI_FLOAT, MPI_MAX, halo.comm());
    return g;
  }
  float allMin(float v) const {
    float g = v;
    MPI_Allreduce(&v, &g, 1, MPI_FLOAT, MPI_MIN, halo.comm());
    return g;
  }
  void gatherGhosts(Particles& P) const {
    // The drift vote at every pair-list rebuild (docs/contact_solve_framework.md §5.5): if some
    // particle sits S or more beyond its owner's block (committed positions; the list criterion
    // covers motion until the next rebuild), move every particle to its current block first. The
    // live Mindlin springs ride the migration pack (gid-keyed) exactly as in a rebalance.
    float e = 0.0f, d = 0.0f;
    driftVoteLocalKokkos(P, halo.blockBox(), 0.0f, e, d);
    int commSize = 1;
    MPI_Comm_size(halo.comm(), &commSize);
    if (allMax(e) >= slack && commSize > 1) {
      halo.migrateToBlocks(P);
      ++P.splitStats.driftMigrations;
      // The owned force accumulators are zeroed as they are consumed (integrate), the ghost ones
      // only at the start of each step (clearGhostScratch) and then left holding the discarded
      // ghost halves. A migration turns former ghost slots into owned ones: clear the owned range
      // so no stale ghost force is integrated (measured: dP 1e-8 -> 9e-4 at the first migration).
      zeroForceScratchKokkos(P.deltaVel, P.deltaAngVel, 0, P.numReal);
    }
    halo.invalidateTopology();  // fresh band + fresh positions at every pair rebuild
    halo.gather(P, band);
    fillWorldRadiiKokkos(P.scale, P.rad, P.globalScale, P.baseRadius, P.numParticles);
  }
  void refreshGhostState(Particles& P, bool needQuat) const {
    halo.forwardPositions(P.pos);
    halo.forward(P.vel);
    halo.forward(P.angVel);
    if (needQuat)
      halo.forward4(P.quat);
  }
  void clearGhostScratch(Particles& P) const {
    zeroForceScratchKokkos(P.deltaVel, P.deltaAngVel, P.numReal, P.numReal + halo.numGhost());
  }
};
 
inline void demStepHertzMpi(Particles& P, ParticleHalo& halo, float dt, int nsteps,
                            float skinFrac) {
  // Ghost band = worst-case pair cutoff (2 R_max,global) + skin + drift slack S (§5.5). The
  // driver's skin is skinFrac * R_min,global <= skinFrac * R_max,global, so this band bounds it;
  // S covers a partner up to S beyond its owner's block (the rebuild vote keeps it there).
  float maxR = P.numReal > 0 ? maxOwnedRadius(P) : 0.0f, maxRg = maxR;
  MPI_Allreduce(&maxR, &maxRg, 1, MPI_FLOAT, MPI_MAX, halo.comm());
  const float slack = kDriftSlack * maxRg;
  const double band = (2.0 + skinFrac) * static_cast<double>(maxRg) + static_cast<double>(slack);
  demStepForce(P, dt, nsteps, skinFrac, HertzMindlinLaw{}, MpiForceHooks{halo, band, slack});
  P.numParticles = P.numReal;  // restore owned-only active count for getters
}
 
}  // namespace peclet::dem
 
#endif  // PECLET_DEM_MPI
#endif  // DEM_STEP_SOLVE_MPI_HPP