Skip to content
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

Combine back-to-back saxpy operations, as well as saxpy next to xpay #4366

Open
wants to merge 2 commits into
base: development
Choose a base branch
from
Open

Combine back-to-back saxpy operations, as well as saxpy next to xpay #4366

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
73018e9
combine saxpy with xpay
mbkuhn Mar 10, 2025
25951a3
set up double saxpy's
mbkuhn Mar 10, 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
149 changes: 149 additions & 0 deletions Src/Base/AMReX_FabArray.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1524,6 +1524,43 @@ public:
value_type a, const FabArray<FAB>& x, int xcomp,
value_type b, const FabArray<FAB>& y, int ycomp,
int dstcomp, int numcomp, const IntVect& nghost);

/**
* \brief y = x2+a2*(y+a1*x1)
*
* \param y FabArray y
* \param a_saxpy scalar a_saxpy
* \param x_saxpy FabArray x_saxpy
* \param a_xpay scalar a_xpay
* \param x_xpay FabArray x_xpay
* \param xcomp starting component of x
* \param ycomp starting component of y
* \param ncomp number of components
* \param nghost number of ghost cells
*/
template <class F=FAB, std::enable_if_t<IsBaseFab<F>::value,int> = 0>
static void Saxpy_Xpay (FabArray<FAB>& y, value_type a_saxpy, FabArray<FAB> const& x_saxpy,
value_type a_xpay, FabArray<FAB> const& x_xpay,
int xcomp, int ycomp, int ncomp, IntVect const& nghost);

/**
* \brief y1 += a1*x1; y2 += a2*x2;
*
* \param y1 FabArray y1
* \param a1 scalar a1
* \param x1 FabArray x1
* \param y2 FabArray y2
* \param a2 scalar a2
* \param x2 FabArray x2
* \param xcomp starting component of x1, x2
* \param ycomp starting component of y1, y2
* \param ncomp number of components
* \param nghost number of ghost cells
*/
template <class F=FAB, std::enable_if_t<IsBaseFab<F>::value,int> = 0>
static void Saxpy_Saxpy (FabArray<FAB>& y1, value_type a1, FabArray<FAB> const& x1,
FabArray<FAB>& y2, value_type a2, FabArray<FAB> const& x2,
int xcomp, int ycomp, int ncomp, IntVect const& nghost);
};


Expand Down Expand Up @@ -2918,6 +2955,118 @@ FabArray<FAB>::Xpay (FabArray<FAB>& y, value_type a, FabArray<FAB> const& x,
}
}

template <class FAB>
template <class F, std::enable_if_t<IsBaseFab<F>::value,int> FOO>
void FabArray<FAB>::Saxpy_Xpay (FabArray<FAB>& y, value_type a_saxpy, FabArray<FAB> const& x_saxpy,
value_type a_xpay, FabArray<FAB> const& x_xpay,
int xcomp, int ycomp, int ncomp, IntVect const& nghost)
{
AMREX_ASSERT(y.boxArray() == x_saxpy.boxArray());
AMREX_ASSERT(y.distributionMap == x_saxpy.distributionMap);
AMREX_ASSERT(y.nGrowVect().allGE(nghost) && x_saxpy.nGrowVect().allGE(nghost));

AMREX_ASSERT(y.boxArray() == x_xpay.boxArray());
AMREX_ASSERT(y.distributionMap == x_xpay.distributionMap);
AMREX_ASSERT(y.nGrowVect().allGE(nghost) && x_xpay.nGrowVect().allGE(nghost));

BL_PROFILE("FabArray::Saxpy_Xpay()");

#ifdef AMREX_USE_GPU
if (Gpu::inLaunchRegion() && y.isFusingCandidate()) {
auto const& yma = y.arrays();
auto const& xma_s = x_saxpy.const_arrays();
auto const& xma_x = x_xpay.const_arrays();
ParallelFor(y, nghost, ncomp,
[=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k, int n) noexcept
{
yma[box_no](i,j,k,ycomp+n) += a_saxpy * xma_s[box_no](i,j,k,xcomp+n);
yma[box_no](i,j,k,ycomp+n) = xma_x[box_no](i,j,k,n+xcomp)
+ a_xpay * yma[box_no](i,j,k,n+ycomp);
});
if (!Gpu::inNoSyncRegion()) {
Gpu::streamSynchronize();
}
} else
#endif
{
#ifdef AMREX_USE_OMP
#pragma omp parallel if (Gpu::notInLaunchRegion())
#endif
for (MFIter mfi(y,TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const Box& bx = mfi.growntilebox(nghost);

if (bx.ok()) {
auto const& xfab_s = x_saxpy.const_array(mfi);
auto const& xfab_x = x_xpay.const_array(mfi);
auto const& yfab = y.array(mfi);
AMREX_HOST_DEVICE_PARALLEL_FOR_4D( bx, ncomp, i, j, k, n,
{
yfab(i,j,k,ycomp+n) += a_saxpy * xfab_s(i,j,k,xcomp+n);
yfab(i,j,k,ycomp+n) = xfab_x(i,j,k,xcomp+n)
+ a_xpay * yfab(i,j,k,ycomp+n);
});
}
}
}
}

template <class FAB>
template <class F, std::enable_if_t<IsBaseFab<F>::value,int> FOO>
void FabArray<FAB>::Saxpy_Saxpy (FabArray<FAB>& y1, value_type a1, FabArray<FAB> const& x1,
FabArray<FAB>& y2, value_type a2, FabArray<FAB> const& x2,
int xcomp, int ycomp, int ncomp, IntVect const& nghost)
{
AMREX_ASSERT(y1.boxArray() == x1.boxArray());
AMREX_ASSERT(y1.distributionMap == x1.distributionMap);
AMREX_ASSERT(y1.nGrowVect().allGE(nghost) && x1.nGrowVect().allGE(nghost));

AMREX_ASSERT(y2.boxArray() == x2.boxArray());
AMREX_ASSERT(y2.distributionMap == x2.distributionMap);
AMREX_ASSERT(y2.nGrowVect().allGE(nghost) && x2.nGrowVect().allGE(nghost));

BL_PROFILE("FabArray::Saxpy_Saxpy()");

#ifdef AMREX_USE_GPU
if (Gpu::inLaunchRegion() && y1.isFusingCandidate()) {
auto const& y1ma = y1.arrays();
auto const& x1ma = x1.const_arrays();
auto const& y2ma = y2.arrays();
auto const& x2ma = x2.const_arrays();
ParallelFor(y1, nghost, ncomp,
[=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k, int n) noexcept
{
y1ma[box_no](i,j,k,ycomp+n) += a1 * x1ma[box_no](i,j,k,xcomp+n);
y2ma[box_no](i,j,k,ycomp+n) += a2 * x2ma[box_no](i,j,k,xcomp+n);
});
if (!Gpu::inNoSyncRegion()) {
Gpu::streamSynchronize();
}
} else
#endif
{
#ifdef AMREX_USE_OMP
#pragma omp parallel if (Gpu::notInLaunchRegion())
#endif
for (MFIter mfi(y1,TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const Box& bx = mfi.growntilebox(nghost);

if (bx.ok()) {
auto const& x1fab = x1.const_array(mfi);
auto const& y1fab = y1.array(mfi);
auto const& x2fab = x2.const_array(mfi);
auto const& y2fab = y2.array(mfi);
AMREX_HOST_DEVICE_PARALLEL_FOR_4D( bx, ncomp, i, j, k, n,
{
y1fab(i,j,k,ycomp+n) += a1 * x1fab(i,j,k,xcomp+n);
y2fab(i,j,k,ycomp+n) += a2 * x2fab(i,j,k,xcomp+n);
});
}
}
}
}

template <class FAB>
template <class F, std::enable_if_t<IsBaseFab<F>::value,int> FOO>
void
Expand Down
18 changes: 18 additions & 0 deletions Src/Base/AMReX_FabArrayUtility.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1858,6 +1858,24 @@ void Xpay (MF& dst, typename MF::value_type a, MF const& src, int scomp, int dco
MF::Xpay(dst, a, src, scomp, dcomp, ncomp, nghost);
}

//! dst += a * src; dst = src + a * dst
template <class MF, std::enable_if_t<IsMultiFabLike_v<MF>,int> = 0>
void Saxpy_Xpay (MF& dst, typename MF::value_type a_saxpy, MF const& src_saxpy,
typename MF::value_type a_xpay, MF const& src_xpay, int scomp, int dcomp,
int ncomp, IntVect const& nghost)
{
MF::Saxpy_Xpay(dst, a_saxpy, src_saxpy, a_xpay, src_xpay, scomp, dcomp, ncomp, nghost);
}

//! dst1 += a1 * src1; dst1 += a1 * src1
template <class MF, std::enable_if_t<IsMultiFabLike_v<MF>,int> = 0>
void Saxpy_Saxpy (MF& dst1, typename MF::value_type a1, MF const& src1,
MF& dst2, typename MF::value_type a2, MF const& src2, int scomp, int dcomp,
int ncomp, IntVect const& nghost)
{
MF::Saxpy_Saxpy(dst1, a1, src1, dst2, a2, src2, scomp, dcomp, ncomp, nghost);
}

//! dst = a*src_a + b*src_b
template <class MF, std::enable_if_t<IsMultiFabLike_v<MF>,int> = 0>
void LinComb (MF& dst,
Expand Down
31 changes: 31 additions & 0 deletions Src/Base/AMReX_MultiFab.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -620,6 +620,37 @@ public:

using FabArray<FArrayBox>::Xpay;

/**
* \brief dst = src2+a2*(dst+a1*src1)
*/
static void Saxpy_Xpay (MultiFab& dst,
Real a_saxpy,
const MultiFab& src_saxpy,
Real a_xpay,
const MultiFab& src_xpay,
int srccomp,
int dstcomp,
int numcomp,
int nghost);

using FabArray<FArrayBox>::Saxpy_Xpay;

/**
* \brief dst1 += a1*src1; dst2 += a2*src2
*/
static void Saxpy_Saxpy (MultiFab& dst1,
Real a1,
const MultiFab& src1,
MultiFab& dst2,
Real a2,
const MultiFab& src2,
int srccomp,
int dstcomp,
int numcomp,
int nghost);

using FabArray<FArrayBox>::Saxpy_Saxpy;

/**
* \brief dst = a*x + b*y
*/
Expand Down
16 changes: 16 additions & 0 deletions Src/Base/AMReX_MultiFab.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,22 @@ MultiFab::Xpay (MultiFab& dst, Real a, const MultiFab& src,
Xpay(dst,a,src,srccomp,dstcomp,numcomp,IntVect(nghost));
}

void
MultiFab::Saxpy_Xpay (MultiFab& dst, Real a_saxpy, const MultiFab& src_saxpy,
Real a_xpay, const MultiFab& src_xpay,
int srccomp, int dstcomp, int numcomp, int nghost)
{
Saxpy_Xpay(dst,a_saxpy,src_saxpy,a_xpay,src_xpay,srccomp,dstcomp,numcomp,IntVect(nghost));
}

void
MultiFab::Saxpy_Saxpy (MultiFab& dst1, Real a1, const MultiFab& src1,
MultiFab& dst2, Real a2, const MultiFab& src2,
int srccomp, int dstcomp, int numcomp, int nghost)
{
Saxpy_Saxpy(dst1,a1,src1,dst2,a2,src2,srccomp,dstcomp,numcomp,IntVect(nghost));
}

void
MultiFab::LinComb (MultiFab& dst,
Real a, const MultiFab& x, int xcomp,
Expand Down
18 changes: 10 additions & 8 deletions Src/LinearSolvers/MLMG/AMReX_MLCGSolver.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,8 +166,9 @@ MLCGSolverT<MF>::solve_bicgstab (MF& sol, const MF& rhs, RT eps_rel, RT eps_abs)
else
{
const RT beta = (rho/rho_1)*(alpha/omega);
Saxpy(p, -omega, v, 0, 0, ncomp, nghost); // p += -omega*v
Xpay(p, beta, r, 0, 0, ncomp, nghost); // p = r + beta*p
// two operations: p += -omega*v; p = r + beta*p
// same as: p = r + beta*(p - omega*v)
Saxpy_Xpay(p, -omega, v, beta, r, 0, 0, ncomp, nghost);
}
Lp.apply(amrlev, mglev, v, p, MLLinOpT<MF>::BCMode::Homogeneous, MLLinOpT<MF>::StateMode::Correction);
Lp.normalize(amrlev, mglev, v);
Expand All @@ -181,8 +182,9 @@ MLCGSolverT<MF>::solve_bicgstab (MF& sol, const MF& rhs, RT eps_rel, RT eps_abs)
{
ret = 2; break;
}
Saxpy(sol, alpha, p, 0, 0, ncomp, nghost); // sol += alpha * p
Saxpy(r, -alpha, v, 0, 0, ncomp, nghost); // r += -alpha * v
// sol += alpha * p; r += -alpha * v
Saxpy_Saxpy(sol, alpha, p, r, -alpha, v, 0, 0, ncomp, nghost);


rnorm = norm_inf(r);

Expand Down Expand Up @@ -217,8 +219,8 @@ MLCGSolverT<MF>::solve_bicgstab (MF& sol, const MF& rhs, RT eps_rel, RT eps_abs)
{
ret = 3; break;
}
Saxpy(sol, omega, r, 0, 0, ncomp, nghost); // sol += omega * r
Saxpy(r, -omega, t, 0, 0, ncomp, nghost); // r += -omega * t
// sol += omega * r; r += -omega * t
Saxpy_Saxpy(sol, omega, r, r, -omega, t, 0, 0, ncomp, nghost);

rnorm = norm_inf(r);

Expand Down Expand Up @@ -359,8 +361,8 @@ MLCGSolverT<MF>::solve_cg (MF& sol, const MF& rhs, RT eps_rel, RT eps_abs)
<< " rho " << rho
<< " alpha " << alpha << '\n';
}
Saxpy(sol, alpha, p, 0, 0, ncomp, nghost); // sol += alpha * p
Saxpy(r, -alpha, q, 0, 0, ncomp, nghost); // r += -alpha * q
// sol += alpha * p; r += -alpha * q
Saxpy_Saxpy(sol, alpha, p, r, -alpha, q, 0, 0, ncomp, nghost);
rnorm = norm_inf(r);

if ( verbose > 2 )
Expand Down
Loading