AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use a map to host scratch buffers (#1004)

Browse Source 
* Use a map to host scratch buffers

* Address review feedbacks. Deliberately keep mscclSetupScratch function.
This commit is contained in:
Wen-Heng (Jack) Chung
2023-12-05 13:15:28 -06:00
committed by GitHub
parent bc44e3faa7
commit 293f0fb752
3 changed files with 16 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/misc/msccl/msccl_lifecycle.cc
+4 -5
View File
@@ -236,8 +236,6 @@ ncclResult_t mscclInit(ncclComm_t comm) {
return ncclSuccess;
}
status.scratchBuffer = nullptr;
status.scratchBufferSize = 0;
status.workIndex = 1;
NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
status.lastStream = nullptr;
@@ -542,13 +540,14 @@ ncclResult_t mscclTeardown() {
for (auto &p : status.devAlgos) {
CUDACHECK(hipFree(p.second));
}
CUDACHECK(hipFree(status.scratchBuffer));
CUDACHECK(hipFree(status.syncFlags));
status.hostAlgos.clear();
status.devAlgos.clear();
status.freeAlgoHandles.clear();
status.scratchBuffer = nullptr;
status.scratchBufferSize = 0;
for (auto &p : status.scratchBuffers) {
CUDACHECK(hipFree(p.second));
}
status.scratchBuffers.clear();
status.connectedAlgos.clear();
if (status.mscclSchedulerPtr) {
NCCLCHECK(status.mscclSchedulerPtr->teardown());