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

Adding User Buffer Registration support for Unit test (#1199)

Browse Source 
* Adding UBR support for UT SendRecv

Signed-off-by: Tim Hu <timhu102@amd.com>

* Update test/common/TestBedChild.cpp

Co-authored-by: corey-derochie-amd <161367113+corey-derochie-amd@users.noreply.github.com>

---------

Signed-off-by: Tim Hu <timhu102@amd.com>
Co-authored-by: corey-derochie-amd <161367113+corey-derochie-amd@users.noreply.github.com>
This commit is contained in:
Tim
2024-07-30 13:39:25 -04:00
committed by GitHub
parent ca5341d419
commit a4793286c7
8 changed files with 137 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/SendRecvTests.cpp
+79
View File
@@ -82,4 +82,83 @@ namespace RcclUnitTesting
}
testBed.Finalize();
}
TEST(SendRecv, UserBufferRegister)
{
setenv("RCCL_ENABLE_INTRANET", "1", 1);
TestBed testBed;
// Configuration
std::vector<ncclDataType_t> const& dataTypes = {ncclInt32, ncclFloat16, ncclFloat64};
std::vector<int> 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");
}
}
test/common/CollectiveArgs.cpp
+11 -9
View File
@@ -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();
test/common/CollectiveArgs.hpp
+4 -1
View File
@@ -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);
test/common/PtrUnion.cpp
+24 -9
View File
@@ -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;
test/common/PtrUnion.hpp
+2 -2
View File
@@ -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);
test/common/TestBed.cpp
+3 -1
View File
@@ -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);
}
test/common/TestBed.hpp
+2 -1
View File
@@ -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
test/common/TestBedChild.cpp
+12 -3
View File
@@ -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)