From 9e35dc19489dc5d312e92781cb0471d282cf8370 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 20 Nov 2023 21:04:03 -0800 Subject: [PATCH] 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 --- Src/Extern/Conduit/AMReX_Conduit_Blueprint.H | 5 +- .../AMReX_Conduit_Blueprint_ParticlesI.H | 207 ++++++++---- .../Ascent_Insitu_SOA/CMakeLists.txt | 13 + Tests/Particles/Ascent_Insitu_SOA/GNUmakefile | 24 ++ Tests/Particles/Ascent_Insitu_SOA/inputs.rt | 10 + Tests/Particles/Ascent_Insitu_SOA/main.cpp | 306 ++++++++++++++++++ 6 files changed, 491 insertions(+), 74 deletions(-) create mode 100644 Tests/Particles/Ascent_Insitu_SOA/CMakeLists.txt create mode 100644 Tests/Particles/Ascent_Insitu_SOA/GNUmakefile create mode 100644 Tests/Particles/Ascent_Insitu_SOA/inputs.rt create mode 100644 Tests/Particles/Ascent_Insitu_SOA/main.cpp diff --git a/Src/Extern/Conduit/AMReX_Conduit_Blueprint.H b/Src/Extern/Conduit/AMReX_Conduit_Blueprint.H index 6d23bcf07ed..9ac8eb31fd0 100644 --- 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 - void ParticleContainerToBlueprint (const ParticleContainer + void ParticleContainerToBlueprint (const ParticleContainer_impl &pc, const Vector &real_comp_names, diff --git a/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H b/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H index f2d7d1ed2da..e4186ba247b 100644 --- a/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H +++ b/Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H @@ -20,10 +20,9 @@ namespace amrex // Note: // This is a helper function, it's not part of the AMReX Blueprint Interface. //---------------------------------------------------------------------------// -template +template void -ParticleTileToBlueprint(const ParticleTile, +ParticleTileToBlueprint(const ParticleTile &ptile, const Vector &real_comp_names, @@ -31,15 +30,11 @@ ParticleTileToBlueprint(const ParticleTile); + 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(pstruct.data()); - char* pbuf = const_cast(pbuf_const); + if constexpr(ParticleType::is_soa_particle) + { + amrex::ignore_unused(pbuf); - ParticleReal* xp = reinterpret_cast(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(&soa.GetRealData(0)[0]), + num_particles); #if AMREX_SPACEDIM > 1 - ParticleReal* yp = reinterpret_cast(pbuf); pbuf += sizeof(ParticleReal); - n_coords["values/y"].set_external(yp, - num_particles, - 0, - struct_size); + n_coords["values/y"].set_external(const_cast(&soa.GetRealData(1)[0]), + num_particles); #endif #if AMREX_SPACEDIM > 2 - ParticleReal* zp = reinterpret_cast(pbuf); pbuf += sizeof(ParticleReal); - n_coords["values/z"].set_external(zp, - num_particles, - 0, - struct_size); + n_coords["values/z"].set_external(const_cast(&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(pstruct.data()); + pbuf = const_cast(pbuf_const); + + ParticleReal* xp = reinterpret_cast(pbuf); pbuf += sizeof(ParticleReal); + n_coords["values/x"].set_external(xp, + num_particles, + 0, + struct_size); +#if AMREX_SPACEDIM > 1 + ParticleReal* yp = reinterpret_cast(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(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,20 +117,26 @@ ParticleTileToBlueprint(const ParticleTile(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(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++; + } } //----------------------------------// @@ -116,44 +144,77 @@ ParticleTileToBlueprint(const ParticleTile(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(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(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(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(&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(&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(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(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 +template void -ParticleContainerToBlueprint(const ParticleContainer &pc, const Vector &real_comp_names, @@ -209,8 +269,13 @@ ParticleContainerToBlueprint(const ParticleContainer; + using MyParConstIter = ParConstIter_impl; // // 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 new file mode 100644 index 00000000000..82216a02af8 --- /dev/null +++ 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 new file mode 100644 index 00000000000..660e4a13f22 --- /dev/null +++ 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 new file mode 100644 index 00000000000..e34fda14923 --- /dev/null +++ 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 diff --git a/Tests/Particles/Ascent_Insitu_SOA/main.cpp b/Tests/Particles/Ascent_Insitu_SOA/main.cpp new file mode 100644 index 00000000000..46e2af98422 --- /dev/null +++ b/Tests/Particles/Ascent_Insitu_SOA/main.cpp @@ -0,0 +1,306 @@ +#include +#include +#include + +#if !defined(AMREX_PARTICLES) || !defined(AMREX_USE_CONDUIT) +#error Incompatible AMReX library configuration! This tutorial requires AMREX_PARTICLES and AMREX_USE_CONDUIT +#endif + +#include + +#include + +#include + + +using namespace amrex; + +static constexpr int NR = 7; +static constexpr int NI = 4; + +int num_runtime_real = 0; +int num_runtime_int = 0; + +bool remove_negative = true; + +void get_position_unit_cell (Real* r, const IntVect& nppc, int i_part) +{ + int nx = nppc[0]; +#if AMREX_SPACEDIM > 1 + int ny = nppc[1]; +#else + int ny = 1; +#endif +#if AMREX_SPACEDIM > 2 + int nz = nppc[2]; +#else + int nz = 1; +#endif + + int ix_part = i_part/(ny * nz); + int iy_part = (i_part % (ny * nz)) % ny; + int iz_part = (i_part % (ny * nz)) / ny; + + r[0] = (0.5+ix_part)/nx; + r[1] = (0.5+iy_part)/ny; + r[2] = (0.5+iz_part)/nz; +} + +class TestParticleContainer + : public amrex::ParticleContainerPureSoA +{ + +public: + + TestParticleContainer (const Vector & a_geom, + const Vector & a_dmap, + const Vector & a_ba, + const Vector & a_rr) + : amrex::ParticleContainerPureSoA(a_geom, a_dmap, a_ba, a_rr) + { + for (int i = 0; i < num_runtime_real; ++i) + { + AddRealComp(true); + } + for (int i = 0; i < num_runtime_int; ++i) + { + AddIntComp(true); + } + } + + void InitParticles (const amrex::IntVect& a_num_particles_per_cell) + { + BL_PROFILE("InitParticles"); + + const int lev = 0; // only add particles on level 0 + const Real* dx = Geom(lev).CellSize(); + const Real* plo = Geom(lev).ProbLo(); + + const int num_ppc = AMREX_D_TERM( a_num_particles_per_cell[0], + *a_num_particles_per_cell[1], + *a_num_particles_per_cell[2]); + + for(MFIter mfi = MakeMFIter(lev); mfi.isValid(); ++mfi) + { + const Box& tile_box = mfi.tilebox(); + + std::array, NR> host_real; + std::array, NI> host_int; + + std::vector > host_runtime_real(NumRuntimeRealComps()); + std::vector > host_runtime_int(NumRuntimeIntComps()); + + for (IntVect iv = tile_box.smallEnd(); iv <= tile_box.bigEnd(); tile_box.next(iv)) + { + for (int i_part=0; i_part(id)); + host_int[1].push_back(ParallelDescriptor::MyProc()); + host_real[0].push_back(static_cast (plo[0] + (iv[0] + r[0])*dx[0])); +#if AMREX_SPACEDIM > 1 + host_real[1].push_back(static_cast (plo[1] + (iv[1] + r[1])*dx[1])); +#endif +#if AMREX_SPACEDIM > 2 + host_real[2].push_back(static_cast (plo[2] + (iv[2] + r[2])*dx[2])); +#endif + + for (int i = AMREX_SPACEDIM; i < NR; ++i) + host_real[i].push_back(static_cast(id)); + for (int i = 2; i < NI; ++i) + host_int[i].push_back(static_cast(id)); + for (int i = 0; i < NumRuntimeRealComps(); ++i) + host_runtime_real[i].push_back(static_cast(id)); + for (int i = 0; i < NumRuntimeIntComps(); ++i) + host_runtime_int[i].push_back(static_cast(id)); + } + } + + auto& particle_tile = DefineAndReturnParticleTile(lev, mfi.index(), mfi.LocalTileIndex()); + auto old_size = particle_tile.size(); + auto new_size = old_size + host_real[0].size(); + particle_tile.resize(new_size); + + auto& soa = particle_tile.GetStructOfArrays(); + for (int i = 0; i < NR; ++i) + { + Gpu::copyAsync(Gpu::hostToDevice, + host_real[i].begin(), + host_real[i].end(), + soa.GetRealData(i).begin() + old_size); + } + + for (int i = 0; i < NI; ++i) + { + Gpu::copyAsync(Gpu::hostToDevice, + host_int[i].begin(), + host_int[i].end(), + soa.GetIntData(i).begin() + old_size); + } + for (int i = 0; i < NumRuntimeRealComps(); ++i) + { + Gpu::copyAsync(Gpu::hostToDevice, + host_runtime_real[i].begin(), + host_runtime_real[i].end(), + soa.GetRealData(NR+i).begin() + old_size); + } + + for (int i = 0; i < NumRuntimeIntComps(); ++i) + { + Gpu::copyAsync(Gpu::hostToDevice, + host_runtime_int[i].begin(), + host_runtime_int[i].end(), + soa.GetIntData(NI+i).begin() + old_size); + } + + Gpu::streamSynchronize(); + } + + Redistribute(); + } +}; + +struct TestParams +{ + IntVect size; + int max_grid_size; + int num_ppc; + int is_periodic; + int nlevs; +}; + +void testAscent (); + +int main (int argc, char* argv[]) +{ + amrex::Initialize(argc,argv); + + amrex::Print() << "Running redistribute test \n"; + testAscent(); + + amrex::Finalize(); +} + +void get_test_params (TestParams& params, const std::string& prefix) +{ + ParmParse pp(prefix); + pp.get("size", params.size); + pp.get("max_grid_size", params.max_grid_size); + pp.get("num_ppc", params.num_ppc); + pp.get("is_periodic", params.is_periodic); + pp.get("nlevs", params.nlevs); + pp.query("num_runtime_real", num_runtime_real); + pp.query("num_runtime_int", num_runtime_int); +} + +void testAscent () +{ + BL_PROFILE("testAscent"); + TestParams params; + get_test_params(params, "ascent"); + + int is_per[BL_SPACEDIM]; + for (int & d : is_per) + d = params.is_periodic; + + Vector rr(params.nlevs-1); + for (int lev = 1; lev < params.nlevs; lev++) + rr[lev-1] = IntVect(AMREX_D_DECL(2,2,2)); + + RealBox real_box; + for (int n = 0; n < BL_SPACEDIM; n++) + { + real_box.setLo(n, 0.0); + real_box.setHi(n, params.size[n]); + } + + IntVect domain_lo(AMREX_D_DECL(0, 0, 0)); + IntVect domain_hi(AMREX_D_DECL(params.size[0]-1,params.size[1]-1,params.size[2]-1)); + const Box base_domain(domain_lo, domain_hi); + + Vector geom(params.nlevs); + geom[0].define(base_domain, &real_box, CoordSys::cartesian, is_per); + for (int lev = 1; lev < params.nlevs; lev++) { + geom[lev].define(amrex::refine(geom[lev-1].Domain(), rr[lev-1]), + &real_box, CoordSys::cartesian, is_per); + } + + Vector ba(params.nlevs); + Vector dm(params.nlevs); + auto lo = IntVect(AMREX_D_DECL(0, 0, 0)); + IntVect size = params.size; + for (int lev = 0; lev < params.nlevs; ++lev) + { + ba[lev].define(Box(lo, lo+params.size-1)); + ba[lev].maxSize(params.max_grid_size); + dm[lev].define(ba[lev]); + lo += size/2; + size *= 2; + } + + TestParticleContainer pc(geom, dm, ba, rr); + + int npc = params.num_ppc; + auto nppc = IntVect(AMREX_D_DECL(npc, npc, npc)); + + amrex::Print() << "About to initialize particles \n"; + + pc.InitParticles(nppc); + + { + conduit::Node bp_mesh; + /* TODO + amrex::MultiLevelToBlueprint( + nlev, + amrex::GetVecOfConstPtrs(mf), + varnames, + geom, + time, + iteration, + warpx.refRatio(), + bp_mesh + ); + */ + + // wrap pc for current species into a blueprint topology + std::string const prefix = "particle"; + Vector particle_varnames; + for (int i = 0; i < pc.NumRealComps(); ++i) { + particle_varnames.push_back(prefix + "_real_" + std::to_string(i)); + } + Vector particle_int_varnames; + for (int i = 0; i < pc.NumIntComps(); ++i) { + particle_int_varnames.push_back(prefix + "_int_" + std::to_string(i)); + } + + amrex::ParticleContainerToBlueprint(pc, + particle_varnames, + particle_int_varnames, + bp_mesh, + prefix); + // publish + ascent::Ascent ascent; + conduit::Node opts; + opts["exceptions"] = "catch"; + opts["mpi_comm"] = MPI_Comm_c2f(ParallelDescriptor::Communicator()); + ascent.open(opts); + ascent.publish(bp_mesh); + + // If you want to save blueprint HDF5 files w/o using an Ascent + // extract, you can call the following AMReX helper: + // const auto step = istep[0]; + // amrex::WriteBlueprintFiles(bp_mesh, "bp_export", step, "hdf5"); + + // render + conduit::Node actions; + ascent.execute(actions); + ascent.close(); + } + + // the way this test is set up, if we make it here we pass + amrex::Print() << "pass \n"; +} +