MPI run on GPU fails

[dineshadepu]

Hi all,

I am looking into MPI and Cabana to do some simulations, which need MPI from cabana side, and for further expansion.

I had followed similar way of using particle communication as CabanaPD does, and was getting into runtime error.

So I tried to run PowderFill using the similar commands. CabanaPD is also failing for a mpi run on a single GPU.

Here is the log:

 

dineshadepu@dwi199a (main) /home/dineshadepu/life/softwares/ecp/CabanaPD/build $  
|  lab desktop => mpirun -n 2 ./examples/dem/PowderFill ../examples/dem/inputs/powder_fill.json 
Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

MPI detected: For OpenMP binding to work as intended, MPI ranks must be bound to exclusive CPU sets.

Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

MPI detected: For OpenMP binding to work as intended, MPI ranks must be bound to exclusive CPU sets.

Kokkos::Cuda::initialize WARNING: running kernels compiled for compute capability 10.0 on device with compute capability 12.0 , this will likely reduce potential performance.
Kokkos::Cuda::initialize WARNING: running kernels compiled for compute capability 10.0 on device with compute capability 12.0 , this will likely reduce potential performance.
[dwi199a:44403:0:44403] Caught signal 11 (Segmentation fault: invalid permissions for mapped object at address 0x7f72cedecc80)
[dwi199a:44402:0:44402] Caught signal 11 (Segmentation fault: invalid permissions for mapped object at address 0x7fb6eeded880)
==== backtrace (tid:  44403) ====
 0  /lib64/libucs.so.0(ucs_handle_error+0x304) [0x7f74fa3c35b4]
 1  /lib64/libucs.so.0(+0x2c691) [0x7f74fa3c6691]
 2  /lib64/libucs.so.0(+0x2c85a) [0x7f74fa3c685a]
 3  /lib64/libc.so.6(+0x41000) [0x7f74f9469000]
 4  /lib64/libc.so.6(+0x16e80d) [0x7f74f959680d]
 5  /usr/lib64/openmpi/lib/libopen-pal.so.80(+0xa5b47) [0x7f74f99a6b47]
 6  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_send_request_schedule_once+0x2d1) [0x7f74fa2620c1]
 7  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_recv_frag_callback_ack+0x115) [0x7f74fa26a4d5]
 8  /usr/lib64/openmpi/lib/libopen-pal.so.80(mca_btl_sm_poll_handle_frag+0x9a) [0x7f74f999df3a]
 9  /usr/lib64/openmpi/lib/libopen-pal.so.80(+0x9fb1c) [0x7f74f99a0b1c]
10  /usr/lib64/openmpi/lib/libopen-pal.so.80(opal_progress+0x34) [0x7f74f992c194]
11  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_send+0x490) [0x7f74fa25d9a0]
12  /usr/lib64/openmpi/lib/libmpi.so.40(MPI_Send+0x21a) [0x7f74fa0e153a]
13  ./examples/dem/PowderFill() [0x41bc82]
14  ./examples/dem/PowderFill() [0x4bbee6]
15  ./examples/dem/PowderFill() [0x4c1167]
16  ./examples/dem/PowderFill() [0x4c2c23]
17  ./examples/dem/PowderFill() [0x428760]
18  ./examples/dem/PowderFill() [0x40e4f8]
19  /lib64/libc.so.6(+0x2a58e) [0x7f74f945258e]
20  /lib64/libc.so.6(__libc_start_main+0x89) [0x7f74f9452649]
21  ./examples/dem/PowderFill() [0x40fa75]
=================================
==== backtrace (tid:  44402) ====
 0  /lib64/libucs.so.0(ucs_handle_error+0x304) [0x7fb91fcf35b4]
 1  /lib64/libucs.so.0(+0x2c691) [0x7fb91fcf6691]
 2  /lib64/libucs.so.0(+0x2c85a) [0x7fb91fcf685a]
 3  /lib64/libc.so.6(+0x41000) [0x7fb918a69000]
 4  /lib64/libc.so.6(+0x16e80d) [0x7fb918b9680d]
 5  /usr/lib64/openmpi/lib/libopen-pal.so.80(+0xa5b47) [0x7fb919da6b47]
 6  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_send_request_schedule_once+0x2d1) [0x7fb9198620c1]
 7  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_recv_frag_callback_ack+0x115) [0x7fb91986a4d5]
 8  /usr/lib64/openmpi/lib/libopen-pal.so.80(mca_btl_sm_poll_handle_frag+0x9a) [0x7fb919d9df3a]
 9  /usr/lib64/openmpi/lib/libopen-pal.so.80(+0x9fb1c) [0x7fb919da0b1c]
10  /usr/lib64/openmpi/lib/libopen-pal.so.80(opal_progress+0x34) [0x7fb919d2c194]
11  /usr/lib64/openmpi/lib/libmpi.so.40(mca_pml_ob1_send+0x490) [0x7fb91985d9a0]
12  /usr/lib64/openmpi/lib/libmpi.so.40(MPI_Send+0x21a) [0x7fb9196e153a]
13  ./examples/dem/PowderFill() [0x41bc82]
14  ./examples/dem/PowderFill() [0x4bbee6]
15  ./examples/dem/PowderFill() [0x4c1167]
16  ./examples/dem/PowderFill() [0x4c2c23]
17  ./examples/dem/PowderFill() [0x428760]
18  ./examples/dem/PowderFill() [0x40e4f8]
19  /lib64/libc.so.6(+0x2a58e) [0x7fb918a5258e]
20  /lib64/libc.so.6(__libc_start_main+0x89) [0x7fb918a52649]
21  ./examples/dem/PowderFill() [0x40fa75]
=================================

It works nicely on single run though:

dineshadepu@dwi199a (main) /home/dineshadepu/life/softwares/ecp/CabanaPD/build $  
|  lab desktop => ./examples/dem/PowderFill ../examples/dem/inputs/powder_fill.json 
Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

Kokkos::Cuda::initialize WARNING: running kernels compiled for compute capability 10.0 on device with compute capability 12.0 , this will likely reduce potential performance.
Local particles: 19141, Maximum neighbors: 21
#Timestep/Total-steps Simulation-time
1000/250000 2.00e-03
2000/250000 4.00e-03
3000/250000 6.00e-03
4000/250000 8.00e-03
5000/250000 1.00e-02
6000/250000 1.20e-02
7000/250000 1.40e-02
8000/250000 1.60e-02
9000/250000 1.80e-02
10000/250000 2.00e-02
11000/250000 2.20e-02
12000/250000 2.40e-02
13000/250000 2.60e-02
14000/250000 2.80e-02

Is this something only I am facing, or you guys are able to get it working?

I am using Rocky machine, with

|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 5080        On  |   00000000:01:00.0  On |                  N/A |
|  0%   49C    P0             48W /  360W |    1601MiB /  16303MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

Any idea where I am going wrong, or is there a specific way of installing it to make it work?

Thank you very much.

[dineshadepu]

I don't want to clutter the conversation, but I think it will help if I put what I am trying to do. Here is the code:

#include <Cabana_Core.hpp>
#include <Cabana_Grid.hpp>
#include <Kokkos_Core.hpp>
#include <cmath>
#include <iostream>
#include <vector>
#include <optional>
#include <random>
#include "neighbors_particles.hpp"
// #include "particle_comm_mpi_dem.hpp"


#include <mpi.h>

template<class CoreState>
struct CoreParticleComm
{
  using memory_space = typename CoreState::memory_space;
  using halo_type = Cabana::Halo<memory_space>;

  std::shared_ptr<halo_type> halo;

  CoreState& particles;
  int halo_width;

  CoreParticleComm(CoreState& p, int hw)
    : particles(p), halo_width(hw) {}

  void migrate()
  {
    auto& local_grid = particles.local_grid;

    Cabana::Grid::particleMigrate(
      *local_grid,
      particles.slicePosition(),
      particles.plist.aosoa(),
      halo_width);

    particles.resize(
                     particles.plist.size(),
                     0);
  }
};


template<class MemorySpace>
struct PhysicsIndexLists
{
  using view_type = Kokkos::View<int*, MemorySpace>;

  view_type owned;
  view_type free;
  view_type boundary;

  int owned_count    = 0;
  int free_count     = 0;
  int boundary_count = 0;

  PhysicsIndexLists(std::size_t max_particles)
    : owned("owned_indices", max_particles),
      free("free_indices", max_particles),
      boundary("boundary_indices", max_particles)
  {}
};

// ============================================================
// Field Tags for ParticleList
// ============================================================
struct Position
{
  using data_type = double[3];

  static std::string label()
  {
    return "position";
  }
};


// ============================================================
// CoreParticlesState
// ============================================================

template<class MemorySpace>
struct CoreParticlesState
{
  using memory_space = MemorySpace;
  using execution_space = typename memory_space::execution_space;

  static constexpr int Dim = 3;

  // ============================================================
  // 1️⃣ MPI Backbone — ParticleList (Position + ID only)
  // ============================================================

using plist_type =
  Cabana::Grid::ParticleList<
    memory_space,
    1,
    Position>;

  plist_type plist;

  // ============================================================
  // 2️⃣ Physics storage (ALL other data)
  // ============================================================

  using CoreTypes = Cabana::MemberTypes<
    double[3], // velocity
    double[3], // acceleration
    double[3], // force
    double[3], // torque
    double,    // mass
    double,    // cutoff radius
    double,    // radius
    int,       // is boundary
    int,       // is frozen
    uint64_t   // global id
  >;

  using AoSoA_t = Cabana::AoSoA<CoreTypes, MemorySpace>;

  AoSoA_t data;

  // ============================================================
  // Ownership
  // ============================================================

  int num_local = 0;
  int num_ghost = 0;

  std::shared_ptr<
    Cabana::Grid::LocalGrid<
      Cabana::Grid::UniformMesh<double,3>>> local_grid;

  // ============================================================
  // Constructor
  // ============================================================
CoreParticlesState(std::size_t n_initial)
  : plist("plist"),
    data("core_data", n_initial),
    num_local(n_initial),
    num_ghost(0)
{
  plist.resize(n_initial);
}

  // ============================================================
  // Resize
  // ============================================================

  void resize(int n_local, int n_ghost)
  {
    num_local = n_local;
    num_ghost = n_ghost;

    plist.resize(n_local + n_ghost);
    data.resize(n_local + n_ghost);
  }

  // ============================================================
  // Accessors
  // ============================================================

  // Position comes from ParticleList
  auto slicePosition()
  {
    return plist.slice(Position());
  }

  auto sliceVelocity()      { return Cabana::slice<0>(data, "velocity"); }
  auto sliceAcceleration()  { return Cabana::slice<1>(data, "acceleration"); }
  auto sliceForce()         { return Cabana::slice<2>(data, "force"); }
  auto sliceTorque()        { return Cabana::slice<3>(data, "torque"); }
  auto sliceMass()          { return Cabana::slice<4>(data, "mass"); }
  auto sliceCutoffRadius()  { return Cabana::slice<5>(data, "cutoff_radius"); }
  auto sliceRadius()        { return Cabana::slice<6>(data, "radius"); }
  auto sliceIsBoundary()    { return Cabana::slice<7>(data, "is_boundary"); }
  auto sliceIsFrozen()      { return Cabana::slice<8>(data, "is_frozen"); }
  auto sliceGlobalId()      { return Cabana::slice<9>(data, "global_id"); }

  // Position comes from ParticleList
  auto slicePosition() const
  {
    return plist.slice(Position());
  }

  auto sliceVelocity()     const { return Cabana::slice<0>(data, "velocity"); }
  auto sliceAcceleration() const { return Cabana::slice<1>(data, "acceleration"); }
  auto sliceForce()        const { return Cabana::slice<2>(data, "force"); }
  auto sliceTorque()       const { return Cabana::slice<3>(data, "torque"); }
  auto sliceMass()         const { return Cabana::slice<4>(data, "mass"); }
  auto sliceCutoffRadius() const { return Cabana::slice<5>(data, "cutoff_radius"); }
  auto sliceRadius()       const { return Cabana::slice<6>(data, "radius"); }
  auto sliceIsBoundary()   const { return Cabana::slice<7>(data, "is_boundary"); }
  auto sliceIsFrozen()     const { return Cabana::slice<8>(data, "is_frozen"); }
  auto sliceGlobalId()     const { return Cabana::slice<9>(data, "global_id"); }

  int numLocal() const { return num_local; }
  int numGhost() const { return num_ghost; }


  template<class ExecSpace,class Functor,class IndexView>
  void iterate_indices(ExecSpace,
                       const Functor& f,
                       const IndexView& indices,
                       std::size_t count,
                       const std::string& label) const
  {
    Kokkos::parallel_for(
                         label,
                         Kokkos::RangePolicy<ExecSpace>(0,count),
                         KOKKOS_LAMBDA(int k){
                           f(indices(k));
                         });
  }
};


// ============================================================
// Leapfrog Integrator (velocity update)
// ============================================================

template<class CoreState>
struct LeapfrogStage1
{
  // ----------------------------
  // Device-safe slices
  // ----------------------------
  decltype(std::declval<CoreState>().sliceVelocity()) u;
  decltype(std::declval<CoreState>().sliceForce())    force;
  decltype(std::declval<CoreState>().sliceTorque())  torque;
  decltype(std::declval<CoreState>().sliceMass())    m;

  double half_dt;

  // ----------------------------
  // Constructor
  // ----------------------------
  LeapfrogStage1(const CoreState& core,
                 double half_dt_)
    : u(core.sliceVelocity())
    , force(core.sliceForce())
    , m(core.sliceMass())
    , half_dt(half_dt_)
  {}

  // ----------------------------
  // Per-particle update
  // ----------------------------
  KOKKOS_INLINE_FUNCTION
  void operator()(const int i) const
  {
    const double inv_m = 1.0 / m(i);

    u(i,0) += half_dt * force(i,0) * inv_m;
    u(i,1) += half_dt * force(i,1) * inv_m;
    u(i,2) += half_dt * force(i,2) * inv_m;
  }
};

template<class CoreState>
struct LeapfrogStage2
{
  // ----------------------------
  // Device-safe slices
  // ----------------------------
  decltype(std::declval<CoreState>().slicePosition()) x;
  decltype(std::declval<CoreState>().sliceVelocity()) u;

  double dt;

  // ----------------------------
  // Constructor
  // ----------------------------
  LeapfrogStage2(const CoreState& core,
                 double dt_)
    : x(core.slicePosition())
    , u(core.sliceVelocity())
    , dt(dt_)
  {}

  // ----------------------------
  // Per-particle update
  // ----------------------------
  KOKKOS_INLINE_FUNCTION
  void operator()(const int i) const
  {
    x(i,0) += dt * u(i,0);
    x(i,1) += dt * u(i,1);
    x(i,2) += dt * u(i,2);
  }
};



template<class CoreState>
struct LeapfrogStage3
{
  // ----------------------------
  // Device-safe slices
  // ----------------------------
  decltype(std::declval<CoreState>().sliceVelocity()) u;
  decltype(std::declval<CoreState>().sliceForce())    force;
  decltype(std::declval<CoreState>().sliceTorque())  torque;
  decltype(std::declval<CoreState>().sliceMass())    m;

  double half_dt;

  // ----------------------------
  // Constructor
  // ----------------------------
  LeapfrogStage3(const CoreState& core,
                 double half_dt_)
    : u(core.sliceVelocity())
    , force(core.sliceForce())
    , m(core.sliceMass())
    , half_dt(half_dt_)
  {}

  // ----------------------------
  // Per-particle update
  // ----------------------------
  KOKKOS_INLINE_FUNCTION
  void operator()(const int i) const
  {
    const double inv_m = 1.0 / m(i);

    u(i,0) += half_dt * force(i,0) * inv_m;
    u(i,1) += half_dt * force(i,1) * inv_m;
    u(i,2) += half_dt * force(i,2) * inv_m;
  }
};

// ============================================================
// MAIN
// ============================================================

int main(int argc,char* argv[])
{
  MPI_Init(&argc,&argv);
  {

    Kokkos::ScopeGuard guard(argc,argv);

    using ExecSpace = Kokkos::DefaultExecutionSpace;
    // using ExecSpace = Kokkos::OpenMP;
    using MemorySpace = typename ExecSpace::memory_space;


    const double R=2.0;
    const double H=4.0;
    const double wall_rad=0.2;
    const double spacing=2.0*wall_rad*0.95;

    std::vector<double> hx,hu,hf,ht,hm,hr;
    std::vector<int> hb,hfrozen;
    std::vector<uint64_t> hid;

    // ---------------- floor ----------------
    int nr = R/spacing;
    for(int ir=0;ir<nr;++ir)
      {
        double rxy=ir*spacing;
        int ntheta=std::max(1,int(2*M_PI*rxy/spacing));

        for(int i=0;i<ntheta;++i)
          {
            double th=2*M_PI*i/ntheta;
            hx.insert(hx.end(),{rxy*cos(th),rxy*sin(th),0});
            hu.insert(hu.end(),{0,0,0});
            hf.insert(hf.end(),{0,0,0});
            ht.insert(ht.end(),{0,0,0});
            hm.push_back(0);
            hr.push_back(wall_rad);
            hb.push_back(1);
            hfrozen.push_back(1);
            hid.push_back(hid.size());
          }
      }

    // ---------------- cylindrical wall ----------------
    int nz = H / spacing;
    int ntheta_wall = std::max(1, int(2 * M_PI * R / spacing));

    for (int k = 0; k < nz; ++k)
      {
        double z = k * spacing;

        for (int i = 0; i < ntheta_wall; ++i)
          {
            double th = 2 * M_PI * i / ntheta_wall;

            double xw = R * cos(th);
            double yw = R * sin(th);

            hx.insert(hx.end(), {xw, yw, z});
            hu.insert(hu.end(), {0, 0, 0});
            hf.insert(hf.end(), {0, 0, 0});
            ht.insert(ht.end(), {0, 0, 0});

            hm.push_back(0.0);          // infinite mass
            hr.push_back(wall_rad);

            hb.push_back(1);            // boundary
            hfrozen.push_back(1);       // frozen

            hid.push_back(hid.size());
          }
      }

    // ---------------- falling balls ----------------
    std::mt19937 gen(0);
    std::uniform_real_distribution<double> dist(0,1);

    const int n_ball=200;
    const double ball_rad=0.15;
    const double min_dist=2.1*ball_rad;

    std::vector<std::array<double,3>> centers;

    while(centers.size()<n_ball)
      {
        double rxy=sqrt(dist(gen))*(R-3*ball_rad);
        double th=2*M_PI*dist(gen);

        double x=rxy*cos(th);
        double y=rxy*sin(th);
        double z=0.1* H+1.0*dist(gen);

        bool ok=true;
        for(auto& c:centers)
          {
            double dx=x-c[0];
            double dy=y-c[1];
            double dz=z-c[2];
            if(dx*dx+dy*dy+dz*dz<min_dist*min_dist){
              ok=false;break;
            }
          }
        if(!ok) continue;

        centers.push_back({x,y,z});
        hx.insert(hx.end(),{x,y,z});
        hu.insert(hu.end(),{0,0,0});
        hf.insert(hf.end(),{0,0,0});
        ht.insert(ht.end(),{0,0,0});
        hm.push_back(1.0);
        hr.push_back(ball_rad);
        hb.push_back(0);
        hfrozen.push_back(0);
        hid.push_back(hid.size());
      }

    // =======================
    // Create the grid for MPI
    // =======================
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    constexpr int Dim = 3;

    // Global domain (must match mesh bounds!)
    std::array<double,Dim> global_lo = {-R-1, -R-1, -H};
    std::array<double,Dim> global_hi = { R+1,  R+1,  H+3 };

    // Cell size ≈ cutoff
    double cutoff = 3. * wall_rad;
    double cell_size = cutoff;

    std::array<int,Dim> num_cells = {
      int((global_hi[0]-global_lo[0]) / cell_size),
      int((global_hi[1]-global_lo[1]) / cell_size),
      int((global_hi[2]-global_lo[2]) / cell_size)
    };

    auto global_mesh =
      Cabana::Grid::createUniformGlobalMesh(
                                            global_lo, global_hi, num_cells);

    std::array<bool,Dim> periodic = {false,false,false};

    Cabana::Grid::DimBlockPartitioner<Dim> partitioner;

    auto global_grid =
      Cabana::Grid::createGlobalGrid(
                                     MPI_COMM_WORLD,
                                     global_mesh,
                                     periodic,
                                     partitioner);

    int halo_width = 1;

    auto local_grid =
      Cabana::Grid::createLocalGrid(global_grid, halo_width);
    // ==================================
    // Create the grid for MPI ends
    // ==================================

    // ---------------- create core  on host ----------------
    using HostMemorySpace = Kokkos::HostSpace;

    CoreParticlesState<HostMemorySpace> host_core(hm.size());
    {
      auto x  = host_core.slicePosition();
      auto u  = host_core.sliceVelocity();
      auto a  = host_core.sliceAcceleration();
      auto f  = host_core.sliceForce();
      auto t  = host_core.sliceTorque();
      auto m  = host_core.sliceMass();
      auto rc = host_core.sliceCutoffRadius();
      auto rad= host_core.sliceRadius();
      auto isb= host_core.sliceIsBoundary();
      auto isf= host_core.sliceIsFrozen();
      auto pid= host_core.sliceGlobalId();
      for (std::size_t i = 0; i < hm.size(); ++i)
        {
          x(i,0) = hx[3*i+0];
          x(i,1) = hx[3*i+1];
          x(i,2) = hx[3*i+2];

          u(i,0)=u(i,1)=u(i,2)=0.0;
          a(i,0)=a(i,1)=a(i,2)=0.0;
          f(i,0)=f(i,1)=f(i,2)=0.0;
          t(i,0)=t(i,1)=t(i,2)=0.0;

          m(i)   = hm[i];
          rc(i)  = 3. * hr[i];
          rad(i) = hr[i];

          isb(i) = hb[i];
          isf(i) = hfrozen[i];
          pid(i) = hid[i];
        }
    }

    // ---------------- create core on the device ----------------
    CoreParticlesState<MemorySpace> core(hm.size());

    // copy physics
    Cabana::deep_copy(core.data, host_core.data);

    // copy positions
    Cabana::deep_copy(core.plist, host_core.plist);

    // attach grid
    core.local_grid = local_grid;

    // =======================
    // Migration
    // =======================
    Cabana::Grid::particleMigrate(
                                  *local_grid,
                                  core.slicePosition(),
                                  core.plist.aosoa(),
                                  halo_width);


    // After migration:
    core.resize(core.plist.size(), 0);
    // =======================
    // Migration ends
    // =======================

    // =======================
    // Build Halo
    // =======================
    CoreParticleComm<CoreParticlesState<MemorySpace>> comm(core, halo_width);
    // comm.build();
    // =======================
    // Build Halo ends
    // =======================

    auto x=core.slicePosition();
    auto u=core.sliceVelocity();
    auto f=core.sliceForce();
    auto m=core.sliceMass();
    auto rc=core.sliceCutoffRadius();
    auto isb=core.sliceIsBoundary();
    auto isf=core.sliceIsFrozen();
    auto pid=core.sliceGlobalId();

    // auto local_mesh =
    //   Cabana::Grid::createLocalMesh<Kokkos::HostSpace>(*local_grid);

    // double lo[3];
    // double hi[3];

    // for(int d=0; d<3; ++d)
    //   {
    //     lo[d] = local_mesh.lowCorner(Cabana::Grid::Ghost(), d);
    //     hi[d] = local_mesh.highCorner(Cabana::Grid::Ghost(), d);
    //   }

    // core.set_mesh_bounds(lo, hi);

    // core.compute_max_cutoff(ExecSpace{});
    // core.create_neighbors();

    // =======================
    // build free indices
    // =======================

    PhysicsIndexLists<MemorySpace> indices(core.numLocal());

    auto owned_h    = Kokkos::create_mirror_view(indices.owned);
    auto free_h     = Kokkos::create_mirror_view(indices.free);
    auto boundary_h = Kokkos::create_mirror_view(indices.boundary);

    indices.owned_count = 0;
    indices.free_count = 0;
    indices.boundary_count = 0;

    for (int i = 0; i < core.numLocal(); ++i)
      {
        owned_h(indices.owned_count++) = i;

        if (isb(i))
          boundary_h(indices.boundary_count++) = i;
        else
          free_h(indices.free_count++) = i;
      }

    Kokkos::deep_copy(indices.owned, owned_h);
    Kokkos::deep_copy(indices.free, free_h);
    Kokkos::deep_copy(indices.boundary, boundary_h);
    // =======================
    // build free indices ends
    // =======================

    const double dt = 1e-4;
    const double kn = 1e6;
    const int steps = 10000;

    // NormalContact contact(core, kn);
    LeapfrogStage1      stage1(core, 0.5 * dt);
    LeapfrogStage2      stage2(core, dt);
    LeapfrogStage3      stage3(core, 0.5 * dt);

    Cabana::Experimental::HDF5ParticleOutput::HDF5Config h5_config;   // no collective MPI

    Cabana::Experimental::HDF5ParticleOutput::writeTimeStep(
                                                            h5_config,
                                                            "particles",
                                                            MPI_COMM_WORLD,   // single rank
                                                            0,
                                                            0.,
                                                            core.numLocal(),
                                                            core.slicePosition(),
                                                            core.sliceVelocity(),
                                                            core.sliceForce(),
                                                            core.sliceMass(),
                                                            core.sliceRadius(),
                                                            core.sliceIsBoundary()
                                                            );

    for (int step = 0; step < steps; ++step)
      {
        LeapfrogStage1      stage1(core, 0.5 * dt);
        core.iterate_indices(
                             ExecSpace{},
                             stage1,
                             indices.free,
                             indices.free_count,
                             "stage1");

        LeapfrogStage2      stage2(core, dt);
        core.iterate_indices(
                             ExecSpace{},
                             stage2,
                             indices.free,
                             indices.free_count,
                             "stage2");

        comm.migrate();
        // comm.build();
        // comm.gather();
        // core.resize();
        // =======================
        // build free indices
        // =======================

        PhysicsIndexLists<MemorySpace> indices(core.numLocal());

        auto owned_h    = Kokkos::create_mirror_view(indices.owned);
        auto free_h     = Kokkos::create_mirror_view(indices.free);
        auto boundary_h = Kokkos::create_mirror_view(indices.boundary);

        indices.owned_count = 0;
        indices.free_count = 0;
        indices.boundary_count = 0;

        for (int i = 0; i < core.numLocal(); ++i)
          {
            owned_h(indices.owned_count++) = i;

            if (isb(i))
              boundary_h(indices.boundary_count++) = i;
            else
              free_h(indices.free_count++) = i;
          }

        Kokkos::deep_copy(indices.owned, owned_h);
        Kokkos::deep_copy(indices.free, free_h);
        Kokkos::deep_copy(indices.boundary, boundary_h);
        // =======================
        // build free indices ends
        // =======================

        auto x=core.slicePosition();
        auto u=core.sliceVelocity();
        auto f=core.sliceForce();
        auto m=core.sliceMass();
        auto rc=core.sliceCutoffRadius();
        auto isb=core.sliceIsBoundary();
        auto isf=core.sliceIsFrozen();
        auto pid=core.sliceGlobalId();

        // 1️⃣ Reset forces (free particles only)
        core.iterate_indices(
                             ExecSpace{},
                             KOKKOS_LAMBDA(int i){
                               f(i,0)=0; f(i,1)=0; f(i,2)=0;
                             },
                             indices.free,
                             indices.free_count,
                             "reset_force");

        // 2️⃣ Gravity
        core.iterate_indices(
                             ExecSpace{},
                             KOKKOS_LAMBDA(int i){
                               f(i,2) -= 9.81 * m(i);
                             },
                             indices.free,
                             indices.free_count,
                             "gravity");

        LeapfrogStage3      stage3(core, 0.5 * dt);
        core.iterate_indices(
                             ExecSpace{},
                             stage3,
                             indices.free,
                             indices.free_count,
                             "stage3");

        if (step % 200 == 0)
          {
            double time = step * dt;

            Cabana::Experimental::HDF5ParticleOutput::writeTimeStep(
                                                                    h5_config,
                                                                    "particles",
                                                                    MPI_COMM_WORLD,   // single rank
                                                                    step,
                                                                    time,
                                                                    core.numLocal(),
                                                                    core.slicePosition(),
                                                                    core.sliceVelocity(),
                                                                    core.sliceAcceleration(),
                                                                    core.sliceForce(),
                                                                    core.sliceMass(),
                                                                    core.sliceRadius(),
                                                                    core.sliceIsBoundary()
                                                                    );
            std::cout << "Step " << step << "\n";
          }

      }
    std::cout<<"Rain in tank complete.\n";

  }
  MPI_Finalize();
  return 0;
}