Skip to content

WarpXSolverVec: Remove duplicates #6415

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Nov 24, 2025
Merged

WarpXSolverVec: Remove duplicates #6415

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1cbac22
WarpXSolverVec: Remove duplicates
WeiqunZhang Nov 20, 2025
5fb36a5
Fix compilation
WeiqunZhang Nov 21, 2025
b5a40cf
dof_lhs_arr -> dof_arr
WeiqunZhang Nov 21, 2025
64dc13e
Fix warning
WeiqunZhang Nov 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions Source/FieldSolver/ImplicitSolvers/WarpXSolverDOF.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,7 @@ class WarpX;
struct WarpXSolverDOF
{
amrex::Vector<std::array<std::unique_ptr<amrex::iMultiFab>,3>> m_array;
amrex::Vector<std::array<std::unique_ptr<amrex::iMultiFab>,3>> m_array_lhs;
amrex::Vector<std::unique_ptr<amrex::iMultiFab>> m_scalar;
amrex::Vector<std::unique_ptr<amrex::iMultiFab>> m_scalar_lhs;

warpx::fields::FieldType m_array_type = warpx::fields::FieldType::None;
warpx::fields::FieldType m_scalar_type = warpx::fields::FieldType::None;
Expand All @@ -35,6 +33,9 @@ struct WarpXSolverDOF
amrex::Long m_nDoFs_g = 0; /*!< Global nDOF */

void Define ( WarpX* const, int, const std::string&, const std::string&);

void fill_local_dof (amrex::iMultiFab& dof, amrex::iMultiFab const& mask);
void fill_global_dof ();
};

#endif
228 changes: 121 additions & 107 deletions Source/FieldSolver/ImplicitSolvers/WarpXSolverDOF.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@
#include <ablastr/utils/SignalHandling.H>
#include <ablastr/warn_manager/WarnManager.H>

#include <AMReX_Scan.H>

using warpx::fields::FieldType;
using namespace amrex;

void WarpXSolverDOF::Define ( WarpX* const a_WarpX,
const int a_num_amr_levels,
Expand Down Expand Up @@ -41,10 +44,7 @@ void WarpXSolverDOF::Define ( WarpX* const a_WarpX,

m_array.resize(a_num_amr_levels);
m_scalar.resize(a_num_amr_levels);
m_array_lhs.resize(a_num_amr_levels);
m_scalar_lhs.resize(a_num_amr_levels);

amrex::Long offset = 0;
m_nDoFs_l = 0;

// Define the 3D vector field data container
Expand All @@ -62,26 +62,9 @@ void WarpXSolverDOF::Define ( WarpX* const a_WarpX,
this_array[n]->DistributionMap(),
2*ncomp, // {local, global} for each comp
this_array[n]->nGrowVect() );
m_nDoFs_g += this_array[n]->boxArray().numPts()*ncomp;

m_array[lev][n]->setVal(-1.0);
amrex::Long offset_mf = 0;
for (amrex::MFIter mfi(*m_array[lev][n]); mfi.isValid(); ++mfi) {
auto bx = mfi.tilebox();
auto dof_arr = m_array[lev][n]->array(mfi);
ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
for (int v = 0; v < ncomp; v++) {
dof_arr(i,j,k,2*v) = bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
+ v
+ offset_mf
+ offset;
}
});
offset_mf += bx.numPts()*ncomp;
}
offset += offset_mf;
m_nDoFs_l += offset_mf;

auto* mask = a_WarpX->getFieldDotMaskPointer(m_array_type, lev, ablastr::fields::Direction{n});
fill_local_dof(*m_array[lev][n], *mask);
}
}

Expand All @@ -101,106 +84,137 @@ void WarpXSolverDOF::Define ( WarpX* const a_WarpX,
this_mf->DistributionMap(),
2*ncomp, // {local, global} for each comp
this_mf->nGrowVect() );
m_nDoFs_g += this_mf->boxArray().numPts()*ncomp;

m_scalar[lev]->setVal(-1.0);
amrex::Long offset_mf = 0;
for (amrex::MFIter mfi(*m_scalar[lev]); mfi.isValid(); ++mfi) {
auto bx = mfi.tilebox();
auto dof_arr = m_scalar[lev]->array(mfi);
ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
for (int v = 0; v < ncomp; v++) {
dof_arr(i,j,k,2*v) = bx.index(amrex::IntVect(AMREX_D_DECL(i, j, k))) * ncomp
+ v
+ offset_mf
+ offset;
}
});
offset_mf += bx.numPts()*ncomp;
}
offset += offset_mf;
m_nDoFs_l += offset_mf;

auto* mask = a_WarpX->getFieldDotMaskPointer(m_scalar_type, lev, ablastr::fields::Direction{0});
fill_local_dof(*m_scalar[lev], *mask);
}

}

auto nDoFs_g = m_nDoFs_l;
amrex::ParallelDescriptor::ReduceLongSum(&nDoFs_g,1);
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
m_nDoFs_g == nDoFs_g,
"WarpXSolverDOF::Define(): something has gone wrong in DoF counting");

auto num_procs = amrex::ParallelDescriptor::NProcs();
auto my_proc = amrex::ParallelDescriptor::MyProc();
amrex::Vector<int> dof_proc_arr(num_procs,0);
dof_proc_arr[my_proc] = m_nDoFs_l;
amrex::ParallelDescriptor::ReduceIntSum(dof_proc_arr.data(), num_procs);
fill_global_dof();

int offset_global = 0;
for (int i = 0; i < my_proc; i++) { offset_global += dof_proc_arr[i]; }

if (m_array_type != FieldType::None) {
for (int lev = 0; lev < a_num_amr_levels; ++lev) {
const ablastr::fields::VectorField this_array = a_WarpX->m_fields.get_alldirs(a_vector_type_name, lev);
for (int n = 0; n < 3; n++) {
auto ncomp = this_array[n]->nComp();
for (amrex::MFIter mfi(*m_array[lev][n]); mfi.isValid(); ++mfi) {
auto bx = mfi.tilebox();
auto dof_arr = m_array[lev][n]->array(mfi);
ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
for (int v = 0; v < ncomp; v++) {
dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + offset_global;
}
});
for (int lev = 0; lev < a_num_amr_levels; ++lev) {
for (int n = 0; n < 3; n++) {
if (auto* dof = m_array[lev][n].get()) {
for (int comp = 1; comp < dof->nComp(); comp += 2) { // Only call this on global id
dof->FillBoundaryAndSync(comp, 1, dof->nGrowVect(), a_WarpX->Geom(lev).periodicity());
}
}
}
if (auto* dof = m_scalar[lev].get()) {
for (int comp = 1; comp < dof->nComp(); comp += 2) { // Only call this on global id
dof->FillBoundaryAndSync(comp, 1, dof->nGrowVect(), a_WarpX->Geom(lev).periodicity());
}
}
}
if (m_scalar_type != FieldType::None) {
for (int lev = 0; lev < a_num_amr_levels; ++lev) {
const amrex::MultiFab* this_mf = a_WarpX->m_fields.get(a_scalar_type_name,lev);
auto ncomp = this_mf->nComp();
for (amrex::MFIter mfi(*m_scalar[lev]); mfi.isValid(); ++mfi) {
auto bx = mfi.tilebox();
auto dof_arr = m_scalar[lev]->array(mfi);
ParallelFor( bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
for (int v = 0; v < ncomp; v++) {
dof_arr(i,j,k,2*v+1) = dof_arr(i,j,k,2*v) + offset_global;

amrex::Print() << "Defined DOF object for linear solves (total DOFs = " << m_nDoFs_g << ").\n";
}

void WarpXSolverDOF::fill_local_dof (iMultiFab& dof, iMultiFab const& mask)
{
int ncomp = dof.nComp() / 2; // /2 because both local and global ids are stored in dof

AMREX_ALWAYS_ASSERT(dof.boxArray().numPts()*ncomp < static_cast<Long>(std::numeric_limits<int>::max()));

dof.setVal(std::numeric_limits<int>::lowest());

#ifdef AMREX_USE_MPI
int nprocs = ParallelDescriptor::NProcs();
#endif

for (MFIter mfi(dof); mfi.isValid(); ++mfi) {
Box const& vbx = mfi.validbox();
int npts = vbx.numPts();
BoxIndexer boxindex(vbx);
auto const& m = mask.const_array(mfi);
auto const& d = dof.array(mfi);
auto start_id = m_nDoFs_l;
auto ndofs = Scan::PrefixSum<int>(
npts,
[=] AMREX_GPU_DEVICE (int offset) -> int
{
auto [i,j,k] = boxindex(offset);
return m(i,j,k) ? 1 : 0;
},
[=] AMREX_GPU_DEVICE (int offset, int ps)
{
auto [i,j,k] = boxindex(offset);
if (m(i,j,k)) {
d(i,j,k,0) = ps + start_id;
#ifdef AMREX_USE_MPI
if (nprocs == 1)
#endif
{
d(i,j,k,1) = ps + start_id;
}
});
}
}
},
Scan::Type::exclusive, Scan::retSum);
if (ncomp > 1) {
ParallelFor(vbx, ncomp-1, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n)
{
if (m(i,j,k)) {
d(i,j,k,2*(n+1)) = d(i,j,k,0) + ndofs*(n+1);
#ifdef AMREX_USE_MPI
if (nprocs == 1)
#endif
{
d(i,j,k,2*(n+1)+1) = d(i,j,k,0) + ndofs*(n+1);
}
}
});
}
m_nDoFs_l += Long(ndofs)*ncomp;
}
}

if (m_array_type != FieldType::None) {
for (int lev = 0; lev < a_num_amr_levels; ++lev) {
const auto& geom = a_WarpX->Geom(lev);
for (int n = 0; n < 3; n++) {
m_array_lhs[lev][n] = std::make_unique<amrex::iMultiFab>(m_array[lev][n]->boxArray(),
m_array[lev][n]->DistributionMap(),
m_array[lev][n]->nComp(),
0 );
amrex::iMultiFab::Copy(*m_array_lhs[lev][n], *m_array[lev][n], 0, 0, m_array[lev][n]->nComp(), 0);
m_array[lev][n]->FillBoundary(geom.periodicity());
// do NOT call FillBoundary() on m_array_lhs
void WarpXSolverDOF::fill_global_dof ()
{
#ifndef AMREX_USE_MPI
m_nDoFs_g = m_nDoFs_l;
#else
int nprocs = ParallelDescriptor::NProcs();
if (nprocs == 1) {
m_nDoFs_g = m_nDoFs_l;
} else {
Vector<Long> ndofs_allprocs(nprocs);
MPI_Allgather(&m_nDoFs_l, 1, ParallelDescriptor::Mpi_typemap<Long>::type(),
ndofs_allprocs.data(), 1, ParallelDescriptor::Mpi_typemap<Long>::type(),
ParallelDescriptor::Communicator());
Long proc_begin = 0;
int myproc = ParallelDescriptor::MyProc();
m_nDoFs_g = 0;
for (int iproc = 0; iproc < nprocs; ++iproc) {
if (iproc < myproc) {
proc_begin += ndofs_allprocs[iproc];
}
m_nDoFs_g += ndofs_allprocs[iproc];
}
}
if (m_scalar_type != FieldType::None) {
for (int lev = 0; lev < a_num_amr_levels; ++lev) {
m_scalar_lhs[lev] = std::make_unique<amrex::iMultiFab>(m_scalar[lev]->boxArray(),
m_scalar[lev]->DistributionMap(),
m_scalar[lev]->nComp(),
0 );
amrex::iMultiFab::Copy(*m_scalar_lhs[lev], *m_scalar[lev], 0, 0, m_scalar[lev]->nComp(), 0);
const auto& geom = a_WarpX->Geom(lev);
m_scalar[lev]->FillBoundary(geom.periodicity());
// do NOT call FillBoundary() on m_scalar_lhs
for (auto& x : m_array) {
for (auto& y : x) {
if (y) {
auto const& dof = y->arrays();
auto ncomp = y->nComp() / 2;
ParallelFor(*y, IntVect(0), ncomp, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k, int n)
{
dof[b](i,j,k,2*n+1) = dof[b](i,j,k,2*n) + int(proc_begin);
});
}
}
}
for (auto& x : m_scalar) {
if (x) {
auto const& dof = x->arrays();
auto ncomp = x->nComp() / 2;
ParallelFor(*x, IntVect(0), ncomp, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k, int n)
{
dof[b](i,j,k,2*n+1) = dof[b](i,j,k,2*n) + int(proc_begin);
});
}
}
Gpu::streamSynchronize();
}
#endif

amrex::Print() << "Defined DOF object for linear solves (total DOFs = " << m_nDoFs_g << ").\n";
}
18 changes: 14 additions & 4 deletions Source/FieldSolver/ImplicitSolvers/WarpXSolverVec.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,13 @@ void WarpXSolverVec::copyFrom ( const amrex::Real* const a_arr)
{
for (int v = 0; v < ncomp; v++) {
int dof = dof_arr(i,j,k,2*v); // local
data_arr(i,j,k,v) = a_arr[dof];
if (dof >= 0) {
data_arr(i,j,k,v) = a_arr[dof];
}
}
});
}
m_array_vec[lev][n]->FillBoundaryAndSync(m_WarpX->Geom(lev).periodicity());
}
}
if (m_scalar_type != FieldType::None) {
Expand All @@ -176,10 +179,13 @@ void WarpXSolverVec::copyFrom ( const amrex::Real* const a_arr)
{
for (int v = 0; v < ncomp; v++) {
int dof = dof_arr(i,j,k,2*v); // local
data_arr(i,j,k,v) = a_arr[dof];
if (dof >= 0) {
data_arr(i,j,k,v) = a_arr[dof];
}
}
});
}
m_scalar_vec[lev]->FillBoundaryAndSync(m_WarpX->Geom(lev).periodicity());
}
}
}
Expand All @@ -205,7 +211,9 @@ void WarpXSolverVec::copyTo ( amrex::Real* const a_arr) const
{
for (int v = 0; v < ncomp; v++) {
int dof = dof_arr(i,j,k,2*v); // local
a_arr[dof] = data_arr(i,j,k,v);
if (dof >= 0) {
a_arr[dof] = data_arr(i,j,k,v);
}
}
});
}
Expand All @@ -221,7 +229,9 @@ void WarpXSolverVec::copyTo ( amrex::Real* const a_arr) const
{
for (int v = 0; v < ncomp; v++) {
int dof = dof_arr(i,j,k,2*v); // local
a_arr[dof] = data_arr(i,j,k,v);
if (dof >= 0) {
a_arr[dof] = data_arr(i,j,k,v);
}
}
});
}
Expand Down
Loading
Loading