Ascent: SoA Particle Support (AMReX-Codes#3350)

## Summary Add support for pure SoA layouted particle containers for Ascent. ## Additional background Follow-up to AMReX-Codes#2878. ## Checklist The proposed changes: - [ ] fix a bug or incorrect behavior in AMReX - [x] add new capabilities to AMReX - [ ] changes answers in the test suite to more than roundoff level - [ ] are likely to significantly affect the results of downstream AMReX users - [ ] include documentation in the code and/or rst files, if appropriate --------- Co-authored-by: Andrew Myers <[email protected]>
guj · Dec 13, 2023 · 7d689d6 · 7d689d6
1 parent 839836f
commit 7d689d6
Show file tree

Hide file tree

Showing 6 changed files with 491 additions and 74 deletions.
diff --git a/Src/Extern/Conduit/AMReX_Conduit_Blueprint.H b/Src/Extern/Conduit/AMReX_Conduit_Blueprint.H
@@ -96,9 +96,8 @@ namespace amrex
  // coordset and fields used to represent the passed particle container.
  // This allows you to use unique names to wrap multiple particle containers
  // into a single blueprint tree.
- template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
- void ParticleContainerToBlueprint (const ParticleContainer<NStructReal,
- NStructInt,
+ template <typename ParticleType, int NArrayReal, int NArrayInt>
+ void ParticleContainerToBlueprint (const ParticleContainer_impl<ParticleType,
  NArrayReal,
  NArrayInt> &pc,
  const Vector<std::string> &real_comp_names,

diff --git a/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H b/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H
@@ -20,26 +20,21 @@ namespace amrex
 // Note:
 // This is a helper function, it's not part of the AMReX Blueprint Interface.
 //---------------------------------------------------------------------------//
-template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
+template <typename ParticleType, int NArrayReal, int NArrayInt>
 void
-ParticleTileToBlueprint(const ParticleTile<amrex::Particle<NStructReal,
- NStructInt>,
+ParticleTileToBlueprint(const ParticleTile<ParticleType,
  NArrayReal,
  NArrayInt> &ptile,
  const Vector<std::string> &real_comp_names,
  const Vector<std::string> &int_comp_names,
  conduit::Node &res,
  const std::string &topology_name)
 {
- int num_particles = ptile.GetArrayOfStructs().size();
- int struct_size = sizeof(Particle<NStructReal, NStructInt>);
+ int num_particles = ptile.size();
 
  // knowing the above, we can zero copy the x,y,z positions + id, cpu
  // and any user fields in the AOS
 
- // get the first particle's struct
- const auto &pstruct = ptile.GetArrayOfStructs();
-
  // setup a blueprint description for the particle mesh
  // create a coordinate set
  std::string coordset_name = topology_name + "_coords";
@@ -63,29 +58,56 @@ ParticleTileToBlueprint(const ParticleTile<amrex::Particle<NStructReal,
  //----------------------------------//
  // point locations from from aos
  //----------------------------------//
+ char* pbuf = nullptr;
 
- const char* pbuf_const = reinterpret_cast<const char*>(pstruct.data());
- char* pbuf = const_cast<char*>(pbuf_const);
+ if constexpr(ParticleType::is_soa_particle)
+ {
+ amrex::ignore_unused(pbuf);
 
- ParticleReal* xp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
- n_coords["values/x"].set_external(xp,
- num_particles,
- 0,
- struct_size);
+ const auto &soa = ptile.GetStructOfArrays();
+
+ // for soa entries, we can use standard strides,
+ // since these are contiguous arrays
+
+ n_coords["values/x"].set_external(const_cast<ParticleReal*>(&soa.GetRealData(0)[0]),
+ num_particles);
 #if AMREX_SPACEDIM > 1
- ParticleReal* yp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
- n_coords["values/y"].set_external(yp,
- num_particles,
- 0,
- struct_size);
+ n_coords["values/y"].set_external(const_cast<ParticleReal*>(&soa.GetRealData(1)[0]),
+ num_particles);
 #endif
 #if AMREX_SPACEDIM > 2
- ParticleReal* zp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
- n_coords["values/z"].set_external(zp,
- num_particles,
- 0,
- struct_size);
+ n_coords["values/z"].set_external(const_cast<ParticleReal*>(&soa.GetRealData(2)[0]),
+ num_particles);
+#endif
+ } else
+ {
+ // get the first particle's struct
+ const auto &pstruct = ptile.GetArrayOfStructs();
+ const int struct_size = sizeof(ParticleType);
+
+ const char* pbuf_const = reinterpret_cast<const char*>(pstruct.data());
+ pbuf = const_cast<char*>(pbuf_const);
+
+ ParticleReal* xp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
+ n_coords["values/x"].set_external(xp,
+ num_particles,
+ 0,
+ struct_size);
+#if AMREX_SPACEDIM > 1
+ ParticleReal* yp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
+ n_coords["values/y"].set_external(yp,
+ num_particles,
+ 0,
+ struct_size);
+#endif
+#if AMREX_SPACEDIM > 2
+ ParticleReal* zp = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
+ n_coords["values/z"].set_external(zp,
+ num_particles,
+ 0,
+ struct_size);
 #endif
+ }
 
  // fields
  conduit::Node &n_fields = res["fields"];
@@ -95,65 +117,104 @@ ParticleTileToBlueprint(const ParticleTile<amrex::Particle<NStructReal,
  // -----------------------------
 
  int vname_real_idx = 0;
- // struct real fields, the first set are always the particle positions
- // which we wrap above
- for (int i = 0; i < NStructReal; i++)
+ if constexpr(!ParticleType::is_soa_particle)
  {
- ParticleReal* val = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
- conduit::Node &n_f = n_fields[real_comp_names.at(vname_real_idx)];
- n_f["topology"] = topology_name;
- n_f["association"] = "element";
- n_f["values"].set_external(val,
- num_particles,
- 0,
- struct_size);
+ constexpr int struct_size = sizeof(ParticleType);
+ constexpr int NStructReal = ParticleType::NReal;
 
- vname_real_idx++;
+ // struct real fields, the first set are always the particle positions
+ // which we wrap above
+ for (int i = 0; i < NStructReal; i++)
+ {
+ ParticleReal* val = reinterpret_cast<ParticleReal*>(pbuf); pbuf += sizeof(ParticleReal);
+ conduit::Node &n_f = n_fields[real_comp_names.at(vname_real_idx)];
+ n_f["topology"] = topology_name;
+ n_f["association"] = "element";
+ n_f["values"].set_external(val,
+ num_particles,
+ 0,
+ struct_size);
+
+ vname_real_idx++;
+ }
  }
 
  //----------------------------------//
  // standard integer fields from aos
  // (id, cpu)
  //----------------------------------//
 
- // id is the first int entry
- int* id = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
- conduit::Node &n_f_id = n_fields[topology_name + "_id"];
+ if constexpr(!ParticleType::is_soa_particle)
+ {
+ const int struct_size = sizeof(ParticleType);
+
+ // id is the first int entry
+ int* id = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
+ conduit::Node &n_f_id = n_fields[topology_name + "_id"];
+
+ n_f_id["topology"] = topology_name;
+ n_f_id["association"] = "element";
+ n_f_id["values"].set_external(id,
+ num_particles,
+ 0,
+ struct_size);
+
+ // cpu is the second int entry
+ int* cpu = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
+ conduit::Node &n_f_cpu = n_fields[topology_name + "_cpu"];
+
+ n_f_cpu["topology"] = topology_name;
+ n_f_cpu["association"] = "element";
+ n_f_cpu["values"].set_external(cpu,
+ num_particles,
+ 0,
+ struct_size);
+ } else {
+ const auto &soa = ptile.GetStructOfArrays();
+
+ // for soa entries, we can use standard strides,
+ // since these are contiguous arrays
 
- n_f_id["topology"] = topology_name;
- n_f_id["association"] = "element";
- n_f_id["values"].set_external(id,
- num_particles,
- 0,
- struct_size);
+ // id is the first int entry
+ conduit::Node &n_f_id = n_fields[topology_name + "_id"];
 
- // cpu is the second int entry
- int* cpu = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
- conduit::Node &n_f_cpu = n_fields[topology_name + "_cpu"];
+ n_f_id["topology"] = topology_name;
+ n_f_id["association"] = "element";
+ n_f_id["values"].set_external(const_cast<int*>(&soa.GetIntData(0)[0]),
+ num_particles);
 
- n_f_cpu["topology"] = topology_name;
- n_f_cpu["association"] = "element";
- n_f_cpu["values"].set_external(cpu,
- num_particles,
- 0,
- struct_size);
+ // cpu is the second int entry
+ conduit::Node &n_f_cpu = n_fields[topology_name + "_cpu"];
+
+ n_f_cpu["topology"] = topology_name;
+ n_f_cpu["association"] = "element";
+ n_f_cpu["values"].set_external(const_cast<int*>(&soa.GetIntData(0)[0]),
+ num_particles);
+
+ }
 
  // --------------------------------
  // user defined, integer aos fields
  // --------------------------------
 
  int vname_int_idx = 0;
- for (int i = 0; i < NStructInt; i++)
+ if constexpr(!ParticleType::is_soa_particle)
  {
- int* val = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
- conduit::Node &n_f = n_fields[int_comp_names.at(vname_int_idx)];
- n_f["topology"] = topology_name;
- n_f["association"] = "element";
- n_f["values"].set_external(val,
- num_particles,
- 0,
- struct_size);
- vname_int_idx++;
+ constexpr int struct_size = sizeof(ParticleType);
+ constexpr int NStructInt = ParticleType::NInt;
+
+ for (int i = 0; i < NStructInt; i++)
+ {
+ int* val = reinterpret_cast<int*>(pbuf); pbuf += sizeof(int);
+ conduit::Node &n_f = n_fields[int_comp_names.at(vname_int_idx)];
+ n_f["topology"] = topology_name;
+ n_f["association"] = "element";
+ n_f["values"].set_external(val,
+ num_particles,
+ 0,
+ struct_size);
+ vname_int_idx++;
+ }
  }
 
  // -------------------------
@@ -193,10 +254,9 @@ ParticleTileToBlueprint(const ParticleTile<amrex::Particle<NStructReal,
 //---------------------------------------------------------------------------//
 // Converts a AMReX Particle Container into a Conduit Mesh Blueprint Hierarchy.
 //---------------------------------------------------------------------------//
-template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
+template <typename ParticleType, int NArrayReal, int NArrayInt>
 void
-ParticleContainerToBlueprint(const ParticleContainer<NStructReal,
- NStructInt,
+ParticleContainerToBlueprint(const ParticleContainer_impl<ParticleType,
  NArrayReal,
  NArrayInt> &pc,
  const Vector<std::string> &real_comp_names,
@@ -209,8 +269,13 @@ ParticleContainerToBlueprint(const ParticleContainer<NStructReal,
  // validate varnames, which are used to provide field names
  // for user defined aos and soa values.
 
- BL_ASSERT(real_comp_names.size() == (NStructReal + NArrayReal) );
- BL_ASSERT(int_comp_names.size() == (NStructInt + NArrayInt) );
+ if constexpr(ParticleType::is_soa_particle) {
+ BL_ASSERT(real_comp_names.size() == NArrayReal);
+ BL_ASSERT(int_comp_names.size() == NArrayInt);
+ } else {
+ BL_ASSERT(real_comp_names.size() == (ParticleType::NReal + NArrayReal) );
+ BL_ASSERT(int_comp_names.size() == (ParticleType::NInt + NArrayInt) );
+ }
 
  int num_levels = pc.maxLevel() + 1;
  int num_domains = 0;
@@ -224,7 +289,7 @@ ParticleContainerToBlueprint(const ParticleContainer<NStructReal,
  int rank = ParallelDescriptor::MyProc();
  int nprocs = ParallelDescriptor::NProcs();
 
- using MyParConstIter = ParConstIter<NStructReal, NStructInt, NArrayReal, NArrayInt>;
+ using MyParConstIter = ParConstIter_impl<ParticleType, NArrayReal, NArrayInt>;
 
  //
  // blueprint expects unique ids for each domain published

diff --git a/Tests/Particles/Ascent_Insitu_SOA/CMakeLists.txt b/Tests/Particles/Ascent_Insitu_SOA/CMakeLists.txt
@@ -0,0 +1,13 @@
+if ( NOT AMReX_ASCENT )
+ return ()
+endif ()
+
+foreach(D IN LISTS AMReX_SPACEDIM)
+ set(_sources main.cpp)
+ set(_input_files inputs.rt )
+
+ setup_test(${D} _sources _input_files NTASKS 2)
+
+ unset(_sources)
+ unset(_input_files)
+endforeach()
diff --git a/Tests/Particles/Ascent_Insitu_SOA/GNUmakefile b/Tests/Particles/Ascent_Insitu_SOA/GNUmakefile
@@ -0,0 +1,24 @@
+AMREX_HOME = ../../../
+
+DEBUG = FALSE
+
+DIM = 3
+
+COMP = gcc
+
+USE_MPI = TRUE
+USE_OMP = FALSE
+USE_CUDA = FALSE
+
+TINY_PROFILE = TRUE
+USE_PARTICLES = TRUE
+USE_ASCENT = TRUE
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include ./Make.package
+include $(AMREX_HOME)/Src/Base/Make.package
+include $(AMREX_HOME)/Src/Particle/Make.package
+include $(AMREX_HOME)/Src/Extern/Conduit/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules
diff --git a/Tests/Particles/Ascent_Insitu_SOA/inputs.rt b/Tests/Particles/Ascent_Insitu_SOA/inputs.rt
@@ -0,0 +1,10 @@
+ascent.size = (32, 64, 64)
+ascent.max_grid_size = 32
+ascent.is_periodic = 1
+ascent.num_ppc = 1
+ascent.nlevs = 1
+
+ascent.num_runtime_real = 0
+ascent.num_runtime_int = 0
+
+particles.do_tiling = 1