Skip to content

Commit

Permalink
Ascent: SoA Particle Support (#3350)
Browse files Browse the repository at this point in the history
## 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 <atmyers2@gmail.com>
  • Loading branch information
ax3l and atmyers authored Nov 21, 2023
1 parent c51e0a3 commit 9e35dc1
Show file tree
Hide file tree
Showing 6 changed files with 491 additions and 74 deletions.
5 changes: 2 additions & 3 deletions Src/Extern/Conduit/AMReX_Conduit_Blueprint.H
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
207 changes: 136 additions & 71 deletions Src/Extern/Conduit/AMReX_Conduit_Blueprint_ParticlesI.H
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -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"];
Expand All @@ -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++;
}
}

// -------------------------
Expand Down Expand Up @@ -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,
Expand All @@ -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;
Expand All @@ -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
Expand Down
13 changes: 13 additions & 0 deletions Tests/Particles/Ascent_Insitu_SOA/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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()
24 changes: 24 additions & 0 deletions Tests/Particles/Ascent_Insitu_SOA/GNUmakefile
Original file line number Diff line number Diff line change
@@ -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
10 changes: 10 additions & 0 deletions Tests/Particles/Ascent_Insitu_SOA/inputs.rt
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 9e35dc1

Please sign in to comment.