Skip to content

Commit

Permalink
simplify loops in vertical viscosity, and call loops in vert_visc/ver…
Browse files Browse the repository at this point in the history 
…t_mean/advance3d only on valid cells
  • Loading branch information
@hklion
hklion committed Dec 5, 2023
1 parent 69f237a commit e289e94
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 34 deletions.
36 changes: 21 additions & 15 deletions Source/TimeIntegration/ROMSX_advance_3d.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,9 @@ ROMSX::advance_3d (int lev,
Box gbx11 = mfi.growntilebox(IntVect(NGROW-1,NGROW-1,NGROW-1));
Box gbx21 = mfi.growntilebox(IntVect(NGROW,NGROW,NGROW-1));

Box xbx = mfi.nodaltilebox(0);
Box ybx = mfi.nodaltilebox(1);

Box gbx2D = gbx2;
gbx2D.makeSlab(2,0);

Expand Down Expand Up @@ -144,16 +147,18 @@ ROMSX::advance_3d (int lev,
cff=0.25*dt_lev*23.0/12.0;
}

amrex::ParallelFor(gbx2,
amrex::ParallelFor(xbx,
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
u(i,j,k) += cff * (pm(i,j,0)+pm(i-1,j,0)) * (pn(i,j,0)+pn(i-1,j,0)) * ru(i,j,k,nrhs);
u(i,j,k) *= 2.0 / (Hz(i-1,j,k) + Hz(i,j,k));
});
amrex::ParallelFor(ybx,
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
u(i,j,k) += tbxp2.contains(i-1,j,0) ? cff * (pm(i,j,0)+pm(i-1,j,0)) * (pn(i,j,0)+pn(i-1,j,0)) * ru(i,j,k,nrhs) :
cff * (2.0 * pm(i,j,0)) * (2.0 * pn(i,j,0)) * ru(i,j,k,nrhs) ;
v(i,j,k) += tbxp2.contains(i,j-1,0) ? cff * (pm(i,j,0)+pm(i,j-1,0)) * (pn(i,j,0)+pn(i,j-1,0)) * rv(i,j,k,nrhs) :
cff * (2.0 * pm(i,j,0)) * (2.0 * pn(i,j,0)) * rv(i,j,k,nrhs);
u(i,j,k) *= tbxp2.contains(i-1,j,0) ? 2.0 / (Hz(i-1,j,k) + Hz(i,j,k)) : 1.0 / (Hz(i,j,k));
v(i,j,k) *= tbxp2.contains(i,j-1,0) ? 2.0 / (Hz(i,j-1,k) + Hz(i,j,k)) : 1.0 / (Hz(i,j,k));
});
{
v(i,j,k) += cff * (pm(i,j,0)+pm(i,j-1,0)) * (pn(i,j,0)+pn(i,j-1,0)) * rv(i,j,k,nrhs);
v(i,j,k) *= 2.0 / (Hz(i,j-1,k) + Hz(i,j,k));
});

{
amrex::Gpu::synchronize();
Expand All @@ -168,26 +173,26 @@ ROMSX::advance_3d (int lev,
#ifdef AMREX_USE_GPU
Gpu::synchronize();
#endif
vert_visc_3d(gbx1,1,0,u,Hz,Hzk,oHz,AK,Akv,BC,DC,FC,CF,nnew,N,dt_lev);
vert_visc_3d(xbx,1,0,u,Hz,Hzk,oHz,AK,Akv,BC,DC,FC,CF,nnew,N,dt_lev);

// Reset to zero
mf_DC[mfi].template setVal<RunOn::Device>(0.,gbx21);
fab_CF.template setVal<RunOn::Device>(0.,gbx21);

vert_visc_3d(gbx1,0,1,v,Hz,Hzk,oHz,AK,Akv,BC,DC,FC,CF,nnew,N,dt_lev);
vert_visc_3d(ybx,0,1,v,Hz,Hzk,oHz,AK,Akv,BC,DC,FC,CF,nnew,N,dt_lev);
}

// Reset to zero
mf_DC[mfi].template setVal<RunOn::Device>(0.,gbx21);
fab_CF.template setVal<RunOn::Device>(0.,gbx21);

vert_mean_3d(gbx1,1,0,u,Hz,DU_avg1,DC,CF,pm,nnew,N);
vert_mean_3d(xbx,1,0,u,Hz,DU_avg1,DC,CF,pm,nnew,N);

// Reset to zero
mf_DC[mfi].template setVal<RunOn::Device>(0.,gbx21);
fab_CF.template setVal<RunOn::Device>(0.,gbx21);

vert_mean_3d(gbx1,0,1,v,Hz,DV_avg1,DC,CF,pn,nnew,N);
vert_mean_3d(ybx,0,1,v,Hz,DV_avg1,DC,CF,pn,nnew,N);

update_massflux_3d(gbx2,1,0,u,ubar,Huon,Hz,on_u,DU_avg1,DU_avg2,DC,FC,CF,nnew);

Expand Down Expand Up @@ -277,6 +282,7 @@ ROMSX::advance_3d (int lev,
//------------------------------------------------------------------------
//
//Should really use gbx3uneven
//TODO: go over these boxes and compare to other spots where we do the same thing
Box gbx1D = gbx1;
gbx1D.makeSlab(2,0);

Expand Down Expand Up @@ -404,7 +410,7 @@ ROMSX::advance_3d (int lev,
fab_on_u.template setVal<RunOn::Device>(dx[1],makeSlab(tbxp2,2,0));
fab_om_v.template setVal<RunOn::Device>(dx[0],makeSlab(tbxp2,2,0));

vert_visc_3d(gbx1,0,0,temp,Hz,Hzk,oHz,AK,Akt,BC,DC,FC,CF,nnew,N,dt_lev);
vert_visc_3d(gbx1,0,0,salt,Hz,Hzk,oHz,AK,Akt,BC,DC,FC,CF,nnew,N,dt_lev);
vert_visc_3d(bx,0,0,temp,Hz,Hzk,oHz,AK,Akt,BC,DC,FC,CF,nnew,N,dt_lev);
vert_visc_3d(bx,0,0,salt,Hz,Hzk,oHz,AK,Akt,BC,DC,FC,CF,nnew,N,dt_lev);
} // MFiter
}
31 changes: 12 additions & 19 deletions Source/TimeIntegration/ROMSX_vert_visc_3d.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
using namespace amrex;

//
// prestep_uv_3d
// vert_visc_3d
//

void
ROMSX::vert_visc_3d (const Box& phi_bx, const int ioff, const int joff,
Array4<Real> phi,
Array4<Real const> Hz, Array4<Real> Hzk,
Array4<Real const> Hz, /*temp var */ Array4<Real> Hzk,
Array4<Real> oHz,
Array4<Real> AK, Array4<Real> Akv,
Array4<Real> BC, Array4<Real> DC,
Expand All @@ -25,13 +25,6 @@ ROMSX::vert_visc_3d (const Box& phi_bx, const int ioff, const int joff,
ParallelFor(phi_bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
Hzk(i,j,k)=0.5*(Hz(i-ioff,j-joff,k)+Hz(i,j,k));
});

//
// Define oHz = (1/Hz)
//
ParallelFor(phi_bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
oHz(i,j,k) = 1.0/ Hzk(i,j,k);
});

Expand Down Expand Up @@ -98,16 +91,16 @@ ROMSX::vert_visc_3d (const Box& phi_bx, const int ioff, const int joff,
ParallelFor(makeSlab(phi_bx,2,0), [=] AMREX_GPU_DEVICE (int i, int j, int )
{
for(int k=0; k<=N; k++) {
//
// Backward substitution.
//
DC(i,j,N)=0.0;

if(N-k+1<=N && N>=k) //-N,1,-1 => kidx =N-k+1
{
if(N+1<k || N+2<k) amrex::Abort("-1 here");
DC(i,j,N-k) -= CF(i,j,N-k)*DC(i,j,N-k+1);
}
//
// Backward substitution.
//
DC(i,j,N)=0.0;

if(N-k+1<=N && N>=k) //-N,1,-1 => kidx =N-k+1
{
if(N+1<k || N+2<k) amrex::Abort("-1 here");
DC(i,j,N-k) -= CF(i,j,N-k)*DC(i,j,N-k+1);
}
}
});

Expand Down

0 comments on commit e289e94

Please sign in to comment.