From c79b80fbbff1b4f2509ec95dce6ff8ea6ca95e20 Mon Sep 17 00:00:00 2001 From: Shahrear Jahan Santho Date: Mon, 29 Jun 2026 06:56:17 -0700 Subject: [PATCH 1/2] Add CUDA-aware MPI halo exchange --- CMakeLists.txt | 2 +- src/CMakeLists.txt | 1 + src/pmpo_MPMesh.cpp | 3 +- src/pmpo_MPMesh.hpp | 869 ++++++++++++++++++++++++++++++++++- src/pmpo_MPMesh_assembly.hpp | 7 +- src/pmpo_c.cpp | 44 ++ 6 files changed, 903 insertions(+), 23 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 080a4186..3f627daa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,4 +50,4 @@ if(IS_TESTING) add_subdirectory (test) endif() -bob_end_package() \ No newline at end of file +bob_end_package() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 76059965..d0fb8e74 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -21,6 +21,7 @@ set(SOURCES ) add_library(polyMPO-core ${SOURCES}) +target_compile_definitions(polyMPO-core PUBLIC CUDA_AWARE_MPI) 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) diff --git a/src/pmpo_MPMesh.cpp b/src/pmpo_MPMesh.cpp index 67cea389..d410a4a0 100644 --- a/src/pmpo_MPMesh.cpp +++ b/src/pmpo_MPMesh.cpp @@ -166,8 +166,7 @@ void MPMesh::calculateStressDivergence(){ if(numProcsTot>1){ //Takes contribution of halo vertices and adds it in owner procs - communicate_and_take_halo_contributions(stress_divUV, numVertices, 2, 0, 0); - communicate_and_take_halo_contributions(divUV_edge, numVertices, 2, 0, 0); + communicate_and_take_halo_contributions1_improved(stress_divUV, numVertices, 2, 0, 0); //Transfer the correct values at owned vertices to halo vertices communicate_and_take_halo_contributions(stress_divUV, numVertices, 2, 1, 1); communicate_and_take_halo_contributions(divUV_edge, numVertices, 2, 1, 1); diff --git a/src/pmpo_MPMesh.hpp b/src/pmpo_MPMesh.hpp index 8f0be27c..dfd922e5 100644 --- a/src/pmpo_MPMesh.hpp +++ b/src/pmpo_MPMesh.hpp @@ -4,6 +4,9 @@ #include "pmpo_utils.hpp" #include "pmpo_mesh.hpp" #include "pmpo_materialPoints.hpp" +#include +#include +#include namespace polyMPO{ @@ -28,26 +31,379 @@ class MPMesh{ int numOwnersTot, numHalosTot; std::vector numOwnersOnOtherProcs; std::vector numHalosOnOtherProcs; - std::vectorhaloOwnerProcs; + std::vector haloOwnerProcs; std::vector> haloOwnerLocalIDs; std::vector> ownerOwnerLocalIDs; std::vector> ownerHaloLocalIDs; void startCommunication(); - void communicateFields(const std::vector>& fieldData, const int numEntities, const int numEntries, int mode, - std::vector>& recvIDVec, std::vector>& recvDataVec); - void communicate_and_take_halo_contributions(const Kokkos::View& meshField, int nEntities, int numEntries, int mode, int op); + + void communicate_and_take_halo_contributions( + const Kokkos::View& meshField, + int nEntities, + int numEntries, + int mode, + int op); + + // Original CPU-staging function + template + void communicate_and_take_halo_contributions1( + const ViewType& meshField, + int nEntities, + int numEntries, + int mode, + int op){ + + int self; + MPI_Comm comm = p_MPs->getMPIComm(); + MPI_Comm_rank(comm, &self); + + Kokkos::Timer timer; + auto reconVals_host = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), meshField); + + pumipic::RecordTime("SD: GPU-CPU copy-" + std::to_string(self), timer.seconds()); + + timer.reset(); + std::vector> recvIDVec; + std::vector> recvDataVec; + + pumipic::RecordTime("SD: Recv Vec Allocation-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + communicateFields1( + reconVals_host, + nEntities, + numEntries, + mode, + recvIDVec, + recvDataVec); + + pumipic::RecordTime("SD: IP Comm-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + int numProcsTot = recvIDVec.size(); + + int totalSize = 0; + std::vector offsets(numProcsTot, 0); + + for(int i = 0; i < numProcsTot; i++){ + offsets[i] = totalSize; + totalSize += recvIDVec[i].size(); + } + + Kokkos::View recvIDGPU("recvIDGPU", totalSize); + auto hostView = + Kokkos::View("recvIDCPU", totalSize); + + for(int i = 0; i < numProcsTot; i++){ + std::copy( + recvIDVec[i].begin(), + recvIDVec[i].end(), + hostView.data() + offsets[i]); + } + + pumipic::RecordTime("SD: Flatten IDs-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + Kokkos::deep_copy(recvIDGPU, hostView); + Kokkos::fence(); + + pumipic::RecordTime("SD: Copy CPU-GPU-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + int totalSize_data = 0; + std::vector offsets_data(numProcsTot, 0); + + for(int i = 0; i < numProcsTot; i++){ + offsets_data[i] = totalSize_data; + totalSize_data += recvDataVec[i].size(); + } + + Kokkos::View recvDataGPU("recvDataGPU", totalSize_data); + auto hostView_data = + Kokkos::View("recvDataCPU", totalSize_data); + + for(int i = 0; i < numProcsTot; i++){ + std::copy( + recvDataVec[i].begin(), + recvDataVec[i].end(), + hostView_data.data() + offsets_data[i]); + } + + pumipic::RecordTime("SD: Flatten Data-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + Kokkos::deep_copy(recvDataGPU, hostView_data); + Kokkos::fence(); + + assert(totalSize_data == totalSize * numEntries); + + for(int i = 0; i < numProcsTot; i++){ + assert(recvDataVec[i].size() == recvIDVec[i].size() * numEntries); + } + + pumipic::RecordTime("SD: Copy CPU-GPU2-" + std::to_string(self), timer.seconds()); + + timer.reset(); + + if(op == 0){ + Kokkos::parallel_for( + "halo contribution add", + recvIDGPU.size(), + KOKKOS_LAMBDA(const int i){ + const int vertex = recvIDGPU(i); + + for(int k = 0; k < numEntries; k++){ + Kokkos::atomic_add( + &meshField(vertex,k), + recvDataGPU(i * numEntries + k)); + } + }); + } + else{ + Kokkos::parallel_for( + "halo contribution assign", + recvIDGPU.size(), + KOKKOS_LAMBDA(const int i){ + const int vertex = recvIDGPU(i); + + for(int k = 0; k < numEntries; k++){ + meshField(vertex,k) = + recvDataGPU(i * numEntries + k); + } + }); + } + + Kokkos::fence(); + + pumipic::RecordTime("SD: Contribution" + std::to_string(self), timer.seconds()); + } + + + void communicateFields( + const std::vector>& fieldData, + const int numEntities, + const int numEntries, + int mode, + std::vector>& recvIDVec, + std::vector>& recvDataVec); + + + template + void communicateFields1( + const ViewType& fieldData, + const int numEntities, + const int numEntries, + int mode, + std::vector>& recvIDVec, + std::vector>& recvDataVec){ + + int self, numProcsTot; + + MPI_Comm comm = p_MPs->getMPIComm(); + + MPI_Comm_rank(comm, &self); + MPI_Comm_size(comm, &numProcsTot); + + assert(numEntities == numOwnersTot + numHalosTot); + + std::vector> sendDataVec(numProcsTot); + + recvIDVec.resize(numProcsTot); + recvDataVec.resize(numProcsTot); + + for(int i = 0; i < numProcsTot; i++){ + if(i == self) continue; + + int numToSend = 0; + int numToRecv = 0; + + if(mode == 0){ + numToSend = numOwnersOnOtherProcs[i]; + numToRecv = numHalosOnOtherProcs[i]; + } + else{ + numToSend = numHalosOnOtherProcs[i]; + numToRecv = numOwnersOnOtherProcs[i]; + } + + if(numToSend > 0){ + sendDataVec[i].reserve(numToSend * numEntries); + } + + if(numToRecv > 0){ + recvDataVec[i].resize(numToRecv * numEntries); + recvIDVec[i].resize(numToRecv); + } + } + + if(mode == 0){ + for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ + auto ownerProc = haloOwnerProcs[iEnt]; + + for(int iDouble = 0; iDouble < numEntries; iDouble++){ + sendDataVec[ownerProc].push_back( + fieldData(numOwnersTot + iEnt, iDouble)); + } + } + } + else if(mode == 1){ + for(size_t iProc = 0; iProc < ownerOwnerLocalIDs.size(); iProc++){ + for(auto& ownerID : ownerOwnerLocalIDs[iProc]){ + for(int iDouble = 0; iDouble < numEntries; iDouble++){ + sendDataVec[iProc].push_back( + fieldData(ownerID, iDouble)); + } + } + } + } + + std::vector requests; + requests.reserve(4 * numProcsTot); + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + + if(mode == 0 && numHalosOnOtherProcs[proc]){ + assert(recvIDVec[proc].size() == + static_cast(numHalosOnOtherProcs[proc])); + + assert(recvDataVec[proc].size() == + recvIDVec[proc].size() * static_cast(numEntries)); + + MPI_Request req3; + MPI_Request req4; + + MPI_Irecv( + recvIDVec[proc].data(), + recvIDVec[proc].size(), + MPI_INT, + proc, + 1, + comm, + &req3); + + MPI_Irecv( + recvDataVec[proc].data(), + recvDataVec[proc].size(), + MPI_DOUBLE, + proc, + 2, + comm, + &req4); + + requests.push_back(req3); + requests.push_back(req4); + } + + if(mode == 0 && numOwnersOnOtherProcs[proc]){ + assert(haloOwnerLocalIDs[proc].size() == + static_cast(numOwnersOnOtherProcs[proc])); + + assert(sendDataVec[proc].size() == + haloOwnerLocalIDs[proc].size() * static_cast(numEntries)); + + MPI_Request req1; + MPI_Request req2; + + MPI_Isend( + haloOwnerLocalIDs[proc].data(), + haloOwnerLocalIDs[proc].size(), + MPI_INT, + proc, + 1, + comm, + &req1); + + MPI_Isend( + sendDataVec[proc].data(), + sendDataVec[proc].size(), + MPI_DOUBLE, + proc, + 2, + comm, + &req2); + + requests.push_back(req1); + requests.push_back(req2); + } + + if(mode == 1 && numOwnersOnOtherProcs[proc]){ + MPI_Request req3; + MPI_Request req4; + + MPI_Irecv( + recvIDVec[proc].data(), + recvIDVec[proc].size(), + MPI_INT, + proc, + 1, + comm, + &req3); + + MPI_Irecv( + recvDataVec[proc].data(), + recvDataVec[proc].size(), + MPI_DOUBLE, + proc, + 2, + comm, + &req4); + + requests.push_back(req3); + requests.push_back(req4); + } + + if(mode == 1 && numHalosOnOtherProcs[proc]){ + MPI_Request req1; + MPI_Request req2; + + MPI_Isend( + ownerHaloLocalIDs[proc].data(), + ownerHaloLocalIDs[proc].size(), + MPI_INT, + proc, + 1, + comm, + &req1); + + MPI_Isend( + sendDataVec[proc].data(), + sendDataVec[proc].size(), + MPI_DOUBLE, + proc, + 2, + comm, + &req2); + + requests.push_back(req1); + requests.push_back(req2); + } + } + + MPI_Waitall(requests.size(), requests.data(), MPI_STATUSES_IGNORE); + } + MPMesh(Mesh* inMesh, MaterialPoints* inMPs): - p_mesh(inMesh), p_MPs(inMPs) { + p_mesh(inMesh), + p_MPs(inMPs) { }; - ~MPMesh() { + + ~MPMesh(){ delete p_mesh; delete p_MPs; } - //MP advection and tracking + + // MP advection and tracking void CVTTrackingEdgeCenterBased(Vec2dView dx); void CVTTrackingElmCenterBased(const int printVTPIndex = -1); void T2LTracking(Vec2dView dx); @@ -57,40 +413,521 @@ class MPMesh{ void push_swap_pos(); void push(); - //Used before advection to interpolate fields from mesh to MPs - //And also before reconstruction + + // Used before advection to interpolate fields from mesh to MPs + // And also before reconstruction void calcBasis(); - //Reconstruction + + // Reconstruction DoubleView assemblyV0(); + template void assemblyVtx0(); + template void assemblyElm0(); + template void assemblyVtx1(); + void reconstruct_coeff_full(); - void invertMatrix(const Kokkos::View& vtxMatrices, const double& radius); + + void invertMatrix( + const Kokkos::View& vtxMatrices, + const double& radius); + Kokkos::View precomputedVtxCoeffs_new; - Kokkos::View nearAnEdge; + Kokkos::View nearAnEdge; Kokkos::View vtxMatrixMass; - //Not used currently - std::map> reconstructSlice = std::map>(); + + // Not used currently + std::map> reconstructSlice = + std::map>(); + template DoubleView wtScaAssembly(); + template Vec2dView wtVec2Assembly(); + template - void assembly(int order, MeshFieldType type, bool basisWeightFlag, bool massWeightFlag); + void assembly( + int order, + MeshFieldType type, + bool basisWeightFlag, + bool massWeightFlag); + template - void setReconstructSlice(int order, MeshFieldType type); + void setReconstructSlice( + int order, + MeshFieldType type); + void reconstructSlices(); void printVTP_mesh(int printVTPIndex); + + void writeMPTrackingVTP( + int printVTPIndex, + int numMPs, + const Vec3dView& history, + const Vec3dView& resultLeft, + const Vec3dView& resultRight, + const Vec3dView& mpTgtPosArray); + void calculateStrain(); void calculateStress(const int constitutive_relation); void calculateStressDivergence(); + + + +#ifdef CUDA_AWARE_MPI + + // Cached CUDA-aware MPI communication metadata and per-neighbor GPU buffers. + // Important change from the previous version: + // MPI is always given the base pointer of a Kokkos allocation, not + // "base pointer + offset". This avoids Cray MPICH/GTL CUDA IPC problems. + bool cudaAwareMPICacheValid = true; + bool cudaAwareMPIDisabled = false; + bool cudaAwareMPIEnvChecked = false; + bool cudaAwareMPIForceCPU = false; + bool cudaAwareMPILogged = false; + + struct CudaAwareMPIFieldCache{ + bool valid = false; + int cachedNumProcs = -1; + + std::vector sendCounts; + std::vector recvCounts; + + std::vector> sendEntityGPUPerProc; + std::vector> recvIDGPUPerProc; + + std::vector> sendDataGPUPerProc; + std::vector> recvDataGPUPerProc; + }; + + std::map, CudaAwareMPIFieldCache> cudaAwareMPICaches; + + bool cudaAwareMPIForceDisabled(){ + if(!cudaAwareMPIEnvChecked){ + const char* value = std::getenv("POLYMPO_DISABLE_CUDA_AWARE_MPI"); + cudaAwareMPIForceCPU = + value != nullptr && value[0] != '\0' && value[0] != '0'; + cudaAwareMPIEnvChecked = true; + } + + return cudaAwareMPIForceCPU; + } + + // Fully CUDA-aware MPI version: + // Field data is sent/received using GPU pointers. Receive IDs are cached + // once from the fixed halo/owner mapping and are not sent every call. + // + // Important: + // This function caches communication metadata and GPU buffers per + // (mode, numEntries). If the communication pattern changes, clear + // cudaAwareMPICaches before the next call. + template + void communicate_and_take_halo_contributions1_improved( + const ViewType& meshField, + int nEntities, + int numEntries, + int mode, + int op){ + + int self, numProcsTot; + + MPI_Comm comm = p_MPs->getMPIComm(); + + MPI_Comm_rank(comm, &self); + MPI_Comm_size(comm, &numProcsTot); + + assert(mode == 0 || mode == 1); + assert(op == 0 || op == 1); + assert(nEntities == numOwnersTot + numHalosTot); + + if(cudaAwareMPIDisabled || cudaAwareMPIForceDisabled()){ + communicate_and_take_halo_contributions1( + meshField, + nEntities, + numEntries, + mode, + op); + return; + } + +#ifdef POLYMPO_VERBOSE_MPI + if(self == 0 && !cudaAwareMPILogged){ + std::cout + << "[CUDA_AWARE_MPI] Using per-proc cached full GPU-aware MPI path in communicate_and_take_halo_contributions1_improved()" + << "\n"; + cudaAwareMPILogged = true; + } +#endif + + Kokkos::Timer timer; + + if(!cudaAwareMPICacheValid){ + cudaAwareMPICaches.clear(); + cudaAwareMPICacheValid = true; + } + + auto& cudaAwareCache = + cudaAwareMPICaches[std::make_pair(mode, numEntries)]; + + const bool needRebuild = + (!cudaAwareCache.valid) || + (cudaAwareCache.cachedNumProcs != numProcsTot); + + if(needRebuild){ + + cudaAwareCache.cachedNumProcs = numProcsTot; + + cudaAwareCache.sendCounts.assign(numProcsTot, 0); + cudaAwareCache.recvCounts.assign(numProcsTot, 0); + + cudaAwareCache.sendEntityGPUPerProc.clear(); + cudaAwareCache.recvIDGPUPerProc.clear(); + cudaAwareCache.sendDataGPUPerProc.clear(); + cudaAwareCache.recvDataGPUPerProc.clear(); + + cudaAwareCache.sendEntityGPUPerProc.resize(numProcsTot); + cudaAwareCache.recvIDGPUPerProc.resize(numProcsTot); + cudaAwareCache.sendDataGPUPerProc.resize(numProcsTot); + cudaAwareCache.recvDataGPUPerProc.resize(numProcsTot); + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + + if(mode == 0){ + cudaAwareCache.sendCounts[proc] = numOwnersOnOtherProcs[proc]; + cudaAwareCache.recvCounts[proc] = numHalosOnOtherProcs[proc]; + } + else{ + cudaAwareCache.sendCounts[proc] = numHalosOnOtherProcs[proc]; + cudaAwareCache.recvCounts[proc] = numOwnersOnOtherProcs[proc]; + } + } + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + + const int sendCount = cudaAwareCache.sendCounts[proc]; + const int recvCount = cudaAwareCache.recvCounts[proc]; + + if(sendCount > 0){ + cudaAwareCache.sendEntityGPUPerProc[proc] = + Kokkos::View( + "cudaAwareMPISendEntityGPUPerProc", + sendCount); + + cudaAwareCache.sendDataGPUPerProc[proc] = + Kokkos::View( + "cudaAwareMPISendDataGPUPerProc", + sendCount * numEntries); + + auto sendEntityCPU = + Kokkos::View( + "sendEntityCPU", + sendCount); + + if(mode == 0){ + assert(haloOwnerLocalIDs[proc].size() == + static_cast(sendCount)); + + int localIndex = 0; + + for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ + int ownerProc = haloOwnerProcs[iEnt]; + + if(ownerProc != proc) continue; + + assert(localIndex < sendCount); + + sendEntityCPU(localIndex) = numOwnersTot + iEnt; + + localIndex++; + } + + assert(localIndex == sendCount); + } + else{ + assert(ownerOwnerLocalIDs[proc].size() == + static_cast(sendCount)); + + for(int i = 0; i < sendCount; i++){ + sendEntityCPU(i) = ownerOwnerLocalIDs[proc][i]; + } + } + + Kokkos::deep_copy( + cudaAwareCache.sendEntityGPUPerProc[proc], + sendEntityCPU); + } + + if(recvCount > 0){ + cudaAwareCache.recvIDGPUPerProc[proc] = + Kokkos::View( + "cudaAwareMPIRecvIDGPUPerProc", + recvCount); + + cudaAwareCache.recvDataGPUPerProc[proc] = + Kokkos::View( + "cudaAwareMPIRecvDataGPUPerProc", + recvCount * numEntries); + + auto recvIDCPU = + Kokkos::View( + "recvIDCPU", + recvCount); + + if(mode == 0){ + assert(ownerOwnerLocalIDs[proc].size() == + static_cast(recvCount)); + + for(int i = 0; i < recvCount; i++){ + recvIDCPU(i) = ownerOwnerLocalIDs[proc][i]; + } + } + else{ + int localIndex = 0; + + for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ + if(haloOwnerProcs[iEnt] != proc) continue; + + assert(localIndex < recvCount); + + recvIDCPU(localIndex) = numOwnersTot + iEnt; + + localIndex++; + } + + assert(localIndex == recvCount); + } + + Kokkos::deep_copy( + cudaAwareCache.recvIDGPUPerProc[proc], + recvIDCPU); + } + } + + Kokkos::fence(); + + cudaAwareCache.valid = true; + + pumipic::RecordTime( + "SD: CUDA-aware MPI Cache Build m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + timer.seconds()); + + timer.reset(); + } + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + if(cudaAwareCache.sendCounts[proc] <= 0) continue; + + auto sendEntityGPU = cudaAwareCache.sendEntityGPUPerProc[proc]; + auto sendDataGPU = cudaAwareCache.sendDataGPUPerProc[proc]; + int sendCount = cudaAwareCache.sendCounts[proc]; + + Kokkos::parallel_for( + "pack cached cuda-aware mpi send buffer per proc", + sendCount, + KOKKOS_LAMBDA(const int i){ + int entity = sendEntityGPU(i); + + for(int k = 0; k < numEntries; k++){ + sendDataGPU(i * numEntries + k) = + meshField(entity, k); + } + }); + } + + Kokkos::fence(); + + pumipic::RecordTime( + "SD: CUDA-aware MPI Pack m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + timer.seconds()); + + timer.reset(); + + Kokkos::Timer mpiTotalTimer; + std::vector requests; + requests.reserve(2 * numProcsTot); + int mpiError = MPI_SUCCESS; + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + + if(cudaAwareCache.recvCounts[proc] > 0){ + MPI_Request reqData; + + mpiError = MPI_Irecv( + cudaAwareCache.recvDataGPUPerProc[proc].data(), + cudaAwareCache.recvCounts[proc] * numEntries, + MPI_DOUBLE, + proc, + 2, + comm, + &reqData); + if(mpiError != MPI_SUCCESS) break; + requests.push_back(reqData); + } + + if(cudaAwareCache.sendCounts[proc] > 0){ + MPI_Request reqData; + + mpiError = MPI_Isend( + cudaAwareCache.sendDataGPUPerProc[proc].data(), + cudaAwareCache.sendCounts[proc] * numEntries, + MPI_DOUBLE, + proc, + 2, + comm, + &reqData); + if(mpiError != MPI_SUCCESS) break; + requests.push_back(reqData); + } + } + + pumipic::RecordTime( + "SD: CUDA-aware MPI Post m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + timer.seconds()); + + timer.reset(); + + if(mpiError == MPI_SUCCESS && !requests.empty()){ + mpiError = MPI_Waitall( + static_cast(requests.size()), + requests.data(), + MPI_STATUSES_IGNORE); + } + + pumipic::RecordTime( + "SD: CUDA-aware MPI Wait m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + timer.seconds()); + + if(mpiError != MPI_SUCCESS){ + cudaAwareMPIDisabled = true; + + if(self == 0){ + std::cout + << "[CUDA_AWARE_MPI] Device-pointer MPI failed." + << std::endl; + } + + if(requests.empty()){ + if(self == 0){ + std::cout + << "[CUDA_AWARE_MPI] Falling back to CPU-staged communication." + << std::endl; + } + + communicate_and_take_halo_contributions1( + meshField, + nEntities, + numEntries, + mode, + op); + return; + } + + if(self == 0){ + std::cout + << "[CUDA_AWARE_MPI] Failure happened after MPI requests were posted. " + << "Set POLYMPO_DISABLE_CUDA_AWARE_MPI=1 before running to force the CPU-staged path." + << std::endl; + } + + MPI_Abort(comm, mpiError); + return; + } + + pumipic::RecordTime( + "SD: CUDA-aware MPI Comm m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + mpiTotalTimer.seconds()); + + timer.reset(); + + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + if(cudaAwareCache.recvCounts[proc] <= 0) continue; + + auto recvIDGPU = cudaAwareCache.recvIDGPUPerProc[proc]; + auto recvDataGPU = cudaAwareCache.recvDataGPUPerProc[proc]; + int recvCount = cudaAwareCache.recvCounts[proc]; + + if(op == 0){ + Kokkos::parallel_for( + "halo add cached cuda-aware mpi per proc", + recvCount, + KOKKOS_LAMBDA(const int i){ + const int vertex = recvIDGPU(i); + + for(int k = 0; k < numEntries; k++){ +#ifdef POLYMPO_ASSUME_UNIQUE_HALO_CONTRIBS + meshField(vertex, k) += + recvDataGPU(i * numEntries + k); +#else + Kokkos::atomic_add( + &meshField(vertex, k), + recvDataGPU(i * numEntries + k)); +#endif + } + }); + } + else{ + Kokkos::parallel_for( + "halo assign cached cuda-aware mpi per proc", + recvCount, + KOKKOS_LAMBDA(const int i){ + const int vertex = recvIDGPU(i); + + for(int k = 0; k < numEntries; k++){ + meshField(vertex, k) = + recvDataGPU(i * numEntries + k); + } + }); + } + } + + Kokkos::fence(); + + pumipic::RecordTime( + "SD: CUDA-aware MPI Contribution m" + std::to_string(mode) + + " e" + std::to_string(numEntries) + "-" + std::to_string(self), + timer.seconds()); + } + +#else + + // Fallback path: + // if CUDA_AWARE_MPI is not defined, use the original GPU-CPU staging function. + template + void communicate_and_take_halo_contributions1_improved( + const ViewType& meshField, + int nEntities, + int numEntries, + int mode, + int op){ + + communicate_and_take_halo_contributions1( + meshField, + nEntities, + numEntries, + mode, + op); + } + +#endif + }; }//namespace polyMPO end diff --git a/src/pmpo_MPMesh_assembly.hpp b/src/pmpo_MPMesh_assembly.hpp index 7bbee9f2..d41aa56d 100644 --- a/src/pmpo_MPMesh_assembly.hpp +++ b/src/pmpo_MPMesh_assembly.hpp @@ -161,10 +161,10 @@ void MPMesh::reconstruct_coeff_full(){ int mode = 0; int op = 0; if (numProcsTot >1){ - communicate_and_take_halo_contributions(vtxMatrices, numVertices, numEntriesMatrix, mode, op); + communicate_and_take_halo_contributions1_improved(vtxMatrices, numVertices, numEntriesMatrix, mode, op); mode=1; op=1; - communicate_and_take_halo_contributions(vtxMatrices, numVertices, numEntriesMatrix, mode, op); + communicate_and_take_halo_contributions1_improved(vtxMatrices, numVertices, numEntriesMatrix, mode, op); } pumipic::RecordTime("Communicate Matrix Values" + std::to_string(self), timer.seconds()); @@ -380,8 +380,7 @@ void MPMesh::assemblyVtx1(){ timer.reset(); if(numProcsTot>1){ - communicate_and_take_halo_contributions(meshField, numVertices, numEntries, 0, 0); - communicate_and_take_halo_contributions(meshField, numVertices, numEntries, 1, 1); + communicate_and_take_halo_contributions1_improved(meshField, numVertices, numEntries, 0, 0); } pumipic::RecordTime("Communicate Field Values" + std::to_string(self), timer.seconds()); } diff --git a/src/pmpo_c.cpp b/src/pmpo_c.cpp index 6d4d8f77..9c0ea163 100644 --- a/src/pmpo_c.cpp +++ b/src/pmpo_c.cpp @@ -1519,6 +1519,50 @@ void polympo_velocity_grid_solve_f(MPMesh_ptr p_mpmesh){ p_mesh->gridSolveGPU(); } +void polympo_set_boundary_normal_vertex_f(MPMesh_ptr p_mpmesh, const int nComps, const int nVertices, double* uArray, double* vArray){ + auto p_mesh = ((polyMPO::MPMesh*)p_mpmesh)->p_mesh; + + kkViewHostU arrayHostU(uArray, nComps, nVertices); + kkViewHostU arrayHostV(vArray, nComps, nVertices); + + Kokkos::View uBdryVtxNormal_device("uBdryVtxNormal", nComps, nVertices); + Kokkos::View vBdryVtxNormal_device("vBdryVtxNormal", nComps, nVertices); + + Kokkos::deep_copy(uBdryVtxNormal_device, arrayHostU); + Kokkos::deep_copy(vBdryVtxNormal_device, arrayHostV); + + auto uBdryVtxNormal = p_mesh->getMeshField(); + auto vBdryVtxNormal = p_mesh->getMeshField(); + + Kokkos::parallel_for("set elm2ElmConn", nVertices, KOKKOS_LAMBDA(const int vtx){ + for(int i=0; ip_mesh; + p_mesh->applyFreeSlipBC(); +} + +void polympo_set_halo_vel_from_owner_f(MPMesh_ptr p_mpmesh){ + + int numProcsTot; + auto p_MPs = ((polyMPO::MPMesh*)p_mpmesh)->p_MPs; + MPI_Comm comm = p_MPs->getMPIComm(); + MPI_Comm_size(comm, &numProcsTot); + if(numProcsTot == 1) return; + + auto mpMesh = ((polyMPO::MPMesh*)p_mpmesh); + auto p_mesh = ((polyMPO::MPMesh*)p_mpmesh)->p_mesh; + int numVertices = p_mesh->getNumVertices(); + auto vtxFieldVel = p_mesh->getMeshField(); + + mpMesh->communicate_and_take_halo_contributions1_improved(vtxFieldVel, numVertices, 2, 1, 1); +} + //Advection Calcualtions void polympo_push_f(MPMesh_ptr p_mpmesh){ checkMPMeshValid(p_mpmesh); From 37b4533dfbf205079f2ef4d7366ce35d16e3ad99 Mon Sep 17 00:00:00 2001 From: Shahrear Jahan Santho Date: Sat, 11 Jul 2026 21:59:58 -0700 Subject: [PATCH 2/2] Add CUDA-aware communication path to MPMesh --- src/CMakeLists.txt | 18 ++ src/pmpo_MPMesh.hpp | 366 ++++++++++++++++++++++++----------- src/pmpo_MPMesh_assembly.hpp | 15 +- 3 files changed, 284 insertions(+), 115 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index d0fb8e74..5eee8a31 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -21,6 +21,24 @@ 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) set_property(TARGET polyMPO-core PROPERTY CXX_STANDARD "17") set_property(TARGET polyMPO-core PROPERTY CXX_STANDARD_REQUIRED ON) diff --git a/src/pmpo_MPMesh.hpp b/src/pmpo_MPMesh.hpp index dfd922e5..094289de 100644 --- a/src/pmpo_MPMesh.hpp +++ b/src/pmpo_MPMesh.hpp @@ -7,6 +7,11 @@ #include #include #include +#include +#include +#ifdef KOKKOS_ENABLE_CUDA +#include +#endif namespace polyMPO{ @@ -484,10 +489,106 @@ class MPMesh{ #ifdef CUDA_AWARE_MPI - // Cached CUDA-aware MPI communication metadata and per-neighbor GPU buffers. - // Important change from the previous version: - // MPI is always given the base pointer of a Kokkos allocation, not - // "base pointer + offset". This avoids Cray MPICH/GTL CUDA IPC problems. + // Use explicit CUDA memory space for MPI device buffers when CUDA is + // enabled. This avoids ambiguity in the default Kokkos::View memory space + // and gives Cray MPICH/GTL plain CUDA allocations to register/export. +#ifdef KOKKOS_ENABLE_CUDA + // Plain cudaMalloc'd device buffer, exposed as an unmanaged Kokkos::View + // via .view(). Used only for the 4 buffers below that get handed + // directly to MPI_Isend/Irecv under CUDA-aware MPI. + // + // Why not just a Kokkos::View: when Kokkos is + // built with Kokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC=ON (the default since + // Kokkos 4.2), View allocations use cudaMallocAsync/memory pools. + // cuIpcGetMemHandle (which Cray MPICH/GTL uses for intra-node GPU-to-GPU + // sends) rejects pool allocations with CUDA_ERROR_INVALID_VALUE. Since + // polyMPO isn't allowed to touch the Kokkos build config, these 4 + // buffers bypass Kokkos's allocator entirely via a direct cudaMalloc, + // which cuIpcGetMemHandle always accepts, regardless of how the rest of + // Kokkos (or the rest of the app's Views) is configured. + template + struct RawCudaMPIBuffer{ + T* ptr = nullptr; + size_t count = 0; + + void allocate(size_t n){ + free(); + count = n; + if(n > 0){ + cudaError_t err = cudaMalloc(&ptr, n * sizeof(T)); + if(err != cudaSuccess){ + throw std::runtime_error( + std::string("RawCudaMPIBuffer: cudaMalloc failed: ") + + cudaGetErrorString(err)); + } + } + } + + void free(){ + if(ptr != nullptr){ cudaFree(ptr); ptr = nullptr; } + count = 0; + } + + T* data() const{ return ptr; } + size_t size() const{ return count; } + + Kokkos::View view() const{ + return Kokkos::View( + ptr, count); + } + + RawCudaMPIBuffer() = default; + ~RawCudaMPIBuffer(){ free(); } + + RawCudaMPIBuffer(const RawCudaMPIBuffer&) = delete; + RawCudaMPIBuffer& operator=(const RawCudaMPIBuffer&) = delete; + + RawCudaMPIBuffer(RawCudaMPIBuffer&& other) noexcept{ + ptr = other.ptr; count = other.count; + other.ptr = nullptr; other.count = 0; + } + + RawCudaMPIBuffer& operator=(RawCudaMPIBuffer&& other) noexcept{ + if(this != &other){ + free(); + ptr = other.ptr; count = other.count; + other.ptr = nullptr; other.count = 0; + } + return *this; + } + }; + + using CudaAwareMPIIntBuffer = RawCudaMPIBuffer; + using CudaAwareMPIDoubleBuffer = RawCudaMPIBuffer; +#else + // No CUDA backend: no CUDA IPC/pool-allocation concern, so just wrap a + // normal Kokkos::View with the same .allocate()/.data()/.view() + // interface as RawCudaMPIBuffer above, so the cache struct and its call + // sites below don't need to branch on KOKKOS_ENABLE_CUDA. + template + struct KokkosMPIBuffer{ + Kokkos::View v; + + void allocate(size_t n){ + v = Kokkos::View("cudaAwareMPIBuffer_batched", n); + } + + T* data() const{ return v.data(); } + size_t size() const{ return v.extent(0); } + Kokkos::View view() const{ return v; } + }; + + using CudaAwareMPIIntBuffer = KokkosMPIBuffer; + using CudaAwareMPIDoubleBuffer = KokkosMPIBuffer; +#endif + + // Cached CUDA-aware MPI communication metadata and batched GPU buffers. + // Every neighbor's data lives in one shared allocation (see + // CudaAwareMPIFieldCache below) and MPI is given "base pointer + byte + // offset" per neighbor rather than a separate allocation per neighbor. + // If you ever need to fall back to one allocation per neighbor (e.g. an + // MPI/GPU stack that mishandles offset device pointers for CUDA IPC), + // restore the per-proc-buffer version from version control. bool cudaAwareMPICacheValid = true; bool cudaAwareMPIDisabled = false; bool cudaAwareMPIEnvChecked = false; @@ -500,12 +601,23 @@ class MPMesh{ std::vector sendCounts; std::vector recvCounts; - - std::vector> sendEntityGPUPerProc; - std::vector> recvIDGPUPerProc; - - std::vector> sendDataGPUPerProc; - std::vector> recvDataGPUPerProc; + std::vector sendOffsets; // prefix sum of sendCounts, in entities + std::vector recvOffsets; // prefix sum of recvCounts, in entities + + int totalSendCount = 0; + int totalRecvCount = 0; + + // Single batched GPU buffers (one allocation each, instead of one + // Kokkos::View per neighbor proc). Per-proc slices are + // [offset, offset + count) for the ID buffers, and + // [offset * numEntries, (offset + count) * numEntries) for the data + // buffers. MPI is given "buffer base pointer + offset", not a + // separate allocation per proc. + CudaAwareMPIIntBuffer sendEntityGPU; + CudaAwareMPIIntBuffer recvIDGPU; + + CudaAwareMPIDoubleBuffer sendDataGPU; + CudaAwareMPIDoubleBuffer recvDataGPU; }; std::map, CudaAwareMPIFieldCache> cudaAwareMPICaches; @@ -521,10 +633,30 @@ class MPMesh{ return cudaAwareMPIForceCPU; } - // Fully CUDA-aware MPI version: + // Fully CUDA-aware MPI version, batched buffer variant: // Field data is sent/received using GPU pointers. Receive IDs are cached // once from the fixed halo/owner mapping and are not sent every call. // + // Every neighbor's send/recv entity-ID list and data live in ONE big + // GPU buffer each (laid out back-to-back in proc order), instead of one + // Kokkos::View allocation per neighbor. Packing/unpacking is a single + // kernel launch over all neighbors' entities at once instead of one + // launch per neighbor, and MPI_Isend/Irecv use "buffer base pointer + + // offset" into that single buffer per proc. This is what actually + // shrinks MPI_Wait time: fewer, larger, more uniform in-flight + // transfers instead of many small independent ones. + // + // Note: an earlier version of this cache used one Kokkos::View + // allocation per neighbor specifically to avoid handing MPI a + // "base pointer + offset" GPU address, out of concern for CUDA IPC + // issues on Cray MPICH/GTL. That failure mode was root-caused to + // Kokkos allocating device Views via cudaMallocAsync (invalid for + // cuIpcGetMemHandle), not to offset pointers themselves, and is fixed + // by building Kokkos with -DKokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC=OFF. + // If you ever do hit IPC trouble that tracks back to offset pointers + // specifically, the per-proc-buffer version can be restored from + // version control. + // // Important: // This function caches communication metadata and GPU buffers per // (mode, numEntries). If the communication pattern changes, clear @@ -561,7 +693,7 @@ class MPMesh{ #ifdef POLYMPO_VERBOSE_MPI if(self == 0 && !cudaAwareMPILogged){ std::cout - << "[CUDA_AWARE_MPI] Using per-proc cached full GPU-aware MPI path in communicate_and_take_halo_contributions1_improved()" + << "[CUDA_AWARE_MPI] Using batched single-buffer GPU-aware MPI path in communicate_and_take_halo_contributions1_improved()" << "\n"; cudaAwareMPILogged = true; } @@ -587,16 +719,8 @@ class MPMesh{ cudaAwareCache.sendCounts.assign(numProcsTot, 0); cudaAwareCache.recvCounts.assign(numProcsTot, 0); - - cudaAwareCache.sendEntityGPUPerProc.clear(); - cudaAwareCache.recvIDGPUPerProc.clear(); - cudaAwareCache.sendDataGPUPerProc.clear(); - cudaAwareCache.recvDataGPUPerProc.clear(); - - cudaAwareCache.sendEntityGPUPerProc.resize(numProcsTot); - cudaAwareCache.recvIDGPUPerProc.resize(numProcsTot); - cudaAwareCache.sendDataGPUPerProc.resize(numProcsTot); - cudaAwareCache.recvDataGPUPerProc.resize(numProcsTot); + cudaAwareCache.sendOffsets.assign(numProcsTot, 0); + cudaAwareCache.recvOffsets.assign(numProcsTot, 0); for(int proc = 0; proc < numProcsTot; proc++){ if(proc == self) continue; @@ -611,109 +735,124 @@ class MPMesh{ } } + int totalSend = 0; + int totalRecv = 0; + for(int proc = 0; proc < numProcsTot; proc++){ - if(proc == self) continue; + cudaAwareCache.sendOffsets[proc] = totalSend; + totalSend += cudaAwareCache.sendCounts[proc]; + + cudaAwareCache.recvOffsets[proc] = totalRecv; + totalRecv += cudaAwareCache.recvCounts[proc]; + } - const int sendCount = cudaAwareCache.sendCounts[proc]; - const int recvCount = cudaAwareCache.recvCounts[proc]; + cudaAwareCache.totalSendCount = totalSend; + cudaAwareCache.totalRecvCount = totalRecv; - if(sendCount > 0){ - cudaAwareCache.sendEntityGPUPerProc[proc] = - Kokkos::View( - "cudaAwareMPISendEntityGPUPerProc", - sendCount); + cudaAwareCache.sendEntityGPU.allocate(totalSend); + cudaAwareCache.sendDataGPU.allocate(totalSend * numEntries); + cudaAwareCache.recvIDGPU.allocate(totalRecv); + cudaAwareCache.recvDataGPU.allocate(totalRecv * numEntries); - cudaAwareCache.sendDataGPUPerProc[proc] = - Kokkos::View( - "cudaAwareMPISendDataGPUPerProc", - sendCount * numEntries); + // ---- Build the flattened send-entity list (host, then one deep_copy) ---- + if(totalSend > 0){ + auto sendEntityCPU = + Kokkos::View( + "sendEntityCPU_batched", totalSend); - auto sendEntityCPU = - Kokkos::View( - "sendEntityCPU", - sendCount); + if(mode == 0){ + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + if(cudaAwareCache.sendCounts[proc] <= 0) continue; - if(mode == 0){ assert(haloOwnerLocalIDs[proc].size() == - static_cast(sendCount)); + static_cast(cudaAwareCache.sendCounts[proc])); + } - int localIndex = 0; + std::vector cursor(cudaAwareCache.sendOffsets); - for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ - int ownerProc = haloOwnerProcs[iEnt]; + for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ + int ownerProc = haloOwnerProcs[iEnt]; + if(ownerProc == self) continue; - if(ownerProc != proc) continue; + sendEntityCPU(cursor[ownerProc]) = numOwnersTot + iEnt; + cursor[ownerProc]++; + } - assert(localIndex < sendCount); + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; - sendEntityCPU(localIndex) = numOwnersTot + iEnt; + assert(cursor[proc] == + cudaAwareCache.sendOffsets[proc] + + cudaAwareCache.sendCounts[proc]); + } + } + else{ + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; - localIndex++; - } + int sendCount = cudaAwareCache.sendCounts[proc]; + if(sendCount <= 0) continue; - assert(localIndex == sendCount); - } - else{ assert(ownerOwnerLocalIDs[proc].size() == static_cast(sendCount)); + int base = cudaAwareCache.sendOffsets[proc]; + for(int i = 0; i < sendCount; i++){ - sendEntityCPU(i) = ownerOwnerLocalIDs[proc][i]; + sendEntityCPU(base + i) = ownerOwnerLocalIDs[proc][i]; } } - - Kokkos::deep_copy( - cudaAwareCache.sendEntityGPUPerProc[proc], - sendEntityCPU); } - if(recvCount > 0){ - cudaAwareCache.recvIDGPUPerProc[proc] = - Kokkos::View( - "cudaAwareMPIRecvIDGPUPerProc", - recvCount); + Kokkos::deep_copy(cudaAwareCache.sendEntityGPU.view(), sendEntityCPU); + } - cudaAwareCache.recvDataGPUPerProc[proc] = - Kokkos::View( - "cudaAwareMPIRecvDataGPUPerProc", - recvCount * numEntries); + // ---- Build the flattened recv-ID list (host, then one deep_copy) ---- + if(totalRecv > 0){ + auto recvIDCPU = + Kokkos::View( + "recvIDCPU_batched", totalRecv); - auto recvIDCPU = - Kokkos::View( - "recvIDCPU", - recvCount); + if(mode == 0){ + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; + + int recvCount = cudaAwareCache.recvCounts[proc]; + if(recvCount <= 0) continue; - if(mode == 0){ assert(ownerOwnerLocalIDs[proc].size() == static_cast(recvCount)); + int base = cudaAwareCache.recvOffsets[proc]; + for(int i = 0; i < recvCount; i++){ - recvIDCPU(i) = ownerOwnerLocalIDs[proc][i]; + recvIDCPU(base + i) = ownerOwnerLocalIDs[proc][i]; } } - else{ - int localIndex = 0; - - for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ - if(haloOwnerProcs[iEnt] != proc) continue; + } + else{ + std::vector cursor(cudaAwareCache.recvOffsets); - assert(localIndex < recvCount); + for(int iEnt = 0; iEnt < numHalosTot; iEnt++){ + int ownerProc = haloOwnerProcs[iEnt]; + if(ownerProc == self) continue; - recvIDCPU(localIndex) = numOwnersTot + iEnt; + recvIDCPU(cursor[ownerProc]) = numOwnersTot + iEnt; + cursor[ownerProc]++; + } - localIndex++; - } + for(int proc = 0; proc < numProcsTot; proc++){ + if(proc == self) continue; - assert(localIndex == recvCount); + assert(cursor[proc] == + cudaAwareCache.recvOffsets[proc] + + cudaAwareCache.recvCounts[proc]); } - - Kokkos::deep_copy( - cudaAwareCache.recvIDGPUPerProc[proc], - recvIDCPU); } - } - Kokkos::fence(); + Kokkos::deep_copy(cudaAwareCache.recvIDGPU.view(), recvIDCPU); + } cudaAwareCache.valid = true; @@ -725,17 +864,14 @@ class MPMesh{ timer.reset(); } - for(int proc = 0; proc < numProcsTot; proc++){ - if(proc == self) continue; - if(cudaAwareCache.sendCounts[proc] <= 0) continue; - - auto sendEntityGPU = cudaAwareCache.sendEntityGPUPerProc[proc]; - auto sendDataGPU = cudaAwareCache.sendDataGPUPerProc[proc]; - int sendCount = cudaAwareCache.sendCounts[proc]; + // ---- Pack: ONE kernel over all neighbors' send entities at once ---- + if(cudaAwareCache.totalSendCount > 0){ + auto sendEntityGPU = cudaAwareCache.sendEntityGPU.view(); + auto sendDataGPU = cudaAwareCache.sendDataGPU.view(); Kokkos::parallel_for( - "pack cached cuda-aware mpi send buffer per proc", - sendCount, + "pack cached cuda-aware mpi send buffer batched", + cudaAwareCache.totalSendCount, KOKKOS_LAMBDA(const int i){ int entity = sendEntityGPU(i); @@ -766,8 +902,13 @@ class MPMesh{ if(cudaAwareCache.recvCounts[proc] > 0){ MPI_Request reqData; + double* recvPtr = + cudaAwareCache.recvDataGPU.data() + + static_cast(cudaAwareCache.recvOffsets[proc]) * + numEntries; + mpiError = MPI_Irecv( - cudaAwareCache.recvDataGPUPerProc[proc].data(), + recvPtr, cudaAwareCache.recvCounts[proc] * numEntries, MPI_DOUBLE, proc, @@ -781,8 +922,13 @@ class MPMesh{ if(cudaAwareCache.sendCounts[proc] > 0){ MPI_Request reqData; + double* sendPtr = + cudaAwareCache.sendDataGPU.data() + + static_cast(cudaAwareCache.sendOffsets[proc]) * + numEntries; + mpiError = MPI_Isend( - cudaAwareCache.sendDataGPUPerProc[proc].data(), + sendPtr, cudaAwareCache.sendCounts[proc] * numEntries, MPI_DOUBLE, proc, @@ -818,7 +964,7 @@ class MPMesh{ if(self == 0){ std::cout - << "[CUDA_AWARE_MPI] Device-pointer MPI failed." + << "[CUDA_AWARE_MPI] Batched device-pointer MPI failed." << std::endl; } @@ -838,6 +984,8 @@ class MPMesh{ return; } + timer.reset(); + if(self == 0){ std::cout << "[CUDA_AWARE_MPI] Failure happened after MPI requests were posted. " @@ -856,18 +1004,15 @@ class MPMesh{ timer.reset(); - for(int proc = 0; proc < numProcsTot; proc++){ - if(proc == self) continue; - if(cudaAwareCache.recvCounts[proc] <= 0) continue; - - auto recvIDGPU = cudaAwareCache.recvIDGPUPerProc[proc]; - auto recvDataGPU = cudaAwareCache.recvDataGPUPerProc[proc]; - int recvCount = cudaAwareCache.recvCounts[proc]; + // ---- Unpack: ONE kernel over all neighbors' recv entities at once ---- + if(cudaAwareCache.totalRecvCount > 0){ + auto recvIDGPU = cudaAwareCache.recvIDGPU.view(); + auto recvDataGPU = cudaAwareCache.recvDataGPU.view(); if(op == 0){ Kokkos::parallel_for( - "halo add cached cuda-aware mpi per proc", - recvCount, + "halo add cached cuda-aware mpi batched", + cudaAwareCache.totalRecvCount, KOKKOS_LAMBDA(const int i){ const int vertex = recvIDGPU(i); @@ -885,8 +1030,8 @@ class MPMesh{ } else{ Kokkos::parallel_for( - "halo assign cached cuda-aware mpi per proc", - recvCount, + "halo assign cached cuda-aware mpi batched", + cudaAwareCache.totalRecvCount, KOKKOS_LAMBDA(const int i){ const int vertex = recvIDGPU(i); @@ -933,4 +1078,3 @@ class MPMesh{ }//namespace polyMPO end #endif - diff --git a/src/pmpo_MPMesh_assembly.hpp b/src/pmpo_MPMesh_assembly.hpp index d41aa56d..d933e7a7 100644 --- a/src/pmpo_MPMesh_assembly.hpp +++ b/src/pmpo_MPMesh_assembly.hpp @@ -130,6 +130,7 @@ void MPMesh::reconstruct_coeff_full(){ radius=p_mesh->getSphereRadius(); //Assemble matrix for each vertex + timer.reset(); auto assemble = PS_LAMBDA(const int& elm, const int& mp, const int& mask) { if(mask) { //if material point is 'active'/'enabled' int nVtxE = elm2VtxConn(elm,0); //number of vertices bounding the element @@ -152,7 +153,7 @@ void MPMesh::reconstruct_coeff_full(){ }; p_MPs->parallel_for(assemble, "assembly"); Kokkos::fence(); - pumipic::RecordTime("Assemble Matrix Per Process" + std::to_string(self), timer.seconds()); + pumipic::RecordTime("VR Assemble Matrix Per Process" + std::to_string(self), timer.seconds()); //Mode 0 is Gather: Halos Send to Owners //Mode 1 is Scatter: Owners Send to Halos //Op 0 is addition @@ -162,20 +163,26 @@ void MPMesh::reconstruct_coeff_full(){ int op = 0; if (numProcsTot >1){ communicate_and_take_halo_contributions1_improved(vtxMatrices, numVertices, numEntriesMatrix, mode, op); + pumipic::RecordTime("VR Matrix Gather Halo/MPI " + std::to_string(self), timer.seconds()); + + timer.reset(); mode=1; op=1; communicate_and_take_halo_contributions1_improved(vtxMatrices, numVertices, numEntriesMatrix, mode, op); + pumipic::RecordTime("VR Matrix Scatter Halo/MPI " + std::to_string(self), timer.seconds()); } - pumipic::RecordTime("Communicate Matrix Values" + std::to_string(self), timer.seconds()); - - //Stroe the 1st matrix element + + //Store the 1st matrix element Kokkos::ViewvtxMatrixMass_l("vtxMass", numVertices); Kokkos::parallel_for("storeMatrixMass", numVertices, KOKKOS_LAMBDA(const int vtx){ vtxMatrixMass_l(vtx) = vtxMatrices(vtx, 0); }); + Kokkos::fence(); this->vtxMatrixMass = vtxMatrixMass_l; + timer.reset(); invertMatrix(vtxMatrices, radius); + pumipic::RecordTime("VR Invert Matrix " + std::to_string(self), timer.seconds()); } void MPMesh::invertMatrix(const Kokkos::View& vtxMatrices, const double& radius){