Skip to content

Commit

Permalink
EBData: New class containing Array4's for all EB data
Browse files Browse the repository at this point in the history 
This can make the code cleaner because we can get access to all available EB
data in one object. This also reduces many GPU kernels' register pressure
because the GPU device lambdas are much smaller without the need to capture
multiple Array4's. For example, the 3D EBABecLap's gsrb kernel's number of
registers reduces from 198 to 130.

Add a new function that returns a random point on EB for WarpX.
  • Loading branch information
@WeiqunZhang
WeiqunZhang committed Nov 21, 2024
1 parent b98eaaf commit c759d7d
Show file tree
Hide file tree
Showing 9 changed files with 232 additions and 99 deletions.
56 changes: 56 additions & 0 deletions Src/EB/AMReX_EBData.H
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
#ifndef AMREX_EB_DATA_H_
#define AMREX_EB_DATA_H_
#include <AMReX_Config.H>

#include <AMReX_EBCellFlag.H>

namespace amrex
{

enum struct EBData_t : int
{
levelset, // level set
volfrac, // volume fraction
centroid, // volume centroid
bndrycent, // boundary centroid
bndrynorm, // boundary normal
bndryarea, // boundary area
AMREX_D_DECL(apx, apy, apz), // area fraction
AMREX_D_DECL(fcx, fcy, fcz), // face centroid
AMREX_D_DECL(ecx, ecy, ecz), // edge centroid
cellflag // EBCellFlag
};

struct EBData
{
template <EBData_t T>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
auto get (int i, int j, int k) const noexcept
{
if constexpr (T == EBData_t::cellflag) {
return (*m_cell_flag)(i,j,k);
} else {
return m_real_data[static_cast<int>(T)](i,j,k);
}
}

template <EBData_t T, std::enable_if_t< T == EBData_t::centroid
|| T == EBData_t::bndrycent
|| T == EBData_t::bndrynorm
AMREX_D_TERM(|| T==EBData_t::fcx,
|| T==EBData_t::fcy,
|| T==EBData_t::fcz), int> = 0>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
auto get (int i, int j, int k, int n) const noexcept
{
return m_real_data[static_cast<int>(T)](i,j,k,n);
}

static constexpr int real_data_size = static_cast<int>(EBData_t::cellflag);

Array4<EBCellFlag const> const* m_cell_flag = nullptr;
Array4<Real const> const* m_real_data = nullptr;
};

}
#endif
3 changes: 3 additions & 0 deletions Src/EB/AMReX_EBDataCollection.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

namespace amrex {

class EBFArrayBoxFactory;
template <class T> class FabArray;
class MultiFab;
class iMultiFab;
Expand All @@ -19,6 +20,8 @@ class EBDataCollection
{
public:

friend class EBFArrayBoxFactory;

EBDataCollection (const EB2::Level& a_level, const Geometry& a_geom,
const BoxArray& a_ba, const DistributionMapping& a_dm,
Vector<int> a_ngrow, EBSupport a_support);
Expand Down
5 changes: 5 additions & 0 deletions Src/EB/AMReX_EBFabFactory.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@

#include <AMReX_FabFactory.H>

#include <AMReX_EBData.H>
#include <AMReX_EBDataCollection.H>
#include <AMReX_Geometry.H>
#include <AMReX_EBSupport.H>
#include <AMReX_Array.H>
#include <AMReX_MFIter.H>

namespace amrex
{
Expand Down Expand Up @@ -89,12 +91,15 @@ public:
//! One should use getMultiEBCellFlagFab for normal levels.
[[nodiscard]] iMultiFab const* getCutCellMask () const noexcept { return m_ebdc->getCutCellMask(); }

[[nodiscard]] EBData getEBData (MFIter const& mfi) const noexcept;

private:

EBSupport m_support;
Geometry m_geom;
std::shared_ptr<EBDataCollection> m_ebdc;
EB2::Level const* m_parent = nullptr;
Gpu::DeviceVector<Array4<Real const>> m_eb_data;
};

std::unique_ptr<EBFArrayBoxFactory>
Expand Down
78 changes: 77 additions & 1 deletion Src/EB/AMReX_EBFabFactory.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,70 @@ EBFArrayBoxFactory::EBFArrayBoxFactory (const EB2::Level& a_level,
m_geom(a_geom),
m_ebdc(std::make_shared<EBDataCollection>(a_level,a_geom,a_ba,a_dm,a_ngrow,a_support)),
m_parent(&a_level)
{}
{
auto const& ebflags = getMultiEBCellFlagFab();
#ifdef AMREX_USE_GPU
m_eb_data.resize(EBData::real_data_size*ebflags.local_size());
Gpu::PinnedVector<Array4<Real const>> eb_data_hv;
#else
auto& eb_data_hv = m_eb_data;
#endif

eb_data_hv.reserve(EBData::real_data_size*ebflags.local_size());

for (MFIter mfi(ebflags,MFItInfo{}.DisableDeviceSync()); mfi.isValid(); ++mfi) {
Array4<Real const> a{};

bool cutfab_is_ok = ebflags[mfi].getType() == FabType::singlevalued;

a = ( m_ebdc->m_levelset )
? m_ebdc->m_levelset->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

a = ( m_ebdc->m_volfrac )
? m_ebdc->m_volfrac->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

a = ( m_ebdc->m_centroid && cutfab_is_ok )
? m_ebdc->m_centroid->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

a = ( m_ebdc->m_bndrycent && cutfab_is_ok )
? m_ebdc->m_bndrycent->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

a = ( m_ebdc->m_bndrynorm && cutfab_is_ok )
? m_ebdc->m_bndrynorm->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

a = ( m_ebdc->m_bndryarea && cutfab_is_ok )
? m_ebdc->m_bndryarea->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);

for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
a = ( m_ebdc->m_areafrac[idim] && cutfab_is_ok )
? m_ebdc->m_areafrac[idim]->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);
}

for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
a = ( m_ebdc->m_facecent[idim] && cutfab_is_ok )
? m_ebdc->m_facecent[idim]->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);
}

for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
a = ( m_ebdc->m_edgecent[idim] && cutfab_is_ok )
? m_ebdc->m_edgecent[idim]->const_array(mfi) : Array4<Real const>{};
eb_data_hv.push_back(a);
}
}

#ifdef AMREX_USE_GPU
Gpu::copyAsync(Gpu::hostToDevice, eb_data_hv.begin(), eb_data_hv.end(), m_eb_data.begin());
Gpu::streamSynchronize();
#endif
}

AMREX_NODISCARD
FArrayBox*
Expand Down Expand Up @@ -115,6 +178,19 @@ EBFArrayBoxFactory::hasEBInfo () const noexcept
return m_parent->hasEBInfo();
}

EBData
EBFArrayBoxFactory::getEBData (MFIter const& mfi) const noexcept
{
int const li = mfi.LocalIndex();
auto const& ebflags_ma = this->getMultiEBCellFlagFab().const_arrays();
#ifdef AMREX_USE_GPU
auto const* pebflag = ebflags_ma.dp + li;
#else
auto const* pebflag = ebflags_ma.hp + li;
#endif
return EBData{pebflag, m_eb_data.data()+EBData::real_data_size*li};
}

std::unique_ptr<EBFArrayBoxFactory>
makeEBFabFactory (const Geometry& a_geom,
const BoxArray& a_ba,
Expand Down
1 change: 1 addition & 0 deletions Src/EB/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ foreach(D IN LISTS AMReX_SPACEDIM)
AMReX_EBMultiFabUtil.cpp
AMReX_EBCellFlag.H
AMReX_EBCellFlag.cpp
AMReX_EBData.H
AMReX_EBDataCollection.H
AMReX_EBDataCollection.cpp
AMReX_MultiCutFab.H
Expand Down
2 changes: 2 additions & 0 deletions Src/EB/Make.package
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ CEXE_sources += AMReX_EBMultiFabUtil.cpp
CEXE_headers += AMReX_EBCellFlag.H
CEXE_sources += AMReX_EBCellFlag.cpp

CEXE_headers += AMReX_EBData.H

CEXE_headers += AMReX_EBDataCollection.H
CEXE_sources += AMReX_EBDataCollection.cpp

Expand Down
59 changes: 29 additions & 30 deletions Src/LinearSolvers/MLMG/AMReX_MLEBABecLap_2D_K.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -378,12 +378,8 @@ void mlebabeclap_gsrb (Box const& box,
Array4<int const> const& m1, Array4<int const> const& m3,
Array4<Real const> const& f0, Array4<Real const> const& f2,
Array4<Real const> const& f1, Array4<Real const> const& f3,
Array4<const int> const& ccm, Array4<EBCellFlag const> const& flag,
Array4<Real const> const& vfrc,
Array4<Real const> const& apx, Array4<Real const> const& apy,
Array4<Real const> const& fcx, Array4<Real const> const& fcy,
Array4<Real const> const& ba, Array4<Real const> const& bc,
Array4<Real const> const& beb,
Array4<const int> const& ccm, Array4<Real const> const& beb,
EBData const& ebdata,
bool is_dirichlet, bool beta_on_centroid, bool phi_on_centroid,
Box const& vbox, int redblack, int ncomp) noexcept
{
Expand All @@ -394,7 +390,8 @@ void mlebabeclap_gsrb (Box const& box,
{
if ((i+j+k+redblack) % 2 == 0)
{
if (flag(i,j,k).isCovered())
auto const flag = ebdata.get<EBData_t::cellflag>(i,j,k);
if (flag.isCovered())
{
phi(i,j,k,n) = Real(0.0);
}
Expand All @@ -409,7 +406,7 @@ void mlebabeclap_gsrb (Box const& box,
Real cf3 = (j == vhi.y && m3(i,vhi.y+1,k) > 0)
? f3(i,vhi.y,k,n) : Real(0.0);

if (flag(i,j,k).isRegular())
if (flag.isRegular())
{
Real gamma = alpha*a(i,j,k)
+ dhx * (bX(i+1,j,k,n) + bX(i,j,k,n))
Expand All @@ -428,19 +425,19 @@ void mlebabeclap_gsrb (Box const& box,
}
else
{
Real kappa = vfrc(i,j,k);
Real apxm = apx(i,j,k);
Real apxp = apx(i+1,j,k);
Real apym = apy(i,j,k);
Real apyp = apy(i,j+1,k);
Real kappa = ebdata.get<EBData_t::volfrac>(i,j,k);
Real apxm = ebdata.get<EBData_t::apx>(i ,j ,k);
Real apxp = ebdata.get<EBData_t::apx>(i+1,j ,k);
Real apym = ebdata.get<EBData_t::apy>(i ,j ,k);
Real apyp = ebdata.get<EBData_t::apy>(i ,j+1,k);

Real fxm = -bX(i,j,k,n)*phi(i-1,j,k,n);
Real oxm = -bX(i,j,k,n)*cf0;
Real sxm = bX(i,j,k,n);
if (apxm != Real(0.0) && apxm != Real(1.0)) {
int jj = j + static_cast<int>(std::copysign(Real(1.0),fcx(i,j,k)));
Real fracy = (ccm(i-1,jj,k) || ccm(i,jj,k))
? std::abs(fcx(i,j,k)) : Real(0.0);
Real fcx = ebdata.get<EBData_t::fcx>(i,j,k);
int jj = j + static_cast<int>(std::copysign(Real(1.0),fcx));
Real fracy = (ccm(i-1,jj,k) || ccm(i,jj,k)) ? std::abs(fcx) : Real(0.0);
if (!beta_on_centroid && !phi_on_centroid)
{
fxm = (Real(1.0)-fracy)*fxm +
Expand All @@ -460,9 +457,9 @@ void mlebabeclap_gsrb (Box const& box,
Real oxp = bX(i+1,j,k,n)*cf2;
Real sxp = -bX(i+1,j,k,n);
if (apxp != Real(0.0) && apxp != Real(1.0)) {
int jj = j + static_cast<int>(std::copysign(Real(1.0),fcx(i+1,j,k)));
Real fracy = (ccm(i,jj,k) || ccm(i+1,jj,k))
? std::abs(fcx(i+1,j,k)) : Real(0.0);
Real fcx = ebdata.get<EBData_t::fcx>(i+1,j,k);
int jj = j + static_cast<int>(std::copysign(Real(1.0),fcx));
Real fracy = (ccm(i,jj,k) || ccm(i+1,jj,k)) ? std::abs(fcx) : Real(0.0);
if (!beta_on_centroid && !phi_on_centroid)
{
fxp = (Real(1.0)-fracy)*fxp +
Expand All @@ -482,9 +479,9 @@ void mlebabeclap_gsrb (Box const& box,
Real oym = -bY(i,j,k,n)*cf1;
Real sym = bY(i,j,k,n);
if (apym != Real(0.0) && apym != Real(1.0)) {
int ii = i + static_cast<int>(std::copysign(Real(1.0),fcy(i,j,k)));
Real fracx = (ccm(ii,j-1,k) || ccm(ii,j,k))
? std::abs(fcy(i,j,k)) : Real(0.0);
Real fcy = ebdata.get<EBData_t::fcy>(i,j,k);
int ii = i + static_cast<int>(std::copysign(Real(1.0),fcy));
Real fracx = (ccm(ii,j-1,k) || ccm(ii,j,k)) ? std::abs(fcy) : Real(0.0);
if (!beta_on_centroid && !phi_on_centroid)
{
fym = (Real(1.0)-fracx)*fym +
Expand All @@ -505,9 +502,9 @@ void mlebabeclap_gsrb (Box const& box,
Real oyp = bY(i,j+1,k,n)*cf3;
Real syp = -bY(i,j+1,k,n);
if (apyp != Real(0.0) && apyp != Real(1.0)) {
int ii = i + static_cast<int>(std::copysign(Real(1.0),fcy(i,j+1,k)));
Real fracx = (ccm(ii,j,k) || ccm(ii,j+1,k))
? std::abs(fcy(i,j+1,k)) : Real(0.0);
Real fcy = ebdata.get<EBData_t::fcy>(i,j+1,k);
int ii = i + static_cast<int>(std::copysign(Real(1.0),fcy));
Real fracx = (ccm(ii,j,k) || ccm(ii,j+1,k)) ? std::abs(fcy) : Real(0.0);
if (!beta_on_centroid && !phi_on_centroid)
{
fyp = (Real(1.0)-fracx)*fyp +
Expand Down Expand Up @@ -543,9 +540,9 @@ void mlebabeclap_gsrb (Box const& box,
Real anrmx = dapx * anorminv;
Real anrmy = dapy * anorminv;

Real bctx = bc(i,j,k,0);
Real bcty = bc(i,j,k,1);
Real dx_eb = get_dx_eb(vfrc(i,j,k));
Real bctx = ebdata.get<EBData_t::bndrycent>(i,j,k,0);
Real bcty = ebdata.get<EBData_t::bndrycent>(i,j,k,1);
Real dx_eb = get_dx_eb(kappa);

Real dg, gx, gy, sx, sy;
if (std::abs(anrmx) > std::abs(anrmy)) {
Expand All @@ -569,10 +566,12 @@ void mlebabeclap_gsrb (Box const& box,
// In gsrb we are always in residual-correction form so phib = 0
Real dphidn = ( -phig)/dg;

Real feb = dphidn * ba(i,j,k) * beb(i,j,k,n);
Real ba = ebdata.get<EBData_t::bndryarea>(i,j,k);

Real feb = dphidn * ba * beb(i,j,k,n);
rho += -vfrcinv*(-dh)*feb;

Real feb_gamma = -phig_gamma/dg * ba(i,j,k) * beb(i,j,k,n);
Real feb_gamma = -phig_gamma/dg * beb(i,j,k,n);
gamma += vfrcinv*(-dh)*feb_gamma;
}

Expand Down
Loading

0 comments on commit c759d7d

Please sign in to comment.