Skip to content

Commit

Permalink
WIP MOCMC refactor to use swarm packs and MeshData and work with mode…
Browse files Browse the repository at this point in the history 
…rn parthenon
  • Loading branch information
@AstroBarker
AstroBarker committed Jan 17, 2025
1 parent 13a84dd commit a2990dd
Show file tree
Hide file tree
Showing 6 changed files with 422 additions and 304 deletions.
109 changes: 87 additions & 22 deletions src/phoebus_driver.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,7 @@ TaskCollection PhoebusDriver::RungeKuttaStage(const int stage) {
}

// Goal: make async regions go away
// TODO(BLB): move moments code to packs / MeshData
TaskRegion &async_region_1 = tc.AddRegion(num_independent_task_lists);
for (int ib = 0; ib < num_independent_task_lists; ib++) {
auto pmb = blocks[ib].get();
Expand Down Expand Up @@ -320,24 +321,68 @@ TaskCollection PhoebusDriver::RungeKuttaStage(const int stage) {
auto fix_flux = tl.AddTask(sndrcv_flux_depend, fixup::FixFluxes, sc0.get());
sndrcv_flux_depend = sndrcv_flux_depend | fix_flux;
}
}

if (rad_mocmc_active) {
using MDT = std::remove_pointer<decltype(sc0.get())>::type;
if (rad_mocmc_active) {
TaskRegion &sync_region_tr = tc.AddRegion(1);
{
for (int i = 0; i < blocks.size(); i++) {
auto &tl = sync_region_tr[0];
auto &pmb = blocks[i];
auto &sc = pmb->meshblock_data.Get()->GetSwarmData();
auto reset_comms =
tl.AddTask(none, &SwarmContainer::ResetCommunication, sc.get());
}
}
}

// TODO(BLB) merge ifs
// if (rad_mocmc_active && stage == integrator->nstages) {
if (rad_mocmc_active) {
TaskRegion &sync_region_mocmc = tc.AddRegion(num_partitions);
for (int n = 0; n < num_partitions; n++) {
auto &tl = sync_region_mocmc[n];

auto &base = pmesh->mesh_data.GetOrAdd("base", n);
using MDT = std::remove_pointer<decltype(base.get())>::type;
// TODO(BRR) stage_name[stage - 1]?
auto &sd0 = pmb->meshblock_data.Get()->GetSwarmData();
auto samples_transport =
tl.AddTask(none, radiation::MOCMCTransport<MDT>, sc0.get(), dt);
auto send = tl.AddTask(samples_transport, &SwarmContainer::Send, sd0.get(),
BoundaryCommSubset::all);
tl.AddTask(none, radiation::MOCMCTransport<MDT>, base.get(), beta * dt);
}
}

if (rad_mocmc_active) {
TaskRegion &sync_region_mocmc_comm = tc.AddRegion(blocks.size());
for (int ib = 0; ib < blocks.size(); ib++) {
auto &tl = sync_region_mocmc_comm[ib];
auto pmb = blocks[ib].get();

// TODO(BRR) stage_name[stage - 1]?
auto &sc = pmb->meshblock_data.Get()->GetSwarmData();
auto send =
tl.AddTask(none, &SwarmContainer::Send, sc.get(), BoundaryCommSubset::all);
auto receive =
tl.AddTask(send, &SwarmContainer::Receive, sd0.get(), BoundaryCommSubset::all);
tl.AddTask(send, &SwarmContainer::Receive, sc.get(), BoundaryCommSubset::all);
}
}

if (rad_mocmc_active) {
TaskRegion &async_region_mocmc_2 = tc.AddRegion(num_partitions);
for (int n = 0; n < num_partitions; n++) {
auto &tl = async_region_mocmc_2[n];

auto &base = pmesh->mesh_data.GetOrAdd("base", n);
auto &sc1 = pmesh->mesh_data.GetOrAdd(stage_name[stage], n);
using MDT = std::remove_pointer<decltype(base.get())>::type;
// TODO(BRR) stage_name[stage - 1]?
auto sample_bounds =
tl.AddTask(receive, radiation::MOCMCSampleBoundaries<MDT>, sc0.get());
auto sample_recon =
tl.AddTask(sample_bounds, radiation::MOCMCReconstruction<MDT>, sc0.get());
tl.AddTask(none, radiation::MOCMCSampleBoundaries<MDT>, base.get(), sc1.get());
auto sample_recon = tl.AddTask(sample_bounds, radiation::MOCMCReconstruction<MDT>,
base.get(), sc1.get());
auto eddington =
tl.AddTask(sample_recon, radiation::MOCMCEddington<MDT>, sc0.get());
tl.AddTask(sample_recon, radiation::MOCMCEddington<MDT>, base.get(), sc1.get());

TaskID geom_src(0);
geom_src = geom_src | eddington;
}
}
Expand Down Expand Up @@ -524,13 +569,16 @@ TaskCollection PhoebusDriver::RungeKuttaStage(const int stage) {
}

// Fix up flux failures
// TODO (DO NOT MERGE) is this loop right
// TODO (BLB) once the fixups can work with MeshData,
// -- then stitch this section back together
TaskRegion &async_region_2 = tc.AddRegion(num_independent_task_lists);
for (int ib = 0; ib < num_independent_task_lists; ib++) {
auto pmb = blocks[ib].get();
auto &tl = async_region_2[ib];

// first make other useful containers
auto &base = pmb->meshblock_data.Get();
// auto &base = pmb->meshblock_data.Get();

// pull out the container we'll use to get fluxes and/or compute RHSs
auto &sc1 = pmb->meshblock_data.Get(stage_name[stage]);
Expand All @@ -548,23 +596,40 @@ TaskCollection PhoebusDriver::RungeKuttaStage(const int stage) {

auto floors =
tl.AddTask(radfixup, fixup::ApplyFloors<MeshBlockData<Real>>, sc1.get());
}

TaskID gas_rad_int(0);
TaskRegion &sync_region_fixup = tc.AddRegion(num_partitions);
TaskID gas_rad_int(0);
for (int ib = 0; ib < num_partitions; ib++) {
auto &tl = sync_region_fixup[ib];

auto &base = pmesh->mesh_data.GetOrAdd("base", ib);
auto &sc1 = pmesh->mesh_data.GetOrAdd(stage_name[stage], ib);
using MDT = std::remove_pointer<decltype(base.get())>::type;
if (rad_mocmc_active) {
auto impl_update =
tl.AddTask(floors, radiation::MOCMCFluidSource<MeshBlockData<Real>>, sc1.get(),
beta * dt, fluid_active);
auto impl_edd = tl.AddTask(
impl_update, radiation::MOCMCEddington<MeshBlockData<Real>>, sc1.get());
auto impl_update = tl.AddTask(none, radiation::MOCMCFluidSource<MDT>, base.get(),
sc1.get(), beta * dt, fluid_active);
auto impl_edd =
tl.AddTask(impl_update, radiation::MOCMCEddington<MDT>, base.get(), sc1.get());
gas_rad_int = gas_rad_int | impl_edd;
} else if (rad_moments_active) {
// auto impl_update =
// tl.AddTask(none, radiation::MomentFluidSource<MeshBlockData<Real>>, sc1.get(),
// beta * dt, fluid_active);
// gas_rad_int = gas_rad_int | impl_update;
}
}

TaskRegion &async_region_fixup2 = tc.AddRegion(num_independent_task_lists);
for (int ib = 0; ib < num_independent_task_lists; ib++) {
auto pmb = blocks[ib].get();
auto &tl = async_region_fixup2[ib];
auto &sc1 = pmb->meshblock_data.Get(stage_name[stage]);
if (rad_moments_active) { // TODO(PRE MERGE) move outside
auto impl_update =
tl.AddTask(floors, radiation::MomentFluidSource<MeshBlockData<Real>>, sc1.get(),
tl.AddTask(none, radiation::MomentFluidSource<MeshBlockData<Real>>, sc1.get(),
beta * dt, fluid_active);
gas_rad_int = gas_rad_int | impl_update;
}

if (rad_moments_active) {
// Only apply floors because MomentFluidSource already ensured that a sensible state
// was returned
auto floors =
Expand Down
24 changes: 24 additions & 0 deletions src/phoebus_utils/variables.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
#ifndef PHOEBUS_UTILS_VARIABLES_HPP_
#define PHOEBUS_UTILS_VARIABLES_HPP_

#include "compile_constants.hpp"

#include <interface/sparse_pack.hpp>
#include <parthenon/package.hpp>

Expand Down Expand Up @@ -41,6 +43,17 @@
static std::string name() { return #varstring; } \
}

#define TENSOR_SWARM(type, ns, varname, ...) \
struct varname : public parthenon::swarm_variable_names::base_t<type, __VA_ARGS__> { \
template <class... Ts> \
KOKKOS_INLINE_FUNCTION varname(Ts &&...args) \
: parthenon::swarm_variable_names::base_t<type, __VA_ARGS__>( \
std::forward<Ts>(args)...) {} \
static std::string name() { return #ns "." #varname; } \
}

using parthenon::variable_names::ANYDIM;

namespace fluid_prim {
VARIABLE(p, density);
VARIABLE(p, velocity);
Expand Down Expand Up @@ -96,6 +109,16 @@ VARIABLE(mocmc.i, Inu1);
VARIABLE(mocmc.i, jinvs);
} // namespace mocmc_internal

namespace mocmc_core {
SWARM_VARIABLE(Real, mocmc.c, t);
SWARM_VARIABLE(Real, mocmc.c, mu_lo);
SWARM_VARIABLE(Real, mocmc.c, mu_hi);
SWARM_VARIABLE(Real, mocmc.c, phi_lo);
SWARM_VARIABLE(Real, mocmc.c, phi_hi);
SWARM_VARIABLE(Real, mocmc.c, ncov);
TENSOR_SWARM(Real, mocmc.c, Inuinv, ANYDIM, MOCMC_NUM_SPECIES);
} // namespace mocmc_core

namespace internal_variables {
VARIABLE_NONS(face_signal_speed);
VARIABLE_NONS(cell_signal_speed);
Expand Down Expand Up @@ -139,6 +162,7 @@ SWARM_VARIABLE(Real, tr, bernoulli);
SWARM_VARIABLE(Real, tr, B_x);
SWARM_VARIABLE(Real, tr, B_y);
SWARM_VARIABLE(Real, tr, B_z);
TENSOR_SWARM(Real, tr, test, ANYDIM, MOCMC_NUM_SPECIES);
} // namespace tracer_variables

namespace diagnostic_variables {
Expand Down
16 changes: 8 additions & 8 deletions src/radiation/geodesics.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,17 @@ void GetXSource(Real &Kcon0, Real &Kcon1, Real &Kcon2, Real &Kcon3, Real src[NDF
src[3] = Kcon3 / Kcon0;
}

KOKKOS_INLINE_FUNCTION
void GetKSource(Real &X0, Real &X1, Real &X2, Real &X3, Real &Kcov0, Real &Kcov1,
Real &Kcov2, Real &Kcov3, Real &Kcon0,
const Geometry::CoordSysMeshBlock &geom, Real source[4]) {
template <class Geom>
KOKKOS_INLINE_FUNCTION void GetKSource(Real &X0, Real &X1, Real &X2, Real &X3,
Real &Kcov0, Real &Kcov1, Real &Kcov2, Real &Kcov3,
Real &Kcon0, Geom &geom, Real source[4]) {
SPACETIMELOOP(mu) { source[mu] = 0.; }
}

KOKKOS_INLINE_FUNCTION
void PushParticle(Real &X0, Real &X1, Real &X2, Real &X3, Real &Kcov0, Real &Kcov1,
Real &Kcov2, Real &Kcov3, const Real &dt,
const Geometry::CoordSysMeshBlock &geom) {
template <class Geom>
KOKKOS_INLINE_FUNCTION void PushParticle(Real &X0, Real &X1, Real &X2, Real &X3,
Real &Kcov0, Real &Kcov1, Real &Kcov2,
Real &Kcov3, const Real &dt, Geom &geom) {
Real c1[NDFULL], c2[NDFULL], d1[NDFULL], d2[NDFULL];
Real Xtmp[NDFULL], Kcontmp[NDFULL], Kcovtmp[NDFULL];
Real Kcov[NDFULL] = {Kcov0, Kcov1, Kcov2, Kcov3};
Expand Down
Loading

0 comments on commit a2990dd

Please sign in to comment.