Skip to content

Commit

Permalink
fix GMRES solver (#2012)
Browse files Browse the repository at this point in the history 
* fix GMRES solver

* update so code compiles without FFT

* add ignore_unused
  • Loading branch information
@asalmgren
asalmgren authored Dec 8, 2024
1 parent b8ec7e5 commit e8c3bd5
Show file tree
Hide file tree
Showing 8 changed files with 132 additions and 117 deletions.
7 changes: 4 additions & 3 deletions Source/LinearSolvers/ERF_ComputeDivergence.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,14 +62,15 @@ void ERF::compute_divergence (int lev, MultiFab& rhs, Vector<MultiFab>& mom_mf,

const Array4<Real const>& ax_arr = ax[lev]->const_array(mfi);
const Array4<Real const>& ay_arr = ay[lev]->const_array(mfi);
const Array4<Real const>& az_arr = az[lev]->const_array(mfi);
const Array4<Real const>& dJ_arr = detJ_cc[lev]->const_array(mfi);

//
// az == 1 for terrain-fitted coordinates
//
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
rhs_arr(i,j,k) = ((ax_arr(i+1,j,k)*rho0u_arr(i+1,j,k) - ax_arr(i,j,k)*rho0u_arr(i,j,k)) * dxInv[0]
+(ay_arr(i,j+1,k)*rho0v_arr(i,j+1,k) - ay_arr(i,j,k)*rho0v_arr(i,j,k)) * dxInv[1]
+(az_arr(i,j,k+1)*rho0w_arr(i,j,k+1) - az_arr(i,j,k)*rho0w_arr(i,j,k)) * dxInv[2]) / dJ_arr(i,j,k);
+( rho0w_arr(i,j,k+1) - rho0w_arr(i,j,k)) * dxInv[2]) / dJ_arr(i,j,k);
});
} // mfi

Expand Down
12 changes: 6 additions & 6 deletions Source/LinearSolvers/ERF_FFTUtils.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ get_fft_bc (Geometry const& lev_geom,
(lev_geom.Domain().smallEnd(dir) == bounding_box.smallEnd(dir)) &&
(lev_geom.Domain().bigEnd(dir) == bounding_box.bigEnd(dir)) ) {
r[dir] = std::make_pair(FFT::Boundary::periodic,FFT::Boundary::periodic);
amrex::Print() << "SETTING " << dir << " TO PERIODIC " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO PERIODIC " << std::endl;
}

else if ( (lev_geom.Domain().smallEnd(dir) == bounding_box.smallEnd(dir)) &&
Expand All @@ -30,7 +30,7 @@ get_fft_bc (Geometry const& lev_geom,
(bc_type_hi != "Outflow" && bc_type_hi != "Open") )
{
r[dir] = std::make_pair(FFT::Boundary::even,FFT::Boundary::even);
amrex::Print() << "SETTING " << dir << " TO EVEN EVEN " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO EVEN EVEN " << std::endl;
}

else if ( (lev_geom.Domain().smallEnd(dir) == bounding_box.smallEnd(dir)) &&
Expand All @@ -39,26 +39,26 @@ get_fft_bc (Geometry const& lev_geom,
(bc_type_hi == "Outflow" || bc_type_hi == "Open") )
{
r[dir] = std::make_pair(FFT::Boundary::odd,FFT::Boundary::odd);
amrex::Print() << "SETTING " << dir << " TO ODD ODD " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO ODD ODD " << std::endl;
}

else if ( (lev_geom.Domain().smallEnd(dir) == bounding_box.smallEnd(dir)) &&
(bc_type_lo == "Outflow" || bc_type_lo == "Open") )
{
r[dir] = std::make_pair(FFT::Boundary::odd,FFT::Boundary::even);
amrex::Print() << "SETTING " << dir << " TO ODD EVEN " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO ODD EVEN " << std::endl;
}

else if ( (lev_geom.Domain().bigEnd(dir) == bounding_box.bigEnd(dir)) &&
(bc_type_hi == "Outflow" || bc_type_hi == "Open") )
{
r[dir] = std::make_pair(FFT::Boundary::even,FFT::Boundary::odd);
amrex::Print() << "SETTING " << dir << " TO EVEN ODD " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO EVEN ODD " << std::endl;
}
else
{
r[dir] = std::make_pair(FFT::Boundary::even,FFT::Boundary::even);
amrex::Print() << "SETTING " << dir << " TO EVEN EVEN " << std::endl;
// amrex::Print() << "SETTING " << dir << " TO EVEN EVEN " << std::endl;
}
} // dir

Expand Down
18 changes: 16 additions & 2 deletions Source/LinearSolvers/ERF_PoissonSolve.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ void ERF::project_velocities (int lev, Real l_dt, Vector<MultiFab>& mom_mf, Mult
Real rhsnorm = rhs[0].norm0();

if (mg_verbose > 0) {
Print() << "Max norm of divergence before solve at level " << lev << " : " << rhsnorm << std::endl;
Print() << "Max/L2 norm of divergence before solve at level " << lev << " : " << rhsnorm << " " << rhs[0].norm2() << std::endl;
}

// ****************************************************************************
Expand Down Expand Up @@ -214,7 +214,21 @@ void ERF::project_velocities (int lev, Real l_dt, Vector<MultiFab>& mom_mf, Mult
{
compute_divergence(lev, rhs[0], mom_mf, geom_tmp[0]);

amrex::Print() << "Max norm of divergence after solve at level " << lev << " : " << rhs[0].norm0() << std::endl;
Print() << "Max/L2 norm of divergence after solve at level " << lev << " : " << rhs[0].norm0() << " " << rhs[0].norm2() << std::endl;

#if 0
// FOR DEBUGGING ONLY
for ( MFIter mfi(rhs[0],TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const Array4<Real const>& rhs_arr = rhs[0].const_array(mfi);
Box bx = mfi.validbox();
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
if (std::abs(rhs_arr(i,j,k)) > 1.e-10) {
amrex::Print() << "RHS AT " << IntVect(i,j,k) << " " << rhs_arr(i,j,k) << std::endl;
}
});
} // mfi
#endif

} // mg_verbose

Expand Down
6 changes: 5 additions & 1 deletion Source/LinearSolvers/ERF_SolveWithGMRES.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,12 @@
using namespace amrex;

/**
* Solve the Poisson equation using GMRES
* Solve the Poisson equation using FFT-preconditioned GMRES
*/
void ERF::solve_with_gmres (int lev, Vector<MultiFab>& rhs, Vector<MultiFab>& phi,
Vector<Array<MultiFab,AMREX_SPACEDIM>>& fluxes)
{
#ifdef ERF_USE_FFT
BL_PROFILE("ERF::solve_with_gmres()");

Real reltol = solverChoice.poisson_reltol;
Expand All @@ -31,4 +32,7 @@ void ERF::solve_with_gmres (int lev, Vector<MultiFab>& rhs, Vector<MultiFab>& ph
gmsolver.solve(phi[0], rhs[0], reltol, abstol);

tp.getFluxes(phi[0], fluxes[0]);
#else
amrex::ignore_unused(lev, rhs, phi, fluxes);
#endif
}
15 changes: 8 additions & 7 deletions Source/LinearSolvers/ERF_TerrainPoisson.H
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,13 @@
#ifndef ERF_TERRAIN_POISSON_H_
#define ERF_TERRAIN_POISSON_H_

#include "ERF_TerrainPoisson_3D_K.H"
#include <AMReX_Geometry.H>
#include <AMReX_MultiFab.H>

#ifdef ERF_USE_FFT
#include <AMReX_FFT_Poisson.H>
#endif

#include <AMReX_Geometry.H>
#include <AMReX_MultiFab.H>
#include "ERF_TerrainPoisson_3D_K.H"
#include <AMReX_FFT_Poisson.H>

class TerrainPoisson
{
Expand All @@ -24,6 +23,8 @@ public:

void apply(amrex::MultiFab& lhs, amrex::MultiFab const& rhs);

void apply_bcs(amrex::MultiFab& phi);

void usePrecond(bool precond_flag);

void getFluxes(amrex::MultiFab& phi, amrex::Array<amrex::MultiFab,AMREX_SPACEDIM>& fluxes);
Expand Down Expand Up @@ -57,9 +58,9 @@ private:
amrex::Gpu::DeviceVector<amrex::Real> m_stretched_dz_d;
const amrex::MultiFab* m_zphys;

#ifdef ERF_USE_FFT
std::unique_ptr<amrex::FFT::PoissonHybrid<amrex::MultiFab>> m_2D_fft_precond;
#endif
amrex::Array<std::pair<amrex::FFT::Boundary,amrex::FFT::Boundary>,AMREX_SPACEDIM> bc_fft;
};

#endif
#endif
107 changes: 54 additions & 53 deletions Source/LinearSolvers/ERF_TerrainPoisson.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#ifdef ERF_USE_FFT

#include "ERF_TerrainPoisson.H"
#include "ERF_FFTUtils.H"

Expand All @@ -14,15 +16,11 @@ TerrainPoisson::TerrainPoisson (Geometry const& geom, BoxArray const& ba,
m_stretched_dz_d(stretched_dz_lev_d),
m_zphys(z_phys_nd)
{
#ifdef ERF_USE_FFT
if (!m_2D_fft_precond) {
Box bounding_box = ba.minimalBox();
auto bc_fft = get_fft_bc(geom,domain_bc_type,bounding_box);
bc_fft = get_fft_bc(geom,domain_bc_type,bounding_box);
m_2D_fft_precond = std::make_unique<FFT::PoissonHybrid<MultiFab>>(geom,bc_fft);
}
#else
amrex::ignore_unused(domain_bc_type);
#endif
}

void TerrainPoisson::usePrecond (bool use_precond_in)
Expand All @@ -34,72 +32,94 @@ void TerrainPoisson::apply (MultiFab& lhs, MultiFab const& rhs)
{
AMREX_ASSERT(rhs.nGrowVect().allGT(0));

auto domlo = lbound(m_geom.Domain());
auto domhi = ubound(m_geom.Domain());

MultiFab& xx = const_cast<MultiFab&>(rhs);

auto const& dxinv = m_geom.InvCellSizeArray();

auto const& y = lhs.arrays();
auto const& zpa = m_zphys->const_arrays();

// Impose periodic and internal boundary conditions
xx.FillBoundary(m_geom.periodicity());
apply_bcs(xx);

auto const& xc = xx.const_arrays();
ParallelFor(rhs, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
terrpoisson_adotx(i, j, k, y[b], xc[b], zpa[b], dxinv[0], dxinv[1]);
});
}

void TerrainPoisson::apply_bcs (MultiFab& phi)
{
auto domlo = lbound(m_geom.Domain());
auto domhi = ubound(m_geom.Domain());

phi.FillBoundary(m_geom.periodicity());

if (!m_geom.isPeriodic(0)) {
for (MFIter mfi(xx,true); mfi.isValid(); ++mfi)
for (MFIter mfi(phi,true); mfi.isValid(); ++mfi)
{
Box bx = mfi.tilebox();
const Array4<Real>& x_arr = xx.array(mfi);
const Array4<Real>& phi_arr = phi.array(mfi);
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
if (i == domlo.x) {
x_arr(i-1,j,k) = x_arr(i,j,k);
} else if (i == domhi.x) { // OUTFLOW
x_arr(i+1,j,k) = -x_arr(i,j,k);
if (bc_fft[0].first == FFT::Boundary::even) {
phi_arr(i-1,j,k) = phi_arr(i,j,k);
} else if (bc_fft[0].first == FFT::Boundary::odd) {
phi_arr(i-1,j,k) = -phi_arr(i,j,k);
}
} else if (i == domhi.x) {
if (bc_fft[0].second == FFT::Boundary::even) {
phi_arr(i+1,j,k) = phi_arr(i,j,k);
} else if (bc_fft[0].second == FFT::Boundary::odd) {
phi_arr(i+1,j,k) = -phi_arr(i,j,k);
}
}
});
}
} // mfi
}
if (!m_geom.isPeriodic(1)) {
for (MFIter mfi(xx,true); mfi.isValid(); ++mfi)
for (MFIter mfi(phi,true); mfi.isValid(); ++mfi)
{
Box bx = mfi.tilebox();
Box bx2(bx); bx2.grow(0,1);
const Array4<Real>& x_arr = xx.array(mfi);
const Array4<Real>& phi_arr = phi.array(mfi);
ParallelFor(bx2, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
if (j == domlo.y) {
x_arr(i,j-1,k) = x_arr(i,j,k);
if (bc_fft[1].first == FFT::Boundary::even) {
phi_arr(i,j-1,k) = phi_arr(i,j,k);
} else if (bc_fft[1].first == FFT::Boundary::odd) {
phi_arr(i,j-1,k) = -phi_arr(i,j,k);
}
} else if (j == domhi.y) {
x_arr(i,j+1,k) = x_arr(i,j,k);
if (bc_fft[1].second == FFT::Boundary::even) {
phi_arr(i,j+1,k) = phi_arr(i,j,k);
} else if (bc_fft[1].second == FFT::Boundary::odd) {
phi_arr(i,j+1,k) = -phi_arr(i,j,k);
}
}
});
} // mfi
}

for (MFIter mfi(xx,true); mfi.isValid(); ++mfi)
for (MFIter mfi(phi,true); mfi.isValid(); ++mfi)
{
Box bx = mfi.tilebox();
Box bx2(bx); bx2.grow(0,1);
Box bx3(bx2); bx3.grow(1,1);
const Array4<Real>& x_arr = xx.array(mfi);
const Array4<Real>& phi_arr = phi.array(mfi);
ParallelFor(bx3, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
if (k == domlo.z) {
x_arr(i,j,k-1) = x_arr(i,j,k);
phi_arr(i,j,k-1) = phi_arr(i,j,k);
} else if (k == domhi.z) {
x_arr(i,j,k+1) = x_arr(i,j,k);
phi_arr(i,j,k+1) = phi_arr(i,j,k);
}
});
} // mfi

auto const& xc = xx.const_arrays();
ParallelFor(rhs, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
terrpoisson_adotx(i, j, k, y[b], xc[b], zpa[b], dxinv[0], dxinv[1]);
});
phi.FillBoundary(m_geom.periodicity());
}

void TerrainPoisson::getFluxes (MultiFab& phi,
Expand All @@ -113,42 +133,24 @@ void TerrainPoisson::getFluxes (MultiFab& phi,
auto const& x = phi.const_arrays();
auto const& zpa = m_zphys->const_arrays();

phi.FillBoundary(m_geom.periodicity());
apply_bcs(phi);

auto const& fx = fluxes[0].arrays();
ParallelFor(fluxes[0], [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
if (i == domlo.x) {
fx[b](i,j,k) = 0.0;
} else if (i == domhi.x+1) {
fx[b](i,j,k) = 0.0;
} else {
fx[b](i,j,k) = terrpoisson_flux_x(i,j,k,x[b],zpa[b],dxinv[0]);
}
fx[b](i,j,k) = terrpoisson_flux_x(i,j,k,x[b],zpa[b],dxinv[0]);
});

auto const& fy = fluxes[1].arrays();
ParallelFor(fluxes[1], [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
if (j == domlo.y) {
fy[b](i,j,k) = 0.0;
} else if (j == domhi.y+1) {
fy[b](i,j,k) = 0.0;
} else {
fy[b](i,j,k) = terrpoisson_flux_y(i,j,k,x[b],zpa[b],dxinv[1]);
}
fy[b](i,j,k) = terrpoisson_flux_y(i,j,k,x[b],zpa[b],dxinv[1]);
});

auto const& fz = fluxes[2].arrays();
ParallelFor(fluxes[2], [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
if (k == domlo.z) {
fz[b](i,j,k) = 0.0;
} else if (k == domhi.z+1) {
fz[b](i,j,k) = 0.0;
} else {
fz[b](i,j,k) = terrpoisson_flux_z(i,j,k,x[b],zpa[b],dxinv[0],dxinv[1]);
}
fz[b](i,j,k) = terrpoisson_flux_z(i,j,k,x[b],zpa[b],dxinv[0],dxinv[1]);
});
}

Expand Down Expand Up @@ -199,8 +201,6 @@ void TerrainPoisson::precond (MultiFab& lhs, MultiFab const& rhs)
#ifdef ERF_USE_FFT
if (m_use_precond)
{
amrex::Print() << "Using the hybrid FFT solver as a preconditioner..." << std::endl;

// Make a version that isn't constant
MultiFab& rhs_tmp = const_cast<MultiFab&>(rhs);

Expand Down Expand Up @@ -249,3 +249,4 @@ void TerrainPoisson::setToZero(MultiFab& v)
{
v.setVal(0);
}
#endif
Loading

0 comments on commit e8c3bd5

Please sign in to comment.