diff --git a/test/SendRecvTests.cpp b/test/SendRecvTests.cpp index b2cd892e10..4f38470cc3 100644 --- a/test/SendRecvTests.cpp +++ b/test/SendRecvTests.cpp @@ -82,4 +82,83 @@ namespace RcclUnitTesting } testBed.Finalize(); } + + TEST(SendRecv, UserBufferRegister) + { + setenv("RCCL_ENABLE_INTRANET", "1", 1); + TestBed testBed; + + // Configuration + std::vector const& dataTypes = {ncclInt32, ncclFloat16, ncclFloat64}; + std::vector const numElements = {1048576, 53327, 1024, 0}; + bool const inPlace = false; + bool const useManagedMem = false; + bool const userRegistered = true; + int const groupCallId = 0; + + OptionalColArgs options; + bool isCorrect = true; + int numGpus = testBed.ev.maxGpus; + for (int rpg=0; rpg < 2 && isCorrect; ++rpg) + for (int isMultiProcess = 0; isMultiProcess <= 1 && isCorrect; ++isMultiProcess) + { + if (!(testBed.ev.processMask & (1 << isMultiProcess))) continue; + int ranksPerGpu = rpg == 0 ? 1 : testBed.ev.maxRanksPerGpu; + int totalRanks = numGpus * ranksPerGpu; + int const numProcesses = isMultiProcess ? numGpus : 1; + testBed.InitComms(TestBed::GetDeviceIdsList(numProcesses, numGpus, ranksPerGpu), 1); + + for (int dataIdx = 0; dataIdx < dataTypes.size() && isCorrect; ++dataIdx) + for (int numIdx = 0; numIdx < numElements.size() && isCorrect; ++numIdx) + for (int sendRank = 0; sendRank < totalRanks; ++sendRank) + { + for (int recvRank = 0; recvRank < totalRanks; ++recvRank) + { + options.root = recvRank; + testBed.SetCollectiveArgs(ncclCollSend, + dataTypes[dataIdx], + numElements[numIdx], + numElements[numIdx], + options, + 0, + groupCallId, + sendRank); + if (recvRank == 0) + { + testBed.AllocateMem(inPlace, useManagedMem, groupCallId, 0, sendRank, userRegistered); + testBed.PrepareData(groupCallId, 0, sendRank); + } + if (recvRank != sendRank) + { + if (testBed.ev.showNames) // Show test names + INFO("%s Datatype: %s SendReceive test Rank %d -> Rank %d for %d Elements\n", + isMultiProcess ? "MP" : "SP", + ncclDataTypeNames[dataTypes[dataIdx]], + sendRank, + recvRank, + numElements[numIdx]); + + options.root = sendRank; + testBed.SetCollectiveArgs(ncclCollRecv, + dataTypes[dataIdx], + numElements[numIdx], + numElements[numIdx], + options, + 0, + groupCallId, + recvRank); + testBed.AllocateMem(inPlace, useManagedMem, groupCallId, 0, recvRank, userRegistered); + testBed.PrepareData(groupCallId, 0, recvRank); + testBed.ExecuteCollectives({sendRank, recvRank}); + testBed.ValidateResults(isCorrect, groupCallId, 0, recvRank); + testBed.DeallocateMem(groupCallId, 0, recvRank); + } + } + testBed.DeallocateMem(groupCallId, 0, sendRank); + } + testBed.DestroyComms(); + } + testBed.Finalize(); + unsetenv("RCCL_ENABLE_INTRANET"); + } } diff --git a/test/common/CollectiveArgs.cpp b/test/common/CollectiveArgs.cpp index 6908f45e95..2c97a76e68 100644 --- a/test/common/CollectiveArgs.cpp +++ b/test/common/CollectiveArgs.cpp @@ -58,7 +58,8 @@ namespace RcclUnitTesting } ErrCode CollectiveArgs::AllocateMem(bool const inPlace, - bool const useManagedMem) + bool const useManagedMem, + bool const userRegistered) { this->numInputBytesAllocated = this->numInputElements * DataTypeToBytes(this->dataType); this->numOutputBytesAllocated = this->numOutputElements * DataTypeToBytes(this->dataType); @@ -66,6 +67,7 @@ namespace RcclUnitTesting this->numOutputElementsAllocated = this->numOutputElements; this->inPlace = inPlace; this->useManagedMem = useManagedMem; + this->userRegistered = userRegistered; if (hipSetDevice(this->deviceId) != hipSuccess) { @@ -77,26 +79,26 @@ namespace RcclUnitTesting { if (this->funcType == ncclCollScatter) { - CHECK_CALL(this->inputGpu.AllocateGpuMem(this->numInputBytesAllocated, useManagedMem)); + CHECK_CALL(this->inputGpu.AllocateGpuMem(this->numInputBytesAllocated, useManagedMem, userRegistered)); this->outputGpu.Attach(this->inputGpu.U1 + (this->globalRank * this->numOutputBytesAllocated)); } else if (this->funcType == ncclCollGather) { - CHECK_CALL(this->outputGpu.AllocateGpuMem(this->numOutputBytesAllocated, useManagedMem)); + CHECK_CALL(this->outputGpu.AllocateGpuMem(this->numOutputBytesAllocated, useManagedMem, userRegistered)); this->inputGpu.Attach(this->outputGpu.U1 + (this->globalRank * this->numInputBytesAllocated)); } else { size_t const numBytes = std::max(this->numInputBytesAllocated, this->numOutputBytesAllocated); - CHECK_CALL(this->inputGpu.AllocateGpuMem(numBytes, useManagedMem)); + CHECK_CALL(this->inputGpu.AllocateGpuMem(numBytes, useManagedMem, userRegistered)); this->outputGpu.Attach(this->inputGpu.ptr); } CHECK_CALL(this->expected.AllocateCpuMem(this->numOutputBytesAllocated)); } else { - CHECK_CALL(this->inputGpu.AllocateGpuMem(this->numInputBytesAllocated, useManagedMem)); - CHECK_CALL(this->outputGpu.AllocateGpuMem(this->numOutputBytesAllocated, useManagedMem)); + CHECK_CALL(this->inputGpu.AllocateGpuMem(this->numInputBytesAllocated, useManagedMem, userRegistered)); + CHECK_CALL(this->outputGpu.AllocateGpuMem(this->numOutputBytesAllocated, useManagedMem, userRegistered)); CHECK_CALL(this->expected.AllocateCpuMem(this->numOutputBytesAllocated)); } CHECK_CALL(this->outputCpu.AllocateCpuMem(this->numOutputBytesAllocated)); @@ -136,12 +138,12 @@ namespace RcclUnitTesting if (this->funcType == ncclCollGather) this->outputGpu.FreeGpuMem(); else - this->inputGpu.FreeGpuMem(); + this->inputGpu.FreeGpuMem(this->userRegistered); } else { - this->inputGpu.FreeGpuMem(); - this->outputGpu.FreeGpuMem(); + this->inputGpu.FreeGpuMem(this->userRegistered); + this->outputGpu.FreeGpuMem(this->userRegistered); } this->outputCpu.FreeCpuMem(); diff --git a/test/common/CollectiveArgs.hpp b/test/common/CollectiveArgs.hpp index 61347d44ee..182f76325a 100644 --- a/test/common/CollectiveArgs.hpp +++ b/test/common/CollectiveArgs.hpp @@ -115,6 +115,8 @@ namespace RcclUnitTesting PtrUnion expected; bool inPlace; bool useManagedMem; + bool userRegistered; + void* commRegHandle; size_t numInputBytesAllocated; size_t numOutputBytesAllocated; size_t numInputElementsAllocated; @@ -134,7 +136,8 @@ namespace RcclUnitTesting // Allocates GPU memory for input/output and CPU memory for expected // When inPlace is true, input and output share the same memory ErrCode AllocateMem(bool const inPlace, - bool const useManagedMem); + bool const useManagedMem, + bool const userRegistered); // Execute the provided data preparation function to fill input and compute expected results ErrCode PrepareData(CollFuncPtr const prepareDataFunc); diff --git a/test/common/PtrUnion.cpp b/test/common/PtrUnion.cpp index 7d67a5cdc0..cb0887164b 100644 --- a/test/common/PtrUnion.cpp +++ b/test/common/PtrUnion.cpp @@ -42,26 +42,38 @@ namespace RcclUnitTesting return TEST_SUCCESS; } - ErrCode PtrUnion::AllocateGpuMem(size_t const numBytes, bool const useManagedMem) + ErrCode PtrUnion::AllocateGpuMem(size_t const numBytes, bool const useManagedMem, bool const userRegistered) { if (numBytes) { - if (useManagedMem) + if (userRegistered) { - if (hipMallocManaged(&I1, numBytes) != hipSuccess) + if (ncclMemAlloc((void**)&I1, numBytes) != ncclSuccess) { - ERROR("Unable to allocate managed memory of GPU memory (%lu bytes)\n", numBytes); + ERROR("Unable to allocate user managed GPU memory (%lu bytes)\n", numBytes); return TEST_FAIL; } } else { - if (hipMalloc(&I1, numBytes) != hipSuccess) + if (useManagedMem) { - ERROR("Unable to allocate memory of GPU memory (%lu bytes)\n", numBytes); - return TEST_FAIL; + if (hipMallocManaged(&I1, numBytes) != hipSuccess) + { + ERROR("Unable to allocate managed memory of GPU memory (%lu bytes)\n", numBytes); + return TEST_FAIL; + } + } + else + { + if (hipMalloc(&I1, numBytes) != hipSuccess) + { + ERROR("Unable to allocate memory of GPU memory (%lu bytes)\n", numBytes); + return TEST_FAIL; + } } } + } return TEST_SUCCESS; } @@ -80,11 +92,14 @@ namespace RcclUnitTesting return TEST_SUCCESS; } - ErrCode PtrUnion::FreeGpuMem() + ErrCode PtrUnion::FreeGpuMem(bool const userRegistered) { if (this->ptr != nullptr) { - hipFree(this->ptr); + if (userRegistered) + ncclMemFree(this->ptr); + else + hipFree(this->ptr); this->ptr = nullptr; } return TEST_SUCCESS; diff --git a/test/common/PtrUnion.hpp b/test/common/PtrUnion.hpp index bed042f176..29d78a376a 100644 --- a/test/common/PtrUnion.hpp +++ b/test/common/PtrUnion.hpp @@ -55,10 +55,10 @@ namespace RcclUnitTesting ErrCode Attach(void *ptr); ErrCode Attach(PtrUnion ptrUnion); - ErrCode AllocateGpuMem(size_t const numBytes, bool const useManagedMem = false); + ErrCode AllocateGpuMem(size_t const numBytes, bool const useManagedMem = false, bool const userRegistered = false); ErrCode AllocateCpuMem(size_t const numBytes); - ErrCode FreeGpuMem(); + ErrCode FreeGpuMem(bool const userRegistered = false); ErrCode FreeCpuMem(); ErrCode ClearGpuMem(size_t const numBytes); diff --git a/test/common/TestBed.cpp b/test/common/TestBed.cpp index 84b4b74ac5..1421ff5c94 100644 --- a/test/common/TestBed.cpp +++ b/test/common/TestBed.cpp @@ -242,7 +242,8 @@ namespace RcclUnitTesting bool const useManagedMem, int const groupId, int const collId, - int const rank) + int const rank, + bool const userRegistered) { InteractiveWait("Starting AllocateMem"); @@ -267,6 +268,7 @@ namespace RcclUnitTesting PIPE_WRITE(childId, collId); PIPE_WRITE(childId, inPlace); PIPE_WRITE(childId, useManagedMem); + PIPE_WRITE(childId, userRegistered); PIPE_WRITE(childId, currGroup); PIPE_CHECK(childId); } diff --git a/test/common/TestBed.hpp b/test/common/TestBed.hpp index e5e27e4b81..d74d10c048 100644 --- a/test/common/TestBed.hpp +++ b/test/common/TestBed.hpp @@ -77,7 +77,8 @@ namespace RcclUnitTesting bool const useManagedMem = false, int const groupId = -1, int const collId = -1, - int const rank = -1); + int const rank = -1, + bool const userRegistered = false); // Initialize input and compute expected results // - requires that SetCollectiveArgs and AllocateMemory have already been called diff --git a/test/common/TestBedChild.cpp b/test/common/TestBedChild.cpp index 892585efe7..48554537ac 100644 --- a/test/common/TestBedChild.cpp +++ b/test/common/TestBedChild.cpp @@ -371,12 +371,14 @@ namespace RcclUnitTesting int collId; bool inPlace; bool useManagedMem; + bool userRegistered; int groupId; PIPE_READ(globalRank); PIPE_READ(collId); PIPE_READ(inPlace); PIPE_READ(useManagedMem); + PIPE_READ(userRegistered); PIPE_READ(groupId); if (globalRank < this->rankOffset || (this->rankOffset + comms.size() <= globalRank)) @@ -392,11 +394,12 @@ namespace RcclUnitTesting if (collId == -1 || collId == collIdx) { CollectiveArgs& collArg = this->collArgs[groupId][localRank][collIdx]; - CHECK_CALL(collArg.AllocateMem(inPlace, useManagedMem)); - if (this->verbose) INFO("Rank %d on child %d allocates memory for collective %d in group %d on device %d (%s,%s) Input: %p Output %p\n", + CHECK_CALL(collArg.AllocateMem(inPlace, useManagedMem, userRegistered)); + if (this->verbose) INFO("Rank %d on child %d allocates memory for collective %d in group %d on device %d (%s,%s,%s) Input: %p Output %p\n", globalRank, this->childId, collIdx, groupId, this->deviceIds[localRank], inPlace ? "in-place" : "out-of-place", useManagedMem ? "managed" : "unmanaged", + userRegistered ? "user registered buffer" : "internal copy", collArg.inputGpu.ptr, collArg.outputGpu.ptr); } @@ -520,7 +523,7 @@ namespace RcclUnitTesting CHECK_HIP_RANK(errCode, hipSetDevice(this->deviceIds[localRank])); - CollectiveArgs const& collArg = this->collArgs[groupId][localRank][collId]; + CollectiveArgs& collArg = this->collArgs[groupId][localRank][collId]; if (this->printValues && !useHipGraph) { @@ -643,6 +646,8 @@ namespace RcclUnitTesting "ncclAllToAllv"); break; case ncclCollSend: + if (collArg.userRegistered) + CHILD_NCCL_CALL_RANK(errCode, ncclCommRegister(this->comms[localRank], collArg.inputGpu.ptr, collArg.numInputBytesAllocated, &(collArg.commRegHandle)),"ncclCommRegister"); CHILD_NCCL_CALL_RANK(errCode, ncclSend( collArg.inputGpu.ptr, collArg.numInputElements, @@ -653,6 +658,8 @@ namespace RcclUnitTesting "ncclSend"); break; case ncclCollRecv: + if (collArg.userRegistered) + CHILD_NCCL_CALL_RANK(errCode, ncclCommRegister(this->comms[localRank], collArg.outputGpu.ptr, collArg.numOutputBytesAllocated, &(collArg.commRegHandle)), "ncclCommRegister"); CHILD_NCCL_CALL_RANK(errCode, ncclRecv( collArg.outputGpu.ptr, collArg.numOutputElements, @@ -884,6 +891,8 @@ namespace RcclUnitTesting for (int collIdx = 0; collIdx < collArgs[groupId][localRank].size(); ++collIdx) { CollectiveArgs& collArg = this->collArgs[groupId][localRank][collIdx]; + if (collArg.userRegistered && (collArg.funcType == ncclCollSend || collArg.funcType == ncclCollRecv)) + CHILD_NCCL_CALL(ncclCommDeregister(this->comms[localRank], collArg.commRegHandle), "ncclCommDeregister"); if (collId == -1 || collId == collIdx) { if (this->verbose)