Adding User Buffer Registration support for Unit test (#1199)
* 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>
[ROCm/rccl commit: a4793286c7]
This commit is contained in:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user