Ascent: SoA Particle Support (#3350)

## Summary

Add support for pure SoA layouted particle containers for Ascent.

## Additional background

Follow-up to #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]>

Author: ax3l

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,

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

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()

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

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