Skip to content

Add CUDA-aware MPI halo exchange for MPMesh #88

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

Open
Shahrear2000 wants to merge 11 commits into
base: dn/Velocity_solver_improvement
Choose a base branch
from
Open

Add CUDA-aware MPI halo exchange for MPMesh #88

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b98b69a
First commit to port
dhyan1272 Jun 29, 2026
12417d3
Conditional for HALO exchange of nProcsTot=1 in polyMPO
dhyan1272 Jun 30, 2026
229c8cb
SetMPArea and modify in strain rate
dhyan1272 Jul 13, 2026
f0c25f7
Passing oceanDragCoeff and using same vel corrections as rebase 29 June
dhyan1272 Jul 14, 2026
3b6fccd
Reconstrucntion of velocity test commented out as now in src involves…
dhyan1272 Jul 14, 2026
215764a
More formatting
dhyan1272 Jul 15, 2026
ab4fa2d
Add CUDA-aware MPI halo exchange
Jun 29, 2026
476b580
Add CUDA-aware communication path to MPMesh
Jul 12, 2026
4d6dcda
Cuda Aware MPI Halo Exchange-Diagnostics lables
Jul 30, 2026
cac3f1e
Cuda Aware MPI Halo Exchange - Optimize: derived-datatype send (pack …
Aug 12, 2026
cb5fcea
Cuda Aware Halo Exchange - Diagnostic labels with Barrier-off
Aug 13, 2026
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
2 changes: 1 addition & 1 deletion CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,4 +50,4 @@ if(IS_TESTING)
add_subdirectory (test)
endif()

bob_end_package()
bob_end_package()
19 changes: 19 additions & 0 deletions src/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,25 @@ set(SOURCES
)

add_library(polyMPO-core ${SOURCES})

if(DEFINED ENV{PE_MPICH_GTL_DIR_nvidia80}
AND DEFINED ENV{PE_MPICH_GTL_LIBS_nvidia80})

target_link_options(polyMPO-core PUBLIC
"$ENV{PE_MPICH_GTL_DIR_nvidia80}"
"$ENV{PE_MPICH_GTL_LIBS_nvidia80}"
)

message(STATUS
"polyMPO: enabling Cray MPICH CUDA GTL linkage")
else()
message(FATAL_ERROR
"Cray MPICH CUDA GTL environment variables are unavailable. "
"Load cudatoolkit and craype-accel-nvidia80, or set "
"CRAY_ACCEL_TARGET=nvidia80 before configuring.")
endif()

target_compile_definitions(polyMPO-core PUBLIC CUDA_AWARE_MPI)

@cwsmith cwsmith Jun 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naming this 'GPU_AWARE_MPI' would be more portable (i.e., to AMD and Intel GPUs).

set_property(TARGET polyMPO-core PROPERTY CXX_STANDARD "17")
set_property(TARGET polyMPO-core PROPERTY CXX_STANDARD_REQUIRED ON)
set_property(TARGET polyMPO-core PROPERTY CXX_EXTENSIONS OFF)
Expand Down
85 changes: 67 additions & 18 deletions src/pmpo_MPMesh.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,13 @@ void MPMesh::calculateStrain(){
auto MPsBasisGrads = p_MPs->getData<MPF_Basis_Grad_Vals>();
auto MPsAppID = p_MPs->getData<MPF_MP_APP_ID>();
auto MPsStrainRate = p_MPs->getData<MPF_Strain_Rate>();
auto MPsArea = p_MPs->getData<polyMPO::MPF_Area>();
//Mesh Fields
auto tanLatVertexRotatedOverRadius = p_mesh->getMeshField<MeshF_TanLatVertexRotatedOverRadius>();
auto elm2VtxConn = p_mesh->getElm2VtxConn();
auto velField = p_mesh->getMeshField<MeshF_Vel>();
auto solveStress = p_mesh->getMeshField<polyMPO::MeshF_SolveStress>();
auto elasticTimeStep = p_mesh->getElasticTimeStep();

auto setMPStrainRate = PS_LAMBDA(const int& elm, const int& mp, const int& mask){
if(mask){
Expand All @@ -38,7 +40,7 @@ void MPMesh::calculateStrain(){
double v22 = 0.0;
double uTanOverR = 0.0;
double vTanOverR = 0.0;

for (int i = 0; i < numVtx; i++){
int iVertex = elm2VtxConn(elm, i+1)-1;
v11 = v11 + MPsBasisGrads(mp, i*2 + 0) * velField(iVertex, 0);
Expand All @@ -52,6 +54,8 @@ void MPMesh::calculateStrain(){
MPsStrainRate(mp, 0) = v11 - vTanOverR;
MPsStrainRate(mp, 1) = v22;
MPsStrainRate(mp, 2) = 0.5*(v12 + v21 + uTanOverR);

MPsArea(mp, 0) = MPsArea(mp, 0) * exp((v11+v22-vTanOverR)*elasticTimeStep);
}
};
p_MPs->parallel_for(setMPStrainRate, "setMPStrainRate");
Expand All @@ -75,25 +79,38 @@ void MPMesh::calculateStress(const int constitutive_relation){
if(mask){
Vec3d strain_rate (MPsStrainRate(mp, 0), MPsStrainRate(mp, 1), MPsStrainRate(mp, 2));
Vec3d stress(MPsStress(mp, 0), MPsStress(mp, 1), MPsStress(mp, 2));
double rep_pressure=MPsRepPressure(mp,0);

if (constitutive_relation == 1)
constitutive_evp(strain_rate, stress, MPsIcePressure(mp,0), MPsRepPressure(mp,0), MPsArea(mp,0), elasticTimeStep, dampingTimescale);
constitutive_evp(strain_rate, stress, MPsIcePressure(mp,0), rep_pressure, MPsArea(mp,0), elasticTimeStep, dampingTimescale);
else if(constitutive_relation == 3)
constitutive_linear(strain_rate, stress);

for (int m=0 ; m<3; m++)
MPsStress(mp, m) = stress[m]*solveStress(elm);
MPsRepPressure(mp,0)=rep_pressure;
}
};
p_MPs->parallel_for(setMPStress, "setMPStress");
}

void MPMesh::calculateStressDivergence(){

Kokkos::Timer timer;
int self, numProcsTot;
MPI_Comm comm = p_MPs->getMPIComm();
MPI_Comm_rank(comm, &self);
MPI_Comm_size(comm, &numProcsTot);

//Kokkos::Timer b0Timer;
//Kokkos::fence();

//const double b0DrainTime = b0Timer.seconds(); //B0: drain any device work left from the previous phase
//MPI_Barrier(comm);
//const double b0Sync = b0Timer.seconds(); //B0: align all ranks so this timed region starts at the same instant
//const double b0BarrierWait = b0Sync - b0DrainTime;

Kokkos::Timer totalTimer;
Kokkos::Timer computeTimer;

//Mesh Information
auto elm2VtxConn = p_mesh->getElm2VtxConn();
Expand Down Expand Up @@ -145,7 +162,7 @@ void MPMesh::calculateStressDivergence(){
(1.0 - ramp) * invM * w_vtx;

factor = factor * tanLatVertexRotatedOverRadius(vID, 0);

auto factor1 = ramp * (w_vtx/radius) * (VtxCoeffs_new(vID, 1, 0) + VtxCoeffs_new(vID, 1, 1)*CoordDiffs[1] +
VtxCoeffs_new(vID, 1, 2)*CoordDiffs[2] +
VtxCoeffs_new(vID, 1, 3)*CoordDiffs[3]) -
Expand All @@ -165,18 +182,44 @@ void MPMesh::calculateStressDivergence(){
}
};
p_MPs->parallel_for(stress_div, " stress_div_assembly");
Kokkos::fence();
pumipic::RecordTime("Stress_Divergence_Reconstruction" + std::to_string(self), timer.seconds());
Kokkos::fence(); //drain device work -> compute really is done on this rank

const double computeTime = computeTimer.seconds(); //T_before: pure local compute time, no waiting

//MPI_Barrier(comm); //B1: fast ranks wait here for the slowest rank

//const double computeTimeSync = computeTimer.seconds(); //time until every rank reached the barrier
//const double b1BarrierWait = computeTimeSync - computeTime; //this rank's wait time = compute load imbalance

Kokkos::Timer communicationTimer;

timer.reset();
if(numProcsTot>1){
//Takes contribution of halo vertices and adds it in owner procs
communicate_and_take_halo_contributions1(stress_divUV, numVertices, 2, 0, 0);
communicate_and_take_halo_contributions1_improved(stress_divUV, numVertices, 2, 0, 0, "Stress_Divergence");
//Transfer the correct values at owned vertices to halo vertices
//communicate_and_take_halo_contributions(stress_divUV, numVertices, 2, 1, 1);
}
Kokkos::fence();
pumipic::RecordTime("Stress_Divergence Communication" + std::to_string(self), timer.seconds());
Kokkos::fence(); //drain device work from the communication step

const double communicationTime = communicationTimer.seconds(); //pure local communication time, no waiting

//MPI_Barrier(comm); //B2: fast ranks wait here for the slowest rank

//const double communicationTimeSync = communicationTimer.seconds();
//const double b2BarrierWait = communicationTimeSync - communicationTime; //communication load imbalance

const double totalTime = totalTimer.seconds();

//pumipic::RecordTime("Stress_Divergence_B0_Drain_" + std::to_string(self), b0DrainTime);
//pumipic::RecordTime("Stress_Divergence_B0_BarrierWait_" + std::to_string(self), b0BarrierWait);

pumipic::RecordTime("Stress_Divergence_Compute_" + std::to_string(self),computeTime);
//pumipic::RecordTime("Stress_Divergence_Compute_B1_Barrier_Wait_Time_" + std::to_string(self),b1BarrierWait);

pumipic::RecordTime("Stress_Divergence_Communication_" + std::to_string(self),communicationTime);
//pumipic::RecordTime("Stress_Divergence_Communication_B2_Wait_Time_" + std::to_string(self), b2BarrierWait);

pumipic::RecordTime("Stress_Divergence_Total_" + std::to_string(self),totalTime);
}

void MPMesh::calcBasis() {
Expand Down Expand Up @@ -280,7 +323,7 @@ void MPMesh::CVTTrackingElmCenterBased(const int printVTPIndex){
Vec3d dx = MPnew-MP;
while(true){
int numConnElms = elm2ElmConn(iElm,0);

Vec3d center(elmCenter(iElm, 0), elmCenter(iElm, 1), elmCenter(iElm, 2));
Vec3d delta = MPnew - center;

Expand Down Expand Up @@ -441,7 +484,13 @@ void MPMesh::startCommunication(){
int self, numProcsTot;
MPI_Comm comm = p_MPs->getMPIComm();
MPI_Comm_rank(comm, &self);
MPI_Comm_size(comm, &numProcsTot);
MPI_Comm_size(comm, &numProcsTot);

std::cout << "[RankSummary] Rank=" << self
<< " Vertices(total)=" << p_mesh->getNumVertices()
<< " Vertices(owned)=" << p_mesh->getNumVerticesOwned()
<< " Elements=" << p_mesh->getNumElements()
<< std::endl;

//The routine should work for elements too, although currently the communication
//is done for vertices. For elements, the follwoing three variables should correspond
Expand Down Expand Up @@ -742,8 +791,8 @@ void MPMesh::T2LTracking(Vec2dView dx){
Vec2d MP(mpPositions(mp,0),mpPositions(mp,1));//XXX:the input is XYZ, but we only support 2d vector
if(mask){
int iElm = elm;
Vec2d MPnew = MP + dx(mp);
Vec2d MPnew = MP + dx(mp);

while(true){
int numVtx = elm2VtxConn(iElm,0);
bool goToNeighbour = false;
Expand All @@ -753,25 +802,25 @@ void MPMesh::T2LTracking(Vec2dView dx){
v[i] = elm2VtxConn(iElm,i+1)-1;
//get edges and perpendiculardx
Vec2d e[maxVtxsPerElm];
double pdx[maxVtxsPerElm];
double pdx[maxVtxsPerElm];
for(int i=0; i< numVtx; i++){
int idx_ip1 = (i+1)%numVtx;
Vec2d v_i(vtxCoords(v[i],0),vtxCoords(v[i],1));
Vec2d v_ip1(vtxCoords(v[idx_ip1],0),vtxCoords(v[idx_ip1],1));
e[i] = v_ip1 - v_i;
pdx[i] = (v_i - MP).cross(dx(mp));
}

for(int i=0; i<numVtx; i++){
int ip1 = (i+1)%numVtx;
//pdx*pdx<0 and edge is acrossed
//pdx*pdx<0 and edge is acrossed
if(pdx[i]*pdx[ip1] <0 && e[i].cross(Vec2d(MPnew[0]-vtxCoords(v[i],0),
MPnew[1]-vtxCoords(v[i],1)))<0){
//go to the next elm
iElm = elm2ElmConn(iElm,i+1);
goToNeighbour = true;
if(iElm <0){
mpStatus(mp) = 0;
mpStatus(mp) = 0;
MPs2Elm(mp) = -1;
goToNeighbour = false;
}
Expand Down
Loading