Added new unit tests for AllReduce with Bias API (#2036)

* Added new unit tests for AllReduce with Bias API

* Address review comments

[ROCm/rccl commit: 7c12b0b76b]
This commit is contained in:
Atul Kulkarni
2025-12-03 17:37:34 -06:00
zatwierdzone przez GitHub
rodzic 3e650467fa
commit e4aef19511
9 zmienionych plików z 586 dodań i 18 usunięć
@@ -102,6 +102,17 @@ namespace RcclUnitTesting
CHECK_CALL(this->expected.AllocateCpuMem(this->numOutputBytesAllocated));
}
CHECK_CALL(this->outputCpu.AllocateCpuMem(this->numOutputBytesAllocated));
// Allocate bias buffers if bias is enabled
if (this->options.useBias)
{
this->numBiasElements = this->options.biasNumElements;
this->numBiasBytesAllocated = this->numBiasElements * DataTypeToBytes(this->dataType);
CHECK_CALL(this->biasGpu.AllocateGpuMem(this->numBiasBytesAllocated, useManagedMem, userRegistered));
CHECK_CALL(this->biasCpu.AllocateCpuMem(this->numBiasBytesAllocated));
this->biasRegHandle = nullptr;
}
return TEST_SUCCESS;
}
@@ -155,6 +166,15 @@ namespace RcclUnitTesting
if (this->options.scalarMode == 1) CHECK_HIP(hipHostFree(this->localScalar.ptr));
this->localScalar.Attach(nullptr);
}
// Deallocate bias buffers if they were allocated
if (this->options.useBias && this->numBiasBytesAllocated > 0)
{
this->biasGpu.FreeGpuMem(this->userRegistered);
this->biasCpu.FreeCpuMem();
this->biasRegHandle = nullptr;
}
return TEST_SUCCESS;
}