derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
2403 lines
93 KiB
C++
2403 lines
93 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
/**
|
|
* @file RegistrationMPITests.cpp
|
|
* @brief Tests for buffer registration functionality in RCCL
|
|
*
|
|
* This file contains tests for:
|
|
* 1. User Buffer Registration (UBR) - explicit buffer registration via ncclCommRegister
|
|
* 2. Graph Capture Registration - automatic buffer registration during HIP graph capture
|
|
*
|
|
* REQUIRED Environment Variables:
|
|
* NCCL_DEBUG=INFO Enable debug logging
|
|
* NCCL_DEBUG_SUBSYS=REG Enable REG subsystem logging
|
|
* NCCL_LOCAL_REGISTER=1 Enable local buffer registration (UBR tests)
|
|
* NCCL_GRAPH_REGISTER=1 Enable graph buffer registration (Graph tests)
|
|
* RCCL_MPI_LOG_ALL_RANKS=1 Enable per-rank logging for log verification
|
|
*
|
|
* Run examples:
|
|
* mpirun -np 8 ./rccl-UnitTestsMPI --gtest_filter=UBR_*
|
|
* mpirun -np 8 ./rccl-UnitTestsMPI --gtest_filter=GraphCapture_*
|
|
*/
|
|
|
|
#include "DeviceBufferHelpers.hpp"
|
|
#include "MPITestBase.hpp"
|
|
#include "MPIHelpers.hpp"
|
|
#include "ResourceGuards.hpp"
|
|
#include "TestChecks.hpp"
|
|
#include <cstdlib>
|
|
#include <regex>
|
|
#include <sstream>
|
|
|
|
#ifdef MPI_TESTS_ENABLED
|
|
|
|
using namespace MPITestConstants;
|
|
using namespace RCCLTestGuards;
|
|
using namespace RCCLTestHelpers;
|
|
|
|
// Test Configuration
|
|
namespace RegTestConfig {
|
|
constexpr size_t SMALL_COUNT = 1024; // 4KB for float
|
|
constexpr size_t MEDIUM_COUNT = 256 * 1024; // 1MB for float
|
|
constexpr size_t LARGE_COUNT = 1024 * 1024; // 4MB for float
|
|
|
|
using DefaultType = hip_bfloat16;
|
|
|
|
constexpr int MIN_RANKS_DEFAULT = 2;
|
|
constexpr int MIN_RANKS_ALLTOALL = 4;
|
|
constexpr int MIN_NODES_MULTINODE = 2;
|
|
}
|
|
|
|
// REG Log Checker - Pattern checking for registration debug output
|
|
class REGLogChecker
|
|
{
|
|
public:
|
|
explicit REGLogChecker(const std::string& logContent)
|
|
: m_content(logContent) {}
|
|
|
|
bool hasIPCRegistration() const
|
|
{
|
|
return hasPattern("IPC register buffer") ||
|
|
hasPattern("IPC registering buffer") ||
|
|
(hasPattern("Proxy rank") && hasPattern("register success"));
|
|
}
|
|
|
|
bool hasIPCReuse() const
|
|
{
|
|
return hasPattern("IPC reuse buffer");
|
|
}
|
|
|
|
bool hasNETReuse() const
|
|
{
|
|
return hasPattern("NET reuse buffer");
|
|
}
|
|
|
|
bool hasNumSegments(int n) const
|
|
{
|
|
const std::regex re("numSegments\\s*:?\\s*" + std::to_string(n) + "\\b");
|
|
return std::regex_search(m_content, re);
|
|
}
|
|
|
|
bool hasNETRegistration() const
|
|
{
|
|
return hasPattern("NET register userbuff");
|
|
}
|
|
|
|
bool hasAnyRegistrationSuccess() const
|
|
{
|
|
return hasIPCRegistration() || hasIPCReuse() || hasNETRegistration() || hasNETReuse();
|
|
}
|
|
|
|
bool hasIPCFailure() const
|
|
{
|
|
return hasPattern("failed to IPC register") ||
|
|
hasPattern("legacy IPC blocked");
|
|
}
|
|
|
|
bool hasNETFailure() const
|
|
{
|
|
return hasPattern("failed to NET register");
|
|
}
|
|
|
|
std::string getSummary() const
|
|
{
|
|
std::ostringstream ss;
|
|
ss << "REG Log: ";
|
|
if (hasIPCReuse()) ss << "[IPC-REUSE] ";
|
|
if (hasNETReuse()) ss << "[NET-REUSE] ";
|
|
if (hasIPCRegistration()) ss << "[IPC-REG] ";
|
|
if (hasNETRegistration()) ss << "[NET-REG] ";
|
|
if (hasIPCFailure()) ss << "[IPC-FAIL] ";
|
|
if (hasNETFailure()) ss << "[NET-FAIL] ";
|
|
if (!hasAnyRegistrationSuccess() && !hasIPCFailure() && !hasNETFailure()) ss << "[NO-REG]";
|
|
return ss.str();
|
|
}
|
|
|
|
size_t getContentLength() const { return m_content.size(); }
|
|
|
|
private:
|
|
bool hasPattern(const std::string& pattern) const
|
|
{
|
|
return m_content.find(pattern) != std::string::npos;
|
|
}
|
|
|
|
std::string m_content;
|
|
};
|
|
|
|
// Registration Test Base Class
|
|
class RegistrationTestBase : public MPITestBase
|
|
{
|
|
protected:
|
|
struct RegInfo {
|
|
void* buffer = nullptr;
|
|
void* handle = nullptr;
|
|
size_t size = 0;
|
|
bool registered = false;
|
|
};
|
|
|
|
// Buffer Management
|
|
RegInfo allocateAndRegister(size_t size)
|
|
{
|
|
RegInfo info;
|
|
info.size = size;
|
|
|
|
if (hipMalloc(&info.buffer, size) != hipSuccess) {
|
|
return info;
|
|
}
|
|
|
|
ncclResult_t result = ncclCommRegister(getActiveCommunicator(),
|
|
info.buffer, size, &info.handle);
|
|
info.registered = (result == ncclSuccess && info.handle != nullptr);
|
|
|
|
return info;
|
|
}
|
|
|
|
void cleanupRegInfo(RegInfo& info)
|
|
{
|
|
if (info.handle) {
|
|
ncclCommDeregister(getActiveCommunicator(), info.handle);
|
|
info.handle = nullptr;
|
|
}
|
|
if (info.buffer) {
|
|
(void)hipFree(info.buffer);
|
|
info.buffer = nullptr;
|
|
}
|
|
info.registered = false;
|
|
}
|
|
|
|
// Test Setup
|
|
bool setupMultiNode(int minRanks = 2, int minNodes = 2)
|
|
{
|
|
int nodeCount = MPITestConstants::detectNodeCount();
|
|
if (nodeCount < minNodes) {
|
|
return false;
|
|
}
|
|
if (!validateTestPrerequisites(minRanks, kNoProcessLimit,
|
|
kNoPowerOfTwoRequired, minNodes, kNoNodeLimit)) {
|
|
return false;
|
|
}
|
|
return (createTestCommunicator() == ncclSuccess);
|
|
}
|
|
|
|
// Environment Checks
|
|
bool isUBREnabled()
|
|
{
|
|
const char* localReg = getenv("NCCL_LOCAL_REGISTER");
|
|
return (localReg && std::string(localReg) == "1");
|
|
}
|
|
|
|
bool isGraphRegisterEnabled()
|
|
{
|
|
const char* graphReg = getenv("NCCL_GRAPH_REGISTER");
|
|
return (graphReg && std::string(graphReg) == "1");
|
|
}
|
|
|
|
bool isCuMemEnabled()
|
|
{
|
|
const char* cuMem = getenv("NCCL_CUMEM_ENABLE");
|
|
return (cuMem && std::string(cuMem) == "1");
|
|
}
|
|
|
|
bool isMultiSegmentRegisterEnabled()
|
|
{
|
|
const char* mseg = getenv("NCCL_MULTI_SEGMENT_REGISTER");
|
|
return (!mseg || std::string(mseg) != "0");
|
|
}
|
|
|
|
bool isWinEnabled()
|
|
{
|
|
const char* win = getenv("NCCL_WIN_ENABLE");
|
|
return (!win || std::string(win) != "0");
|
|
}
|
|
|
|
bool isGinEnabled()
|
|
{
|
|
const char* en = getenv("NCCL_GIN_ENABLE");
|
|
if (en && std::string(en) == "0") return false;
|
|
const char* type = getenv("NCCL_GIN_TYPE");
|
|
return (type && std::string(type) == "2");
|
|
}
|
|
|
|
bool isElasticBufferRegisterEnabled()
|
|
{
|
|
const char* en = getenv("NCCL_ELASTIC_BUFFER_REGISTER");
|
|
return (!en || std::string(en) != "0");
|
|
}
|
|
|
|
bool isPerRankLoggingEnabled() { return MPIHelpers::isPerRankLoggingEnabled(); }
|
|
|
|
// Log File Access
|
|
std::string readRankLogFile()
|
|
{
|
|
return MPIHelpers::readRankLogFile(getTestMpiRank());
|
|
}
|
|
|
|
REGLogChecker getLogChecker()
|
|
{
|
|
return REGLogChecker(readRankLogFile());
|
|
}
|
|
|
|
// Data Initialization and Verification
|
|
template<typename T>
|
|
void initSendBuffer(void* buffer, size_t count, int rank)
|
|
{
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(buffer, count,
|
|
[rank](size_t) { return static_cast<T>(static_cast<float>(rank + 1)); }));
|
|
}
|
|
|
|
template<typename T>
|
|
bool verifyAllReduceResult(void* buffer, size_t count, int nRanks)
|
|
{
|
|
T expected = static_cast<T>(static_cast<float>(nRanks * (nRanks + 1) / 2));
|
|
return verifyBufferData<T>(buffer, count, [expected](size_t) { return expected; });
|
|
}
|
|
|
|
template<typename T>
|
|
bool verifyReduceScatterResult(void* buffer, size_t count, int nRanks)
|
|
{
|
|
T expected = static_cast<T>(static_cast<float>(nRanks * (nRanks + 1) / 2));
|
|
return verifyBufferData<T>(buffer, count, [expected](size_t) { return expected; });
|
|
}
|
|
|
|
template<typename T>
|
|
bool verifyAllGatherResult(void* buffer, size_t countPerRank, int nRanks)
|
|
{
|
|
return verifyBufferData<T>(buffer, countPerRank * nRanks,
|
|
[countPerRank](size_t i) {
|
|
int srcRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(srcRank + 1));
|
|
});
|
|
}
|
|
|
|
template<typename T>
|
|
bool verifyBroadcastResult(void* buffer, size_t count, T rootValue)
|
|
{
|
|
return verifyBufferData<T>(buffer, count,
|
|
[rootValue](size_t) { return rootValue; });
|
|
}
|
|
};
|
|
|
|
// ============================================================================
|
|
// User Buffer Registration (UBR) Tests
|
|
// ============================================================================
|
|
|
|
class UBR_AllReduce : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_AllReduce, OutOfPlace_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t count = RegTestConfig::MEDIUM_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
RegInfo sendInfo = allocateAndRegister(count * sizeof(T));
|
|
RegInfo recvInfo = allocateAndRegister(count * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
initSendBuffer<T>(sendInfo.buffer, count, rank);
|
|
|
|
ncclResult_t result = ncclAllReduce(sendInfo.buffer, recvInfo.buffer, count,
|
|
getNcclDataType<T>(), ncclSum,
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvInfo.buffer, count, nRanks));
|
|
}
|
|
|
|
class UBR_AllGather : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_AllGather, OutOfPlace_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t countPerRank = RegTestConfig::SMALL_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
RegInfo sendInfo = allocateAndRegister(countPerRank * sizeof(T));
|
|
RegInfo recvInfo = allocateAndRegister(countPerRank * nRanks * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
initSendBuffer<T>(sendInfo.buffer, countPerRank, rank);
|
|
|
|
ncclResult_t result = ncclAllGather(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(),
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
ASSERT_TRUE(verifyAllGatherResult<T>(recvInfo.buffer, countPerRank, nRanks));
|
|
}
|
|
|
|
class UBR_ReduceScatter : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_ReduceScatter, OutOfPlace_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t countPerRank = RegTestConfig::SMALL_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
RegInfo sendInfo = allocateAndRegister(countPerRank * nRanks * sizeof(T));
|
|
RegInfo recvInfo = allocateAndRegister(countPerRank * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
initSendBuffer<T>(sendInfo.buffer, countPerRank * nRanks, rank);
|
|
|
|
ncclResult_t result = ncclReduceScatter(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(), ncclSum,
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
ASSERT_TRUE(verifyReduceScatterResult<T>(recvInfo.buffer, countPerRank, nRanks));
|
|
}
|
|
|
|
class UBR_Broadcast : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_Broadcast, NonZeroRoot_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes with 4+ ranks";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t count = RegTestConfig::MEDIUM_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const int root = nRanks - 1;
|
|
|
|
RegInfo bufInfo = allocateAndRegister(count * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(bufInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(bufInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(bufInfo.handle, nullptr);
|
|
|
|
const T rootValue = static_cast<T>(99.0f);
|
|
if (rank == root) {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(bufInfo.buffer, count,
|
|
[rootValue](size_t) { return rootValue; }));
|
|
} else {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(bufInfo.buffer, count,
|
|
[](size_t) { return static_cast<T>(0.0f); }));
|
|
}
|
|
|
|
ncclResult_t result = ncclBroadcast(bufInfo.buffer, bufInfo.buffer, count,
|
|
getNcclDataType<T>(), root,
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
ASSERT_TRUE(verifyBroadcastResult<T>(bufInfo.buffer, count, rootValue));
|
|
}
|
|
|
|
class UBR_AllToAll : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_AllToAll, OutOfPlace_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes with 4+ ranks";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t countPerRank = RegTestConfig::SMALL_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const size_t totalCount = countPerRank * nRanks;
|
|
|
|
RegInfo sendInfo = allocateAndRegister(totalCount * sizeof(T));
|
|
RegInfo recvInfo = allocateAndRegister(totalCount * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendInfo.buffer, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int destRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(rank * 100 + destRank));
|
|
}));
|
|
|
|
ncclResult_t result = ncclAllToAll(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(),
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
bool verified = verifyBufferData<T>(recvInfo.buffer, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int srcRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(srcRank * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(verified);
|
|
}
|
|
|
|
class UBR_SendRecv : public RegistrationTestBase {};
|
|
|
|
TEST_F(UBR_SendRecv, RingPattern_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes for SendRecv UBR";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t count = RegTestConfig::SMALL_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
RegInfo sendInfo = allocateAndRegister(count * sizeof(T));
|
|
RegInfo recvInfo = allocateAndRegister(count * sizeof(T));
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
int sendPeer = (rank + 1) % nRanks;
|
|
int recvPeer = (rank - 1 + nRanks) % nRanks;
|
|
|
|
initSendBuffer<T>(sendInfo.buffer, count, rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
ncclResult_t sendResult = ncclSend(sendInfo.buffer, count, getNcclDataType<T>(),
|
|
sendPeer, getActiveCommunicator(), getActiveStream());
|
|
ncclResult_t recvResult = ncclRecv(recvInfo.buffer, count, getNcclDataType<T>(),
|
|
recvPeer, getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
|
|
ASSERT_EQ(ncclSuccess, sendResult);
|
|
ASSERT_EQ(ncclSuccess, recvResult);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
T expected = static_cast<T>(static_cast<float>(recvPeer + 1));
|
|
bool verified = verifyBufferData<T>(recvInfo.buffer, count,
|
|
[expected](size_t) { return expected; });
|
|
ASSERT_TRUE(verified);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Multi-Segment Registration Tests
|
|
// ============================================================================
|
|
|
|
/**
|
|
* @brief Tests for buffer registration when the user buffer spans multiple
|
|
* underlying physical allocations ("multi-segment" buffers).
|
|
*
|
|
* A multi-segment buffer is built by mapping N separate physical handles
|
|
* (hipMemCreate) contiguously into a single reserved virtual address range
|
|
* (hipMemAddressReserve + hipMemMap). cuMemGetAddressRange() on the head of
|
|
* such a buffer returns only the first segment, so the registration path
|
|
* detects the cross-boundary case and walks every segment when both
|
|
* ncclCuMemEnable() and NCCL_MULTI_SEGMENT_REGISTER (default 1) are true.
|
|
*
|
|
*/
|
|
class UBR_MultiSegment : public RegistrationTestBase
|
|
{
|
|
protected:
|
|
using T = RegTestConfig::DefaultType;
|
|
|
|
// Multi-segment registration currently relies on dmabuf support from the
|
|
// runtime/HSA layer for the inter-node (NET/GIN) path. Until that lands,
|
|
// restrict every UBR_MultiSegment test to a single node so the multi-node
|
|
// tests are not exercised.
|
|
void SetUp() override
|
|
{
|
|
RegistrationTestBase::SetUp();
|
|
if (::testing::Test::IsSkipped() || ::testing::Test::HasFatalFailure()) {
|
|
return;
|
|
}
|
|
if (MPITestConstants::detectNodeCount() != 1) {
|
|
GTEST_SKIP() << "UBR_MultiSegment is limited to single node until "
|
|
"dmabuf support is available from the HIP/HSA layer";
|
|
}
|
|
}
|
|
|
|
struct MultiSegmentBuffer
|
|
{
|
|
hipDeviceptr_t vaBase = 0;
|
|
size_t segmentSize = 0;
|
|
size_t totalSize = 0;
|
|
std::vector<hipMemGenericAllocationHandle_t> handles;
|
|
};
|
|
|
|
void createMultiSegmentBuffer(int dev,
|
|
size_t requestedSegmentSize,
|
|
int numSegments,
|
|
MultiSegmentBuffer& buf)
|
|
{
|
|
buf = MultiSegmentBuffer{};
|
|
|
|
hipMemAllocationProp prop = {};
|
|
prop.type = hipMemAllocationTypePinned;
|
|
prop.location.type = hipMemLocationTypeDevice;
|
|
prop.location.id = dev;
|
|
prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
|
|
|
|
size_t granularity = 0;
|
|
HIP_CHECK(hipMemGetAllocationGranularity(&granularity, &prop, hipMemAllocationGranularityMinimum));
|
|
|
|
const size_t segSize = ((requestedSegmentSize + granularity - 1) / granularity) * granularity;
|
|
const size_t totalSize = segSize * numSegments;
|
|
|
|
hipDeviceptr_t vaBase = 0;
|
|
HIP_CHECK(hipMemAddressReserve(&vaBase, totalSize, granularity, 0, 0));
|
|
|
|
std::vector<hipMemGenericAllocationHandle_t> handles(numSegments, 0);
|
|
char* vaBaseBytes = static_cast<char*>(vaBase);
|
|
for (int i = 0; i < numSegments; i++) {
|
|
HIP_CHECK(hipMemCreate(&handles[i], segSize, &prop, 0));
|
|
HIP_CHECK(hipMemMap(vaBaseBytes + i * segSize, segSize, 0, handles[i], 0));
|
|
}
|
|
|
|
hipMemAccessDesc accessDesc = {};
|
|
accessDesc.location.type = hipMemLocationTypeDevice;
|
|
accessDesc.location.id = dev;
|
|
accessDesc.flags = hipMemAccessFlagsProtReadWrite;
|
|
HIP_CHECK(hipMemSetAccess(vaBase, totalSize, &accessDesc, 1));
|
|
|
|
buf.vaBase = vaBase;
|
|
buf.segmentSize = segSize;
|
|
buf.totalSize = totalSize;
|
|
buf.handles = std::move(handles);
|
|
}
|
|
|
|
void releaseMultiSegmentBuffer(MultiSegmentBuffer& buf)
|
|
{
|
|
if (buf.totalSize == 0) return;
|
|
HIP_EXPECT(hipMemUnmap(buf.vaBase, buf.totalSize));
|
|
for (auto h : buf.handles) {
|
|
if (h != 0) HIP_EXPECT(hipMemRelease(h));
|
|
}
|
|
HIP_EXPECT(hipMemAddressFree(buf.vaBase, buf.totalSize));
|
|
buf.vaBase = 0;
|
|
buf.segmentSize = 0;
|
|
buf.totalSize = 0;
|
|
buf.handles.clear();
|
|
}
|
|
|
|
/**
|
|
* @brief Like createMultiSegmentBuffer, but backs the trailing
|
|
* `numHostSegments` segments with host memory
|
|
* (hipMemLocationTypeHost) rather than device memory, producing a
|
|
* mixed device/host "elastic" buffer.
|
|
*
|
|
* HIP/CLR has no host-NUMA VMM type and only accepts a plain Host location
|
|
* (id must be 0). Host VMM is only available on recent runtimes (ROCm >=
|
|
* 7.12); if it is unsupported, buf.totalSize is left 0 so the caller can
|
|
* GTEST_SKIP() instead of failing.
|
|
*/
|
|
void createMixedMultiSegmentBuffer(int dev,
|
|
size_t requestedSegmentSize,
|
|
int numSegments,
|
|
int numHostSegments,
|
|
MultiSegmentBuffer& buf)
|
|
{
|
|
buf = MultiSegmentBuffer{};
|
|
ASSERT_GE(numSegments, 1);
|
|
ASSERT_GE(numHostSegments, 0);
|
|
ASSERT_LE(numHostSegments, numSegments);
|
|
|
|
hipMemAllocationProp devProp = {};
|
|
devProp.type = hipMemAllocationTypePinned;
|
|
devProp.location.type = hipMemLocationTypeDevice;
|
|
devProp.location.id = dev;
|
|
devProp.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
|
|
|
|
hipMemAllocationProp hostProp = {};
|
|
hostProp.type = hipMemAllocationTypePinned;
|
|
hostProp.location.type = hipMemLocationTypeHost;
|
|
hostProp.location.id = 0;
|
|
hostProp.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
|
|
|
|
size_t devGran = 0;
|
|
size_t hostGran = 0;
|
|
HIP_CHECK(hipMemGetAllocationGranularity(&devGran, &devProp, hipMemAllocationGranularityMinimum));
|
|
if (numHostSegments > 0 &&
|
|
hipMemGetAllocationGranularity(&hostGran, &hostProp, hipMemAllocationGranularityMinimum) != hipSuccess) {
|
|
return;
|
|
}
|
|
size_t gran = devGran;
|
|
if (hostGran > gran) gran = hostGran;
|
|
|
|
const size_t segSize = ((requestedSegmentSize + gran - 1) / gran) * gran;
|
|
const size_t totalSize = segSize * numSegments;
|
|
const int numDevSegments = numSegments - numHostSegments;
|
|
|
|
hipDeviceptr_t vaBase = 0;
|
|
HIP_CHECK(hipMemAddressReserve(&vaBase, totalSize, gran, 0, 0));
|
|
char* vaBaseBytes = static_cast<char*>(vaBase);
|
|
|
|
std::vector<hipMemGenericAllocationHandle_t> handles(numSegments, 0);
|
|
bool hostUnsupported = false;
|
|
int mapped = 0;
|
|
for (int i = 0; i < numSegments; i++) {
|
|
const bool isHost = (i >= numDevSegments);
|
|
hipMemAllocationProp& prop = isHost ? hostProp : devProp;
|
|
hipError_t err = hipMemCreate(&handles[i], segSize, &prop, 0);
|
|
if (isHost && err != hipSuccess) {
|
|
hostUnsupported = true;
|
|
break;
|
|
}
|
|
HIP_CHECK(err);
|
|
HIP_CHECK(hipMemMap(vaBaseBytes + i * segSize, segSize, 0, handles[i], 0));
|
|
mapped++;
|
|
}
|
|
|
|
if (hostUnsupported) {
|
|
for (int i = 0; i < mapped; i++) {
|
|
HIP_EXPECT(hipMemUnmap(vaBaseBytes + i * segSize, segSize));
|
|
}
|
|
for (auto h : handles) {
|
|
if (h != 0) HIP_EXPECT(hipMemRelease(h));
|
|
}
|
|
HIP_EXPECT(hipMemAddressFree(vaBase, totalSize));
|
|
return;
|
|
}
|
|
|
|
hipMemAccessDesc accessDesc = {};
|
|
accessDesc.location.type = hipMemLocationTypeDevice;
|
|
accessDesc.location.id = dev;
|
|
accessDesc.flags = hipMemAccessFlagsProtReadWrite;
|
|
HIP_CHECK(hipMemSetAccess(vaBase, totalSize, &accessDesc, 1));
|
|
|
|
buf.vaBase = vaBase;
|
|
buf.segmentSize = segSize;
|
|
buf.totalSize = totalSize;
|
|
buf.handles = std::move(handles);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* @brief Out-of-place AllReduce on a buffer that spans multiple VMM segments.
|
|
*
|
|
* Exercises the multi-segment registration branch on Ring AllReduce.
|
|
*
|
|
* Layout (N = kSegmentsPerHalf, total 2 * N physical segments allocated):
|
|
* Total reserved VA = 2 * N * kSegmentSize
|
|
* sendbuff = [0, N * kSegmentSize) covers first N segments
|
|
* recvbuff = [N * kSegmentSize, 2N * kSegmentSize) covers last N segments
|
|
*
|
|
* Confirmation in the logs (NCCL_DEBUG=TRACE NCCL_DEBUG_SUBSYS=REG):
|
|
* "... numSegments 4"
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Generic)
|
|
{
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
ASSERT_TRUE(isCuMemEnabled()) << "NCCL_CUMEM_ENABLE must be set to 1";
|
|
ASSERT_TRUE(isMultiSegmentRegisterEnabled()) << "NCCL_MULTI_SEGMENT_REGISTER must be set to 1";
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
constexpr size_t kSegmentSize = 32 * 1024 * 1024;
|
|
constexpr int kNumSegments = 8;
|
|
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
SCOPED_TRACE("kNumSegments=" + std::to_string(kNumSegments));
|
|
|
|
MultiSegmentBuffer buf;
|
|
ASSERT_NO_FATAL_FAILURE(
|
|
createMultiSegmentBuffer(dev, kSegmentSize, kNumSegments, buf));
|
|
if (buf.totalSize == 0) {
|
|
GTEST_SKIP() << "Raw VMM (hipMemCreate / Reserve / Map) not supported on this runtime";
|
|
}
|
|
auto vmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(buf); });
|
|
|
|
const size_t halfSize = buf.totalSize / 2;
|
|
ASSERT_EQ(halfSize % sizeof(T), 0u);
|
|
char* base = reinterpret_cast<char*>(buf.vaBase);
|
|
void* sendBuf = base;
|
|
void* recvBuf = base + halfSize;
|
|
size_t count = halfSize / sizeof(T);
|
|
|
|
void* regHandle = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommRegister(getActiveCommunicator(), buf.vaBase, buf.totalSize, ®Handle));
|
|
auto regCleanup = makeScopeGuard([&]() {
|
|
if (regHandle) HIP_EXPECT(ncclCommDeregister(getActiveCommunicator(), regHandle));
|
|
});
|
|
ASSERT_MPI_NE(regHandle, nullptr);
|
|
|
|
initSendBuffer<T>(sendBuf, count, rank);
|
|
|
|
ncclResult_t result = ncclAllReduce(sendBuf, recvBuf, count, getNcclDataType<T>(), ncclSum, getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvBuf, count, nRanks));
|
|
|
|
REGLogChecker checker = getLogChecker();
|
|
TEST_INFO("SpansMultipleSegments: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
ASSERT_TRUE(checker.hasNumSegments(kNumSegments / 2))
|
|
<< "Expected 'numSegments " << (kNumSegments / 2)
|
|
<< "' in log - the multi-segment registration branch did not fire";
|
|
}
|
|
|
|
/**
|
|
* @brief Register once, AllReduce twice - exercises the registration reuse fast path.
|
|
*
|
|
* After the first collective on a registered multi-segment buffer, a subsequent
|
|
* collective on the same buffer must hit the cache entry and skip re-registration.
|
|
*
|
|
* The fast path is transport-dependent: intra-node ranks reuse via the P2P/IPC
|
|
* cache (p2p.cc, "IPC reuse buffer"), while the inter-node leg reuses the NIC
|
|
* MR via the NET cache (net.cc, "NET reuse buffer"). The topology decides which
|
|
* one fires, so this test asserts on "either IPC or NET" rather than assuming a
|
|
* transport:
|
|
* - the first call performs an initial registration (IPC or NET)
|
|
* - a subsequent call hits the cached entry (IPC or NET reuse)
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Generic_Reuse)
|
|
{
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
ASSERT_TRUE(isCuMemEnabled())
|
|
<< "NCCL_CUMEM_ENABLE must be set to 1 (gates the multi-segment IPC branch in p2p.cc:1071)";
|
|
ASSERT_TRUE(isMultiSegmentRegisterEnabled())
|
|
<< "NCCL_MULTI_SEGMENT_REGISTER must not be 0 (gates the multi-segment IPC branch in p2p.cc:1071)";
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
constexpr size_t kRequestedSegmentSize = 128 * 1024 * 1024;
|
|
constexpr int kNumSegments = 4;
|
|
|
|
MultiSegmentBuffer buf;
|
|
ASSERT_NO_FATAL_FAILURE(createMultiSegmentBuffer(dev, kRequestedSegmentSize, kNumSegments, buf));
|
|
if (buf.totalSize == 0) {
|
|
GTEST_SKIP() << "Raw VMM (hipMemCreate / Reserve / Map) not supported on this runtime";
|
|
}
|
|
auto vmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(buf); });
|
|
|
|
const size_t halfSize = buf.totalSize / 2;
|
|
ASSERT_EQ(halfSize % sizeof(T), 0u);
|
|
char* base = reinterpret_cast<char*>(buf.vaBase);
|
|
void* sendBuf = base;
|
|
void* recvBuf = base + halfSize;
|
|
size_t count = halfSize / sizeof(T);
|
|
|
|
void* regHandle = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommRegister(getActiveCommunicator(), buf.vaBase, buf.totalSize, ®Handle));
|
|
auto regCleanup = makeScopeGuard([&]() {
|
|
if (regHandle) HIP_EXPECT(ncclCommDeregister(getActiveCommunicator(), regHandle));
|
|
});
|
|
ASSERT_MPI_NE(regHandle, nullptr);
|
|
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
initSendBuffer<T>(sendBuf, count, rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclAllReduce(sendBuf, recvBuf, count,getNcclDataType<T>(), ncclSum, getActiveCommunicator(), getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvBuf, count, nRanks));
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemset(recvBuf, 0, halfSize));
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclAllReduce(sendBuf, recvBuf, count, getNcclDataType<T>(), ncclSum,getActiveCommunicator(), getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvBuf, count, nRanks));
|
|
|
|
REGLogChecker checker = getLogChecker();
|
|
TEST_INFO("RegisterOnceUseTwice: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
ASSERT_TRUE(checker.hasIPCRegistration() || checker.hasNETRegistration())
|
|
<< "Expected an initial registration log entry (IPC or NET) from the first AllReduce";
|
|
ASSERT_TRUE(checker.hasIPCReuse() || checker.hasNETReuse())
|
|
<< "Expected a reuse log entry (IPC or NET) from the second AllReduce - "
|
|
"the registration cache (p2p.cc / net.cc) was not hit";
|
|
}
|
|
|
|
/**
|
|
* @brief Multi-segment registration on the symmetric-window path.
|
|
*
|
|
* Validates GPU-only multi-segment registration for symmetric windows
|
|
* (ncclCommWindowRegister, NCCL_WIN_COLL_SYMMETRIC). The reserved VA spans
|
|
* several physical cuMem segments, so the LSA team maps each segment
|
|
* individually (dev_runtime.cc symMemoryMapLsaTeam) and logs, per segment:
|
|
* "[<lsaRank>] Segment <i>, Type : <t>, numSegments <N>, Segment size ..."
|
|
*
|
|
* Unlike the ncclCommRegister path, the multi-segment registration happens once
|
|
* at ncclCommWindowRegister time. The two AllReduce calls then reuse the same
|
|
* window, confirming it stays valid and correct across collectives.
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Symmetric_Lsa)
|
|
{
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isCuMemEnabled()) << "NCCL_CUMEM_ENABLE must be set to 1";
|
|
ASSERT_TRUE(isWinEnabled()) << "NCCL_WIN_ENABLE must not be set to 0";
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
constexpr size_t kSegmentSize = 32 * 1024 * 1024;
|
|
constexpr int kNumSegments = 4;
|
|
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
SCOPED_TRACE("kNumSegments=" + std::to_string(kNumSegments));
|
|
|
|
MultiSegmentBuffer buf;
|
|
ASSERT_NO_FATAL_FAILURE(createMultiSegmentBuffer(dev, kSegmentSize, kNumSegments, buf));
|
|
if (buf.totalSize == 0) {
|
|
GTEST_SKIP() << "Raw VMM (hipMemCreate / Reserve / Map) not supported on this runtime";
|
|
}
|
|
auto vmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(buf); });
|
|
|
|
const size_t halfSize = buf.totalSize / 2;
|
|
ASSERT_EQ(halfSize % sizeof(T), 0u);
|
|
char* base = reinterpret_cast<char*>(buf.vaBase);
|
|
void* sendBuf = base;
|
|
void* recvBuf = base + halfSize;
|
|
size_t count = halfSize / sizeof(T);
|
|
|
|
ncclWindow_t win = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommWindowRegister(getActiveCommunicator(), buf.vaBase, buf.totalSize, &win, NCCL_WIN_COLL_SYMMETRIC));
|
|
auto winCleanup = makeScopeGuard([&]() {
|
|
if (win) HIP_EXPECT(ncclCommWindowDeregister(getActiveCommunicator(), win));
|
|
});
|
|
ASSERT_MPI_NE(win, nullptr);
|
|
|
|
initSendBuffer<T>(sendBuf, count, rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclAllReduce(sendBuf, recvBuf, count, getNcclDataType<T>(), ncclSum, getActiveCommunicator(), getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvBuf, count, nRanks));
|
|
|
|
REGLogChecker checker = getLogChecker();
|
|
TEST_INFO("SymmetricWindow_MultiSegment: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
ASSERT_TRUE(checker.hasNumSegments(kNumSegments))
|
|
<< "Expected 'numSegments " << kNumSegments
|
|
<< "' in log - the symmetric-window multi-segment LSA registration "
|
|
"(symMemoryMapLsaTeam) did not fire";
|
|
}
|
|
|
|
/**
|
|
* @brief Multi-segment registration on the combined LSA + GIN symmetric path.
|
|
*
|
|
* Cross-node ReduceScatter over symmetric windows (NCCL_WIN_COLL_SYMMETRIC)
|
|
* whose VA spans several physical cuMem segments.
|
|
* - LSA: symMemoryMapLsaTeam maps each segment per peer and logs
|
|
* "... numSegments <N>" (dev_runtime.cc).
|
|
* - GIN: symMemoryRegisterGin registers the (all-device, contiguous)
|
|
* multi-segment buffer for the inter-node proxy.
|
|
*
|
|
* Requires >=2 nodes (a GIN inter-node leg) and >=2 ranks per node (an LSA
|
|
* intra-node leg); otherwise the path under test is not exercised and the
|
|
* test SKIPs.
|
|
*
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Symmetric_LsaGin)
|
|
{
|
|
const int nodeCount = MPITestConstants::detectNodeCount();
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
if (nodeCount < 2) {
|
|
GTEST_SKIP() << "LSA+GIN ReduceScatter requires >=2 nodes";
|
|
}
|
|
if (!isGinEnabled()) {
|
|
GTEST_SKIP() << "Requires GIN (NCCL_GIN_ENABLE!=0 and NCCL_GIN_TYPE=2)";
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isCuMemEnabled()) << "NCCL_CUMEM_ENABLE must be set to 1";
|
|
ASSERT_TRUE(isWinEnabled()) << "NCCL_WIN_ENABLE must not be set to 0";
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
if (nRanks / nodeCount < 2) {
|
|
GTEST_SKIP() << "LSA+GIN ReduceScatter requires >=2 ranks per node (LSA intra-node leg)";
|
|
}
|
|
|
|
constexpr size_t kSegmentSize = 4 * 1024 * 1024;
|
|
constexpr int kNumSegments = 4;
|
|
|
|
SCOPED_TRACE("kNumSegments=" + std::to_string(kNumSegments) +
|
|
" nodeCount=" + std::to_string(nodeCount) +
|
|
" nRanks=" + std::to_string(nRanks));
|
|
|
|
MultiSegmentBuffer recvSeg;
|
|
ASSERT_NO_FATAL_FAILURE(createMultiSegmentBuffer(dev, kSegmentSize, kNumSegments, recvSeg));
|
|
if (recvSeg.totalSize == 0) {
|
|
GTEST_SKIP() << "Raw VMM (hipMemCreate / Reserve / Map) not supported on this runtime";
|
|
}
|
|
auto recvVmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(recvSeg); });
|
|
|
|
ASSERT_EQ(recvSeg.totalSize % sizeof(T), 0u);
|
|
const size_t recvCount = recvSeg.totalSize / sizeof(T);
|
|
|
|
MultiSegmentBuffer sendSeg;
|
|
ASSERT_NO_FATAL_FAILURE(
|
|
createMultiSegmentBuffer(dev, recvSeg.segmentSize * static_cast<size_t>(nRanks),
|
|
kNumSegments, sendSeg));
|
|
if (sendSeg.totalSize == 0) {
|
|
GTEST_SKIP() << "Raw VMM (hipMemCreate / Reserve / Map) not supported on this runtime";
|
|
}
|
|
auto sendVmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(sendSeg); });
|
|
|
|
const size_t sendCount = recvCount * static_cast<size_t>(nRanks);
|
|
ASSERT_GE(sendSeg.totalSize, sendCount * sizeof(T));
|
|
|
|
ncclWindow_t sendWin = nullptr;
|
|
ncclWindow_t recvWin = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommWindowRegister(getActiveCommunicator(), sendSeg.vaBase, sendSeg.totalSize, &sendWin, NCCL_WIN_COLL_SYMMETRIC));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommWindowRegister(getActiveCommunicator(), recvSeg.vaBase, recvSeg.totalSize, &recvWin, NCCL_WIN_COLL_SYMMETRIC));
|
|
auto winCleanup = makeScopeGuard([&]() {
|
|
if (sendWin) HIP_EXPECT(ncclCommWindowDeregister(getActiveCommunicator(), sendWin));
|
|
if (recvWin) HIP_EXPECT(ncclCommWindowDeregister(getActiveCommunicator(), recvWin));
|
|
});
|
|
ASSERT_MPI_NE(sendWin, nullptr);
|
|
ASSERT_MPI_NE(recvWin, nullptr);
|
|
|
|
initSendBuffer<T>(sendSeg.vaBase, sendCount, rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclReduceScatter(sendSeg.vaBase, recvSeg.vaBase, recvCount, getNcclDataType<T>(), ncclSum, getActiveCommunicator(), getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
ASSERT_TRUE(verifyReduceScatterResult<T>(recvSeg.vaBase, recvCount, nRanks));
|
|
|
|
REGLogChecker checker = getLogChecker();
|
|
TEST_INFO("LsaGin_MultiSegment_ReduceScatter: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
ASSERT_TRUE(checker.hasNumSegments(kNumSegments))
|
|
<< "Expected 'numSegments " << kNumSegments
|
|
<< "' in log - the symmetric-window multi-segment registration "
|
|
"(symMemoryMapLsaTeam) did not fire for the LSA+GIN ReduceScatter path";
|
|
}
|
|
|
|
/**
|
|
* @brief LSA-only elastic buffer: symmetric window with a
|
|
* host-backed segment.
|
|
*
|
|
* Builds a multi-segment symmetric window (NCCL_WIN_COLL_SYMMETRIC) whose
|
|
* trailing segment is CPU memory (hipMemLocationTypeHost). With elastic
|
|
* registration enabled, ncclCommWindowRegister must accept the mixed buffer,
|
|
* the LSA team maps every segment (symMemoryMapLsaTeam logs per segment), and
|
|
* an AllReduce over the window must produce correct results - the symmetric
|
|
* scheduler falls back to legacy kernels for windows with sysmem segments
|
|
* (symmetric_sched.cc).
|
|
*
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Symmetric_Elastic_Lsa)
|
|
{
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isCuMemEnabled()) << "NCCL_CUMEM_ENABLE must be set to 1";
|
|
ASSERT_TRUE(isWinEnabled()) << "NCCL_WIN_ENABLE must not be set to 0";
|
|
ASSERT_TRUE(isElasticBufferRegisterEnabled())
|
|
<< "NCCL_ELASTIC_BUFFER_REGISTER must not be 0 for the elastic (host-backed) path";
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
constexpr size_t kSegmentSize = 4 * 1024 * 1024;
|
|
constexpr int kNumSegments = 4;
|
|
constexpr int kNumHostSegments = 1; // trailing segment is host-backed
|
|
|
|
SCOPED_TRACE("kNumSegments=" + std::to_string(kNumSegments) +
|
|
" kNumHostSegments=" + std::to_string(kNumHostSegments) +
|
|
" nRanks=" + std::to_string(nRanks));
|
|
|
|
MultiSegmentBuffer buf;
|
|
ASSERT_NO_FATAL_FAILURE(
|
|
createMixedMultiSegmentBuffer(dev, kSegmentSize, kNumSegments, kNumHostSegments, buf));
|
|
if (buf.totalSize == 0) {
|
|
GTEST_SKIP() << "Host VMM (hipMemCreate with hipMemLocationTypeHost) not supported on this runtime";
|
|
}
|
|
auto vmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(buf); });
|
|
|
|
// Split the window into send/recv halves. With kNumHostSegments=1 the recv
|
|
// half straddles the host-backed trailing segment, exercising host access.
|
|
const size_t halfSize = buf.totalSize / 2;
|
|
ASSERT_EQ(halfSize % sizeof(T), 0u);
|
|
char* base = reinterpret_cast<char*>(buf.vaBase);
|
|
void* sendBuf = base;
|
|
void* recvBuf = base + halfSize;
|
|
size_t count = halfSize / sizeof(T);
|
|
|
|
ncclWindow_t win = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommWindowRegister(getActiveCommunicator(), buf.vaBase, buf.totalSize, &win, NCCL_WIN_COLL_SYMMETRIC));
|
|
auto winCleanup = makeScopeGuard([&]() {
|
|
if (win) HIP_EXPECT(ncclCommWindowDeregister(getActiveCommunicator(), win));
|
|
});
|
|
ASSERT_MPI_NE(win, nullptr);
|
|
|
|
initSendBuffer<T>(sendBuf, count, rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclAllReduce(sendBuf, recvBuf, count, getNcclDataType<T>(), ncclSum, getActiveCommunicator(), getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
ASSERT_TRUE(verifyAllReduceResult<T>(recvBuf, count, nRanks));
|
|
|
|
REGLogChecker checker = getLogChecker();
|
|
TEST_INFO("LsaElastic_MultiSegment_AllReduce: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
ASSERT_TRUE(checker.hasNumSegments(kNumSegments))
|
|
<< "Expected 'numSegments " << kNumSegments
|
|
<< "' in log - the symmetric-window multi-segment LSA registration "
|
|
"(symMemoryMapLsaTeam) did not fire for the elastic (host-backed) buffer";
|
|
}
|
|
|
|
/**
|
|
* @brief Elastic buffer registration gating: a host-backed symmetric window must be rejected when
|
|
* NCCL_ELASTIC_BUFFER_REGISTER=0.
|
|
*
|
|
* Only meaningful when elastic registration is disabled; otherwise SKIPs. The
|
|
* rejection happens in ncclCommWindowRegister before any collective bootstrap,
|
|
* so all ranks fail symmetrically.
|
|
*/
|
|
TEST_F(UBR_MultiSegment, Symmetric_Elastic_Gating)
|
|
{
|
|
if (!validateTestPrerequisites(/*min_processes=*/2)) {
|
|
GTEST_SKIP() << "Requires 2+ ranks";
|
|
}
|
|
if (isElasticBufferRegisterEnabled()) {
|
|
GTEST_SKIP() << "Run with NCCL_ELASTIC_BUFFER_REGISTER=0 to exercise the rejection path";
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
|
|
ASSERT_TRUE(isCuMemEnabled()) << "NCCL_CUMEM_ENABLE must be set to 1";
|
|
ASSERT_TRUE(isWinEnabled()) << "NCCL_WIN_ENABLE must not be set to 0";
|
|
|
|
int dev = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGetDevice(&dev));
|
|
|
|
constexpr size_t kSegmentSize = 4 * 1024 * 1024;
|
|
constexpr int kNumSegments = 2;
|
|
constexpr int kNumHostSegments = 1;
|
|
|
|
MultiSegmentBuffer buf;
|
|
ASSERT_NO_FATAL_FAILURE(
|
|
createMixedMultiSegmentBuffer(dev, kSegmentSize, kNumSegments, kNumHostSegments, buf));
|
|
if (buf.totalSize == 0) {
|
|
GTEST_SKIP() << "Host VMM (hipMemCreate with hipMemLocationTypeHost) not supported on this runtime";
|
|
}
|
|
auto vmmCleanup = makeScopeGuard([&]() { releaseMultiSegmentBuffer(buf); });
|
|
|
|
SCOPED_TRACE("Host-backed symmetric window registration must be rejected "
|
|
"when NCCL_ELASTIC_BUFFER_REGISTER=0");
|
|
ncclWindow_t win = nullptr;
|
|
ncclResult_t rc = ncclCommWindowRegister(getActiveCommunicator(), buf.vaBase, buf.totalSize, &win, NCCL_WIN_COLL_SYMMETRIC);
|
|
if (win) HIP_EXPECT(ncclCommWindowDeregister(getActiveCommunicator(), win));
|
|
ASSERT_MPI_NE(rc, ncclSuccess);
|
|
}
|
|
|
|
// ============================================================================
|
|
// UBR Concurrent Registration Hang Reproduction Test
|
|
// ============================================================================
|
|
|
|
/**
|
|
* @brief Test to reproduce UBR hang with multiple separate buffer registrations
|
|
*
|
|
* This test reproduces a hang that occurs when:
|
|
* 1. Multiple SEPARATE buffers are allocated (not views of a contiguous buffer)
|
|
* 2. Each buffer is registered with ncclCommRegister
|
|
* 3. All buffers are used inside ncclGroupStart/ncclGroupEnd with send/recv
|
|
*
|
|
* The hang is caused by concurrent blocking IPC registration calls
|
|
* (ncclProxyCallBlocking) inside the grouped operation creating circular
|
|
* wait conditions.
|
|
*
|
|
* ROOT CAUSE:
|
|
* - ipcRegisterBuffer() in p2p.cc calls ncclProxyCallBlocking() which is a
|
|
* synchronous RPC that blocks waiting for peer to import the IPC handle
|
|
* - When all ranks simultaneously try to register multiple different buffers
|
|
* with each other inside a group, they create circular dependencies
|
|
* - Rank 0 waits for Rank 1 to import buffer A, while Rank 1 waits for Rank 0
|
|
* to import buffer B -> DEADLOCK
|
|
*
|
|
* EXPECTED BEHAVIOR:
|
|
* - SeparateBuffers_Grouped: HANGS (demonstrates the bug)
|
|
* - ContiguousBuffer_Grouped: WORKS (demonstrates the fix)
|
|
*
|
|
* To run this test with a timeout to detect the hang:
|
|
* timeout 60 mpirun -np 4 ./rccl-UnitTestsMPI --gtest_filter=UBR_ConcurrentRegHang*
|
|
*/
|
|
class UBR_ConcurrentRegHang : public RegistrationTestBase
|
|
{
|
|
protected:
|
|
using T = RegTestConfig::DefaultType;
|
|
static constexpr size_t ELEMENTS_PER_PEER = 1024; // Small size to trigger hang quickly
|
|
|
|
// Allocate and register multiple SEPARATE buffers (one per peer)
|
|
// This is the pattern that causes the hang
|
|
struct MultiBufferInfo {
|
|
std::vector<void*> buffers;
|
|
std::vector<void*> handles;
|
|
size_t countPerBuffer = 0;
|
|
int nPeers = 0;
|
|
bool allocated = false;
|
|
};
|
|
|
|
MultiBufferInfo allocateSeparateBuffers(size_t countPerBuffer, int nPeers)
|
|
{
|
|
MultiBufferInfo info;
|
|
info.countPerBuffer = countPerBuffer;
|
|
info.nPeers = nPeers;
|
|
info.buffers.resize(nPeers, nullptr);
|
|
info.handles.resize(nPeers, nullptr);
|
|
|
|
size_t bufSize = countPerBuffer * sizeof(T);
|
|
|
|
for (int i = 0; i < nPeers; i++) {
|
|
// Allocate SEPARATE buffer for each peer
|
|
if (hipMalloc(&info.buffers[i], bufSize) != hipSuccess) {
|
|
cleanupMultiBuffers(info);
|
|
return info;
|
|
}
|
|
|
|
// Register each buffer separately
|
|
ncclResult_t result = ncclCommRegister(getActiveCommunicator(),
|
|
info.buffers[i], bufSize,
|
|
&info.handles[i]);
|
|
if (result != ncclSuccess || info.handles[i] == nullptr) {
|
|
TEST_WARN("Failed to register buffer %d", i);
|
|
}
|
|
}
|
|
|
|
info.allocated = true;
|
|
return info;
|
|
}
|
|
|
|
void cleanupMultiBuffers(MultiBufferInfo& info)
|
|
{
|
|
for (int i = 0; i < info.nPeers; i++) {
|
|
if (info.handles[i]) {
|
|
ncclCommDeregister(getActiveCommunicator(), info.handles[i]);
|
|
info.handles[i] = nullptr;
|
|
}
|
|
if (info.buffers[i]) {
|
|
(void)hipFree(info.buffers[i]);
|
|
info.buffers[i] = nullptr;
|
|
}
|
|
}
|
|
info.allocated = false;
|
|
}
|
|
|
|
// Allocate a SINGLE contiguous buffer and create views (the fix)
|
|
struct ContiguousBufferInfo {
|
|
void* contiguousBuffer = nullptr;
|
|
void* handle = nullptr;
|
|
std::vector<void*> views; // Pointers into contiguous buffer
|
|
size_t countPerView = 0;
|
|
int nViews = 0;
|
|
bool allocated = false;
|
|
};
|
|
|
|
ContiguousBufferInfo allocateContiguousWithViews(size_t countPerView, int nViews)
|
|
{
|
|
ContiguousBufferInfo info;
|
|
info.countPerView = countPerView;
|
|
info.nViews = nViews;
|
|
info.views.resize(nViews, nullptr);
|
|
|
|
size_t totalSize = countPerView * nViews * sizeof(T);
|
|
|
|
// Allocate ONE contiguous buffer
|
|
if (hipMalloc(&info.contiguousBuffer, totalSize) != hipSuccess) {
|
|
return info;
|
|
}
|
|
|
|
// Register the contiguous buffer ONCE
|
|
ncclResult_t result = ncclCommRegister(getActiveCommunicator(),
|
|
info.contiguousBuffer, totalSize,
|
|
&info.handle);
|
|
if (result != ncclSuccess) {
|
|
(void)hipFree(info.contiguousBuffer);
|
|
info.contiguousBuffer = nullptr;
|
|
return info;
|
|
}
|
|
|
|
// Create views (pointers) into the contiguous buffer
|
|
char* base = static_cast<char*>(info.contiguousBuffer);
|
|
for (int i = 0; i < nViews; i++) {
|
|
info.views[i] = base + (i * countPerView * sizeof(T));
|
|
}
|
|
|
|
info.allocated = true;
|
|
return info;
|
|
}
|
|
|
|
void cleanupContiguousBuffer(ContiguousBufferInfo& info)
|
|
{
|
|
if (info.handle) {
|
|
ncclCommDeregister(getActiveCommunicator(), info.handle);
|
|
info.handle = nullptr;
|
|
}
|
|
if (info.contiguousBuffer) {
|
|
(void)hipFree(info.contiguousBuffer);
|
|
info.contiguousBuffer = nullptr;
|
|
}
|
|
info.views.clear();
|
|
info.allocated = false;
|
|
}
|
|
};
|
|
|
|
/**
|
|
* @brief THIS TEST SHOULD HANG - Demonstrates the UBR concurrent registration bug
|
|
*
|
|
* Pattern that causes hang:
|
|
* - N separate send buffers, each registered separately
|
|
* - N separate recv buffers, each registered separately
|
|
* - ncclGroupStart()
|
|
* - N x ncclSend() + N x ncclRecv() with different buffers
|
|
* - ncclGroupEnd()
|
|
*
|
|
* The ncclProxyCallBlocking() calls inside the group create circular waits.
|
|
*
|
|
* IMPORTANT: Run with timeout to detect hang:
|
|
* timeout 30 mpirun -np 4 ./rccl-UnitTestsMPI --gtest_filter=*SeparateBuffers_Grouped*
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, SeparateBuffers_Grouped_EXPECTED_HANG)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
TEST_INFO("Rank %d: Testing SEPARATE buffer pattern (expected to HANG)", rank);
|
|
TEST_INFO("Rank %d: Allocating %d separate send buffers + %d separate recv buffers",
|
|
rank, nRanks, nRanks);
|
|
|
|
// Allocate SEPARATE buffers for each peer (this causes the hang)
|
|
MultiBufferInfo sendInfo = allocateSeparateBuffers(ELEMENTS_PER_PEER, nRanks);
|
|
MultiBufferInfo recvInfo = allocateSeparateBuffers(ELEMENTS_PER_PEER, nRanks);
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupMultiBuffers(sendInfo);
|
|
cleanupMultiBuffers(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_TRUE(sendInfo.allocated);
|
|
ASSERT_MPI_TRUE(recvInfo.allocated);
|
|
|
|
// Initialize send buffers
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendInfo.buffers[peer], ELEMENTS_PER_PEER,
|
|
[rank, peer](size_t) {
|
|
return static_cast<T>(static_cast<float>(rank * 100 + peer));
|
|
}));
|
|
}
|
|
|
|
TEST_INFO("Rank %d: Starting ncclGroupStart/ncclGroupEnd with %d send/recv pairs",
|
|
rank, nRanks);
|
|
TEST_INFO("Rank %d: >>> If this hangs, the test has reproduced the bug <<<", rank);
|
|
|
|
// This grouped operation with separate buffers should HANG
|
|
// because of concurrent blocking IPC registration
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
ncclSend(sendInfo.buffers[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
ncclRecv(recvInfo.buffers[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
}
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd()); // THIS IS WHERE IT HANGS
|
|
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
TEST_INFO("Rank %d: Completed (if you see this, the hang did not occur)", rank);
|
|
|
|
// Verify results
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
bool verified = verifyBufferData<T>(recvInfo.buffers[peer], ELEMENTS_PER_PEER,
|
|
[peer, rank](size_t) {
|
|
return static_cast<T>(static_cast<float>(peer * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(verified) << "Data verification failed for peer " << peer;
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief THIS TEST SHOULD WORK - Demonstrates the fix using contiguous buffer
|
|
*
|
|
* Pattern that works:
|
|
* - 1 contiguous send buffer with views for each peer
|
|
* - 1 contiguous recv buffer with views for each peer
|
|
* - Only 2 buffer registrations total (not 2*N)
|
|
* - ncclGroupStart()
|
|
* - N x ncclSend() + N x ncclRecv() with views
|
|
* - ncclGroupEnd()
|
|
*
|
|
* Works because UBR only needs to register 2 buffers, not 2*N.
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, ContiguousBuffer_Grouped_SHOULD_WORK)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
TEST_INFO("Rank %d: Testing CONTIGUOUS buffer pattern (should work)", rank);
|
|
TEST_INFO("Rank %d: Allocating 1 contiguous send buffer + 1 contiguous recv buffer with %d views each",
|
|
rank, nRanks);
|
|
|
|
// Allocate contiguous buffers with views (this is the fix)
|
|
ContiguousBufferInfo sendInfo = allocateContiguousWithViews(ELEMENTS_PER_PEER, nRanks);
|
|
ContiguousBufferInfo recvInfo = allocateContiguousWithViews(ELEMENTS_PER_PEER, nRanks);
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupContiguousBuffer(sendInfo);
|
|
cleanupContiguousBuffer(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_TRUE(sendInfo.allocated);
|
|
ASSERT_MPI_TRUE(recvInfo.allocated);
|
|
|
|
// Initialize send buffer views
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendInfo.views[peer], ELEMENTS_PER_PEER,
|
|
[rank, peer](size_t) {
|
|
return static_cast<T>(static_cast<float>(rank * 100 + peer));
|
|
}));
|
|
}
|
|
|
|
TEST_INFO("Rank %d: Starting ncclGroupStart/ncclGroupEnd with %d send/recv pairs using views",
|
|
rank, nRanks);
|
|
|
|
// This grouped operation with contiguous buffer views should WORK
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
ncclSend(sendInfo.views[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
ncclRecv(recvInfo.views[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
}
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
TEST_INFO("Rank %d: Completed successfully", rank);
|
|
|
|
// Verify results
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
bool verified = verifyBufferData<T>(recvInfo.views[peer], ELEMENTS_PER_PEER,
|
|
[peer, rank](size_t) {
|
|
return static_cast<T>(static_cast<float>(peer * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(verified) << "Data verification failed for peer " << peer;
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief ncclAllToAll with CONTIGUOUS registered buffers - SHOULD WORK
|
|
*
|
|
* This test uses ncclAllToAll directly (not grouped send/recv) with
|
|
* a single contiguous registered buffer. This matches the correct
|
|
* pattern used after fixing the PARAM benchmark.
|
|
*
|
|
* ncclAllToAll internally uses contiguous send/recv buffers, so
|
|
* only 2 buffer registrations are needed.
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, AllToAll_ContiguousBuffer_SHOULD_WORK)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const size_t countPerRank = ELEMENTS_PER_PEER;
|
|
const size_t totalCount = countPerRank * nRanks;
|
|
const size_t totalSize = totalCount * sizeof(T);
|
|
|
|
TEST_INFO("Rank %d: Testing ncclAllToAll with CONTIGUOUS registered buffers", rank);
|
|
|
|
// Allocate single contiguous buffers
|
|
RegInfo sendInfo = allocateAndRegister(totalSize);
|
|
RegInfo recvInfo = allocateAndRegister(totalSize);
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
// Initialize send buffer: each chunk destined for peer i contains (rank*100 + i)
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendInfo.buffer, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int destRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(rank * 100 + destRank));
|
|
}));
|
|
|
|
TEST_INFO("Rank %d: Calling ncclAllToAll with registered contiguous buffers", rank);
|
|
|
|
ncclResult_t result = ncclAllToAll(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(),
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, result);
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
TEST_INFO("Rank %d: ncclAllToAll completed successfully", rank);
|
|
|
|
// Verify: chunk i should contain (srcRank*100 + myRank) from srcRank=i
|
|
bool verified = verifyBufferData<T>(recvInfo.buffer, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int srcRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(srcRank * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(verified);
|
|
|
|
TEST_INFO("Rank %d: Data verification passed", rank);
|
|
}
|
|
|
|
/**
|
|
* @brief Stress test: Multiple iterations of ncclAllToAll with contiguous buffers
|
|
*
|
|
* This test runs many iterations of ncclAllToAll with registered contiguous
|
|
* buffers to verify stability under repeated use.
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, AllToAll_ContiguousBuffer_Stress_SHOULD_WORK)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const size_t countPerRank = ELEMENTS_PER_PEER;
|
|
const size_t totalCount = countPerRank * nRanks;
|
|
const size_t totalSize = totalCount * sizeof(T);
|
|
const int NUM_ITERATIONS = 50;
|
|
|
|
TEST_INFO("Rank %d: Stress test - %d iterations of ncclAllToAll with contiguous buffers",
|
|
rank, NUM_ITERATIONS);
|
|
|
|
RegInfo sendInfo = allocateAndRegister(totalSize);
|
|
RegInfo recvInfo = allocateAndRegister(totalSize);
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
cleanupRegInfo(sendInfo);
|
|
cleanupRegInfo(recvInfo);
|
|
});
|
|
|
|
ASSERT_MPI_NE(sendInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.buffer, nullptr);
|
|
ASSERT_MPI_NE(sendInfo.handle, nullptr);
|
|
ASSERT_MPI_NE(recvInfo.handle, nullptr);
|
|
|
|
int errors = 0;
|
|
|
|
for (int iter = 0; iter < NUM_ITERATIONS; iter++) {
|
|
// Initialize with iteration-dependent pattern
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendInfo.buffer, totalCount,
|
|
[rank, countPerRank, iter](size_t i) {
|
|
int destRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(rank * 100 + destRank + iter));
|
|
}));
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemset(recvInfo.buffer, 0, totalSize));
|
|
|
|
ncclResult_t result = ncclAllToAll(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(),
|
|
getActiveCommunicator(), getActiveStream());
|
|
if (result != ncclSuccess) {
|
|
errors++;
|
|
continue;
|
|
}
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
bool verified = verifyBufferData<T>(recvInfo.buffer, totalCount,
|
|
[rank, countPerRank, iter](size_t i) {
|
|
int srcRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(srcRank * 100 + rank + iter));
|
|
});
|
|
|
|
if (!verified) {
|
|
errors++;
|
|
if (rank == 0) {
|
|
TEST_WARN("Iteration %d: verification failed", iter);
|
|
}
|
|
}
|
|
|
|
if ((iter + 1) % 10 == 0 && rank == 0) {
|
|
TEST_INFO("Completed %d/%d iterations, errors=%d", iter + 1, NUM_ITERATIONS, errors);
|
|
}
|
|
}
|
|
|
|
ASSERT_EQ(0, errors);
|
|
TEST_INFO("Rank %d: Stress test passed - %d iterations with 0 errors", rank, NUM_ITERATIONS);
|
|
}
|
|
|
|
/**
|
|
* @brief Side-by-side comparison: Contiguous vs Separate buffer patterns
|
|
*
|
|
* This test runs both patterns back-to-back to clearly demonstrate
|
|
* that the contiguous buffer pattern works while separate buffers
|
|
* would hang (the separate buffer test is commented out to avoid hang).
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, ContiguousVsSeparate_Comparison_SHOULD_WORK)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
TEST_INFO("Rank %d: Comparison test - demonstrating contiguous buffer pattern works", rank);
|
|
|
|
// ========================================
|
|
// PATTERN 1: Contiguous buffer (WORKS)
|
|
// ========================================
|
|
TEST_INFO("Rank %d: --- PATTERN 1: Contiguous buffer (should work) ---", rank);
|
|
|
|
ContiguousBufferInfo sendContig = allocateContiguousWithViews(ELEMENTS_PER_PEER, nRanks);
|
|
ContiguousBufferInfo recvContig = allocateContiguousWithViews(ELEMENTS_PER_PEER, nRanks);
|
|
|
|
auto cleanupContig = makeScopeGuard([&]() {
|
|
cleanupContiguousBuffer(sendContig);
|
|
cleanupContiguousBuffer(recvContig);
|
|
});
|
|
|
|
ASSERT_MPI_TRUE(sendContig.allocated);
|
|
ASSERT_MPI_TRUE(recvContig.allocated);
|
|
|
|
// Log buffer info
|
|
TEST_INFO("Rank %d: Contiguous send buffer: base=%p, handle=%p, %d views",
|
|
rank, sendContig.contiguousBuffer, sendContig.handle, nRanks);
|
|
TEST_INFO("Rank %d: Contiguous recv buffer: base=%p, handle=%p, %d views",
|
|
rank, recvContig.contiguousBuffer, recvContig.handle, nRanks);
|
|
|
|
// Initialize
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendContig.views[peer], ELEMENTS_PER_PEER,
|
|
[rank, peer](size_t) {
|
|
return static_cast<T>(static_cast<float>(rank * 100 + peer));
|
|
}));
|
|
}
|
|
|
|
// Execute grouped send/recv with contiguous views
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
ncclSend(sendContig.views[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
ncclRecv(recvContig.views[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
}
|
|
}
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
// Verify
|
|
bool pattern1_ok = true;
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
bool verified = verifyBufferData<T>(recvContig.views[peer], ELEMENTS_PER_PEER,
|
|
[peer, rank](size_t) {
|
|
return static_cast<T>(static_cast<float>(peer * 100 + rank));
|
|
});
|
|
if (!verified) {
|
|
pattern1_ok = false;
|
|
TEST_WARN("Rank %d: Contiguous pattern verification failed for peer %d", rank, peer);
|
|
}
|
|
}
|
|
}
|
|
|
|
ASSERT_TRUE(pattern1_ok);
|
|
TEST_INFO("Rank %d: PATTERN 1 (Contiguous) completed successfully!", rank);
|
|
|
|
// ========================================
|
|
// PATTERN 2: Separate buffers (WOULD HANG)
|
|
// ========================================
|
|
// NOTE: We skip actually running this pattern because it would hang.
|
|
// The SeparateBuffers_Grouped_EXPECTED_HANG test demonstrates this.
|
|
TEST_INFO("Rank %d: --- PATTERN 2: Separate buffers (skipped - would hang) ---", rank);
|
|
TEST_INFO("Rank %d: To see the hang, run: UBR_ConcurrentRegHang.SeparateBuffers_Grouped_EXPECTED_HANG", rank);
|
|
|
|
// Summary
|
|
TEST_INFO("Rank %d: ===== COMPARISON SUMMARY =====", rank);
|
|
TEST_INFO("Rank %d: Contiguous buffer pattern: PASSED", rank);
|
|
TEST_INFO("Rank %d: Separate buffer pattern: SKIPPED (known to hang)", rank);
|
|
}
|
|
|
|
/**
|
|
* @brief Comparison test - Run same operation without UBR registration
|
|
*
|
|
* This test allocates separate buffers but does NOT register them.
|
|
* It should work because without explicit registration, RCCL doesn't
|
|
* attempt IPC registration during the collective.
|
|
*/
|
|
TEST_F(UBR_ConcurrentRegHang, SeparateBuffers_NoRegistration_SHOULD_WORK)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
}
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
TEST_INFO("Rank %d: Testing SEPARATE buffers WITHOUT registration (baseline)", rank);
|
|
|
|
// Allocate separate buffers but DO NOT register them
|
|
std::vector<void*> sendBufs(nRanks, nullptr);
|
|
std::vector<void*> recvBufs(nRanks, nullptr);
|
|
size_t bufSize = ELEMENTS_PER_PEER * sizeof(T);
|
|
|
|
auto cleanup = makeScopeGuard([&]() {
|
|
for (int i = 0; i < nRanks; i++) {
|
|
if (sendBufs[i]) (void)hipFree(sendBufs[i]);
|
|
if (recvBufs[i]) (void)hipFree(recvBufs[i]);
|
|
}
|
|
});
|
|
|
|
for (int i = 0; i < nRanks; i++) {
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&sendBufs[i], bufSize));
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&recvBufs[i], bufSize));
|
|
}
|
|
|
|
// Initialize
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendBufs[peer], ELEMENTS_PER_PEER,
|
|
[rank, peer](size_t) {
|
|
return static_cast<T>(static_cast<float>(rank * 100 + peer));
|
|
}));
|
|
}
|
|
|
|
TEST_INFO("Rank %d: Starting grouped send/recv without UBR registration", rank);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
ncclSend(sendBufs[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
ncclRecv(recvBufs[peer], ELEMENTS_PER_PEER, getNcclDataType<T>(),
|
|
peer, getActiveCommunicator(), getActiveStream());
|
|
}
|
|
}
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
TEST_INFO("Rank %d: Completed successfully (no UBR = no hang)", rank);
|
|
|
|
// Verify results
|
|
for (int peer = 0; peer < nRanks; peer++) {
|
|
if (peer != rank) {
|
|
bool verified = verifyBufferData<T>(recvBufs[peer], ELEMENTS_PER_PEER,
|
|
[peer, rank](size_t) {
|
|
return static_cast<T>(static_cast<float>(peer * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(verified) << "Data verification failed for peer " << peer;
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Graph Capture Registration Tests
|
|
// ============================================================================
|
|
|
|
class GraphCapture_AllToAll : public RegistrationTestBase {};
|
|
|
|
TEST_F(GraphCapture_AllToAll, MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
ASSERT_TRUE(isGraphRegisterEnabled()) << "NCCL_GRAPH_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t countPerRank = RegTestConfig::MEDIUM_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const size_t totalCount = countPerRank * nRanks;
|
|
const size_t bufSize = totalCount * sizeof(T);
|
|
|
|
void* sendBuf = nullptr;
|
|
void* recvBuf = nullptr;
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&sendBuf, bufSize));
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&recvBuf, bufSize));
|
|
|
|
auto bufCleanup = makeScopeGuard([&]() {
|
|
if (sendBuf) (void)hipFree(sendBuf);
|
|
if (recvBuf) (void)hipFree(recvBuf);
|
|
});
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, initializeBufferWithPattern<T>(sendBuf, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int destRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(rank * 100 + destRank));
|
|
}));
|
|
|
|
hipGraph_t graph = nullptr;
|
|
hipGraphExec_t graphExec = nullptr;
|
|
|
|
// Graph capture
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamBeginCapture(getActiveStream(), hipStreamCaptureModeThreadLocal));
|
|
|
|
ncclResult_t ncclErr = ncclAllToAll(sendBuf, recvBuf, countPerRank,
|
|
getNcclDataType<T>(),
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclErr);
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamEndCapture(getActiveStream(), &graph));
|
|
ASSERT_MPI_NE(nullptr, graph);
|
|
|
|
size_t numNodes = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphGetNodes(graph, nullptr, &numNodes));
|
|
ASSERT_MPI_GT(numNodes, 0u);
|
|
TEST_INFO("Graph captured with %zu nodes", numNodes);
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));
|
|
|
|
auto graphCleanup = makeScopeGuard([&]() {
|
|
if (graphExec) (void)hipGraphExecDestroy(graphExec);
|
|
if (graph) (void)hipGraphDestroy(graph);
|
|
});
|
|
|
|
// Graph execution
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemset(recvBuf, 0, bufSize));
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphLaunch(graphExec, getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
// Verify registration
|
|
REGLogChecker checker = getLogChecker();
|
|
bool registrationDetected = checker.hasAnyRegistrationSuccess();
|
|
TEST_INFO("AllToAll_MultiNode: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
|
|
ASSERT_TRUE(registrationDetected);
|
|
|
|
// Verify results
|
|
bool resultValid = verifyBufferData<T>(recvBuf, totalCount,
|
|
[rank, countPerRank](size_t i) {
|
|
int srcRank = i / countPerRank;
|
|
return static_cast<T>(static_cast<float>(srcRank * 100 + rank));
|
|
});
|
|
ASSERT_TRUE(resultValid);
|
|
TEST_INFO("AllToAll graph test completed successfully");
|
|
}
|
|
|
|
class GraphCapture_AllReduce : public RegistrationTestBase {};
|
|
|
|
TEST_F(GraphCapture_AllReduce, MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_DEFAULT, RegTestConfig::MIN_NODES_MULTINODE)) {
|
|
GTEST_SKIP() << "Requires 2+ nodes";
|
|
}
|
|
|
|
ASSERT_TRUE(isPerRankLoggingEnabled()) << "RCCL_MPI_LOG_ALL_RANKS must be set to 1";
|
|
|
|
ASSERT_TRUE(isGraphRegisterEnabled()) << "NCCL_GRAPH_REGISTER must be set to 1";
|
|
|
|
using T = RegTestConfig::DefaultType;
|
|
const size_t count = RegTestConfig::LARGE_COUNT;
|
|
|
|
int rank, nRanks;
|
|
ncclCommUserRank(getActiveCommunicator(), &rank);
|
|
ncclCommCount(getActiveCommunicator(), &nRanks);
|
|
|
|
const size_t bufSize = count * sizeof(T);
|
|
|
|
void* sendBuf = nullptr;
|
|
void* recvBuf = nullptr;
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&sendBuf, bufSize));
|
|
ASSERT_MPI_EQ(hipSuccess, hipMalloc(&recvBuf, bufSize));
|
|
|
|
auto bufCleanup = makeScopeGuard([&]() {
|
|
if (sendBuf) (void)hipFree(sendBuf);
|
|
if (recvBuf) (void)hipFree(recvBuf);
|
|
});
|
|
|
|
initSendBuffer<T>(sendBuf, count, rank);
|
|
|
|
hipGraph_t graph = nullptr;
|
|
hipGraphExec_t graphExec = nullptr;
|
|
|
|
// Graph capture
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamBeginCapture(getActiveStream(), hipStreamCaptureModeThreadLocal));
|
|
|
|
ncclResult_t ncclErr = ncclAllReduce(sendBuf, recvBuf, count,
|
|
getNcclDataType<T>(), ncclSum,
|
|
getActiveCommunicator(), getActiveStream());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclErr);
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamEndCapture(getActiveStream(), &graph));
|
|
ASSERT_MPI_NE(nullptr, graph);
|
|
|
|
size_t numNodes = 0;
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphGetNodes(graph, nullptr, &numNodes));
|
|
ASSERT_MPI_GT(numNodes, 0u);
|
|
TEST_INFO("AllReduce graph captured with %zu nodes", numNodes);
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0));
|
|
|
|
auto graphCleanup = makeScopeGuard([&]() {
|
|
if (graphExec) (void)hipGraphExecDestroy(graphExec);
|
|
if (graph) (void)hipGraphDestroy(graph);
|
|
});
|
|
|
|
// Graph execution
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemset(recvBuf, 0, bufSize));
|
|
ASSERT_MPI_EQ(hipSuccess, hipGraphLaunch(graphExec, getActiveStream()));
|
|
ASSERT_EQ(hipSuccess, hipStreamSynchronize(getActiveStream()));
|
|
|
|
// Verify registration
|
|
REGLogChecker checker = getLogChecker();
|
|
bool registrationDetected = checker.hasAnyRegistrationSuccess();
|
|
TEST_INFO("AllReduce_MultiNode: %s (log size: %zu bytes)",
|
|
checker.getSummary().c_str(), checker.getContentLength());
|
|
|
|
ASSERT_TRUE(registrationDetected);
|
|
|
|
// Verify results
|
|
bool resultValid = verifyAllReduceResult<T>(recvBuf, count, nRanks);
|
|
ASSERT_TRUE(resultValid);
|
|
TEST_INFO("AllReduce graph test completed successfully");
|
|
}
|
|
|
|
// ============================================================================
|
|
// AllToAll/AllToAllv Interleaved Stress Test with UBR
|
|
// ============================================================================
|
|
|
|
/**
|
|
* @brief Stress test for User Buffer Registration with interleaved AllToAll/AllToAllv
|
|
*
|
|
* This test exercises buffer registration under stress conditions:
|
|
* - Interleaves AllToAll and AllToAllv calls
|
|
* - Runs hundreds of iterations
|
|
* - Tests small, medium, and large message sizes
|
|
* - Uses multiple communicators (primary + split)
|
|
* - Detects data mismatches, buffer corruption, and hangs
|
|
*/
|
|
class UBR_AllToAllStress : public RegistrationTestBase
|
|
{
|
|
protected:
|
|
// Use bfloat16 as the test type (same as RegTestConfig::DefaultType)
|
|
using T = RegTestConfig::DefaultType; // hip_bfloat16
|
|
|
|
// Test parameters
|
|
static constexpr int ITERATIONS_PER_SIZE = 100;
|
|
static constexpr size_t STRESS_SMALL_COUNT = 256; // 512B for bfloat16
|
|
static constexpr size_t STRESS_MEDIUM_COUNT = 64 * 1024; // 128KiB for bfloat16
|
|
static constexpr size_t STRESS_LARGE_COUNT = 4 * 1024 * 1024; // 8MiB for bfloat16
|
|
|
|
struct CommContext {
|
|
ncclComm_t comm = nullptr;
|
|
int rank = 0;
|
|
int nRanks = 0;
|
|
std::string name;
|
|
};
|
|
|
|
// Create CommContext from a communicator
|
|
CommContext makeCommContext(ncclComm_t comm, const std::string& name)
|
|
{
|
|
CommContext ctx;
|
|
ctx.comm = comm;
|
|
ctx.name = name;
|
|
ncclCommUserRank(comm, &ctx.rank);
|
|
ncclCommCount(comm, &ctx.nRanks);
|
|
return ctx;
|
|
}
|
|
|
|
// Allocate and register a pair of send/recv buffers
|
|
bool allocateBufferPair(size_t size, ncclComm_t comm,
|
|
RegInfo& sendInfo, RegInfo& recvInfo)
|
|
{
|
|
sendInfo.size = recvInfo.size = size;
|
|
|
|
if (hipMalloc(&sendInfo.buffer, size) != hipSuccess) return false;
|
|
if (hipMalloc(&recvInfo.buffer, size) != hipSuccess) {
|
|
(void)hipFree(sendInfo.buffer);
|
|
sendInfo.buffer = nullptr;
|
|
return false;
|
|
}
|
|
|
|
ncclResult_t r1 = ncclCommRegister(comm, sendInfo.buffer, size, &sendInfo.handle);
|
|
ncclResult_t r2 = ncclCommRegister(comm, recvInfo.buffer, size, &recvInfo.handle);
|
|
sendInfo.registered = (r1 == ncclSuccess && sendInfo.handle);
|
|
recvInfo.registered = (r2 == ncclSuccess && recvInfo.handle);
|
|
|
|
return true;
|
|
}
|
|
|
|
void cleanupBufferPair(RegInfo& sendInfo, RegInfo& recvInfo, ncclComm_t comm)
|
|
{
|
|
auto cleanup = [&](RegInfo& info) {
|
|
if (info.handle) { ncclCommDeregister(comm, info.handle); info.handle = nullptr; }
|
|
if (info.buffer) { (void)hipFree(info.buffer); info.buffer = nullptr; }
|
|
info.registered = false;
|
|
};
|
|
cleanup(sendInfo);
|
|
cleanup(recvInfo);
|
|
}
|
|
|
|
// Compute pattern value - use integer patterns for bfloat16 precision
|
|
// Pattern: rank * 100 + dest * 10 + (iter % 10) + offset
|
|
static float computePattern(int rank, int peerRank, int iter, float offset = 0.0f)
|
|
{
|
|
return static_cast<float>(rank * 100 + peerRank * 10 + (iter % 10)) + offset;
|
|
}
|
|
|
|
void initBuffer(void* buffer, size_t countPerRank, int nRanks, int rank, int iter, float offset = 0.0f)
|
|
{
|
|
(void)initializeBufferWithPattern<T>(buffer, countPerRank * nRanks,
|
|
[rank, countPerRank, iter, offset](size_t i) {
|
|
int dest = static_cast<int>(i / countPerRank);
|
|
return static_cast<T>(computePattern(rank, dest, iter, offset));
|
|
});
|
|
}
|
|
|
|
bool verifyBuffer(void* buffer, size_t countPerRank, int nRanks, int myRank, int iter, float offset = 0.0f)
|
|
{
|
|
std::vector<T> data(countPerRank * nRanks);
|
|
if (hipMemcpy(data.data(), buffer, data.size() * sizeof(T), hipMemcpyDeviceToHost) != hipSuccess)
|
|
return false;
|
|
|
|
for (int src = 0; src < nRanks; src++) {
|
|
float expected = computePattern(src, myRank, iter, offset);
|
|
for (size_t i = 0; i < countPerRank; i++) {
|
|
float actual = static_cast<float>(data[src * countPerRank + i]);
|
|
// bfloat16 has ~3 decimal digits precision, use tolerance of 1.0
|
|
if (std::abs(actual - expected) > 1.0f) {
|
|
TEST_WARN("Mismatch src=%d idx=%zu exp=%f got=%f iter=%d",
|
|
src, i, expected, actual, iter);
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
// Variable-size init/verify for AllToAllv
|
|
void initBufferV(void* buffer, const std::vector<size_t>& counts,
|
|
const std::vector<size_t>& displs, int nRanks, int rank, int iter)
|
|
{
|
|
size_t total = 0;
|
|
for (auto c : counts) total += c;
|
|
std::vector<T> data(total);
|
|
|
|
for (int dest = 0; dest < nRanks; dest++) {
|
|
T value = static_cast<T>(computePattern(rank, dest, iter, 0.5f));
|
|
for (size_t i = 0; i < counts[dest]; i++)
|
|
data[displs[dest] + i] = value;
|
|
}
|
|
(void)hipMemcpy(buffer, data.data(), data.size() * sizeof(T), hipMemcpyHostToDevice);
|
|
}
|
|
|
|
bool verifyBufferV(void* buffer, const std::vector<size_t>& counts,
|
|
const std::vector<size_t>& displs, int nRanks, int myRank, int iter)
|
|
{
|
|
size_t total = 0;
|
|
for (auto c : counts) total += c;
|
|
std::vector<T> data(total);
|
|
if (hipMemcpy(data.data(), buffer, data.size() * sizeof(T), hipMemcpyDeviceToHost) != hipSuccess)
|
|
return false;
|
|
|
|
for (int src = 0; src < nRanks; src++) {
|
|
float expected = computePattern(src, myRank, iter, 0.5f);
|
|
for (size_t i = 0; i < counts[src]; i++) {
|
|
float actual = static_cast<float>(data[displs[src] + i]);
|
|
if (std::abs(actual - expected) > 1.0f) return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
// Core stress test runner
|
|
bool runInterleavedStress(CommContext& ctx, size_t countPerRank, int iterations, const std::string& label)
|
|
{
|
|
RegInfo sendInfo, recvInfo;
|
|
size_t totalSize = countPerRank * ctx.nRanks * sizeof(T);
|
|
|
|
if (!allocateBufferPair(totalSize, ctx.comm, sendInfo, recvInfo)) {
|
|
TEST_WARN("[%s] Buffer allocation failed", label.c_str());
|
|
return false;
|
|
}
|
|
auto cleanup = makeScopeGuard([&]() { cleanupBufferPair(sendInfo, recvInfo, ctx.comm); });
|
|
|
|
// Setup AllToAllv counts
|
|
std::vector<size_t> sendcounts(ctx.nRanks), sdispls(ctx.nRanks);
|
|
std::vector<size_t> recvcounts(ctx.nRanks), rdispls(ctx.nRanks);
|
|
size_t baseCount = countPerRank / 4, sendTotal = 0, recvTotal = 0;
|
|
|
|
for (int r = 0; r < ctx.nRanks; r++) {
|
|
sendcounts[r] = ((ctx.rank + r + 1) % ctx.nRanks + 1) * baseCount;
|
|
sdispls[r] = sendTotal; sendTotal += sendcounts[r];
|
|
recvcounts[r] = ((r + ctx.rank + 1) % ctx.nRanks + 1) * baseCount;
|
|
rdispls[r] = recvTotal; recvTotal += recvcounts[r];
|
|
}
|
|
|
|
RegInfo sendInfoV, recvInfoV;
|
|
if (!allocateBufferPair(std::max(sendTotal, recvTotal) * sizeof(T), ctx.comm, sendInfoV, recvInfoV)) {
|
|
TEST_WARN("[%s] AllToAllv buffer allocation failed", label.c_str());
|
|
return false;
|
|
}
|
|
auto cleanupV = makeScopeGuard([&]() { cleanupBufferPair(sendInfoV, recvInfoV, ctx.comm); });
|
|
|
|
int errors = 0;
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
// Debug: Log buffer addresses at start (only rank 0)
|
|
if (ctx.rank == 0) {
|
|
TEST_INFO("[%s] Buffer addresses - sendBuf=%p recvBuf=%p sendBufV=%p recvBufV=%p",
|
|
label.c_str(), sendInfo.buffer, recvInfo.buffer, sendInfoV.buffer, recvInfoV.buffer);
|
|
}
|
|
|
|
for (int iter = 0; iter < iterations; iter++) {
|
|
bool ok;
|
|
const char* opType;
|
|
|
|
if (iter % 2 == 0) {
|
|
opType = "AllToAll";
|
|
initBuffer(sendInfo.buffer, countPerRank, ctx.nRanks, ctx.rank, iter);
|
|
|
|
// Debug: Show sample data before collective (first 2 iterations, rank 0)
|
|
if (iter < 2 && ctx.rank == 0) {
|
|
T sample;
|
|
(void)hipMemcpy(&sample, sendInfo.buffer, sizeof(T), hipMemcpyDeviceToHost);
|
|
TEST_INFO("[%s] iter=%d PRE-AllToAll: sendBuf[0]=%.1f (expected=%.1f)",
|
|
label.c_str(), iter, static_cast<float>(sample),
|
|
computePattern(ctx.rank, 0, iter, 0.0f));
|
|
}
|
|
|
|
(void)hipMemset(recvInfo.buffer, 0, totalSize);
|
|
if (ncclAlltoAll(sendInfo.buffer, recvInfo.buffer, countPerRank,
|
|
getNcclDataType<T>(), ctx.comm, stream) != ncclSuccess) { errors++; continue; }
|
|
(void)hipStreamSynchronize(stream);
|
|
|
|
// Debug: Show sample data after collective (first 2 iterations, rank 0)
|
|
if (iter < 2 && ctx.rank == 0) {
|
|
T recvSample;
|
|
(void)hipMemcpy(&recvSample, recvInfo.buffer, sizeof(T), hipMemcpyDeviceToHost);
|
|
TEST_INFO("[%s] iter=%d POST-AllToAll: recvBuf[0]=%.1f (expected from rank0=%.1f)",
|
|
label.c_str(), iter, static_cast<float>(recvSample),
|
|
computePattern(0, ctx.rank, iter, 0.0f));
|
|
}
|
|
|
|
ok = verifyBuffer(recvInfo.buffer, countPerRank, ctx.nRanks, ctx.rank, iter);
|
|
} else {
|
|
opType = "AllToAllv";
|
|
initBufferV(sendInfoV.buffer, sendcounts, sdispls, ctx.nRanks, ctx.rank, iter);
|
|
|
|
// Debug: Show sample data before collective (first 2 odd iterations, rank 0)
|
|
if (iter < 3 && ctx.rank == 0) {
|
|
T sample;
|
|
(void)hipMemcpy(&sample, sendInfoV.buffer, sizeof(T), hipMemcpyDeviceToHost);
|
|
TEST_INFO("[%s] iter=%d PRE-AllToAllv: sendBufV[0]=%.1f (expected=%.1f)",
|
|
label.c_str(), iter, static_cast<float>(sample),
|
|
computePattern(ctx.rank, 0, iter, 0.5f));
|
|
}
|
|
|
|
(void)hipMemset(recvInfoV.buffer, 0, recvTotal * sizeof(T));
|
|
if (ncclAlltoAllv(sendInfoV.buffer, sendcounts.data(), sdispls.data(),
|
|
recvInfoV.buffer, recvcounts.data(), rdispls.data(),
|
|
getNcclDataType<T>(), ctx.comm, stream) != ncclSuccess) { errors++; continue; }
|
|
(void)hipStreamSynchronize(stream);
|
|
|
|
// Debug: Show sample data after collective (first 2 odd iterations, rank 0)
|
|
if (iter < 3 && ctx.rank == 0) {
|
|
T recvSample;
|
|
(void)hipMemcpy(&recvSample, recvInfoV.buffer, sizeof(T), hipMemcpyDeviceToHost);
|
|
TEST_INFO("[%s] iter=%d POST-AllToAllv: recvBufV[0]=%.1f (expected from rank0=%.1f)",
|
|
label.c_str(), iter, static_cast<float>(recvSample),
|
|
computePattern(0, ctx.rank, iter, 0.5f));
|
|
}
|
|
|
|
ok = verifyBufferV(recvInfoV.buffer, recvcounts, rdispls, ctx.nRanks, ctx.rank, iter);
|
|
}
|
|
|
|
if (!ok) {
|
|
errors++;
|
|
if (ctx.rank == 0)
|
|
TEST_WARN("[%s] iter=%d %s VERIFICATION FAILED", label.c_str(), iter, opType);
|
|
}
|
|
|
|
// Log progress every 10 iterations (rank 0 only)
|
|
if ((iter + 1) % 10 == 0 && ctx.rank == 0)
|
|
TEST_INFO("[%s] Completed %d/%d iterations, errors=%d", label.c_str(), iter + 1, iterations, errors);
|
|
}
|
|
|
|
if (ctx.rank == 0)
|
|
TEST_INFO("[%s] %s: %d iterations", label.c_str(), errors ? "FAILED" : "PASSED", iterations);
|
|
return errors == 0;
|
|
}
|
|
};
|
|
|
|
TEST_F(UBR_AllToAllStress, InterleavedMultiSize_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE))
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
CommContext ctx = makeCommContext(getActiveCommunicator(), "Primary");
|
|
TEST_INFO("Starting stress test: %d ranks, %d iterations/size", ctx.nRanks, ITERATIONS_PER_SIZE);
|
|
|
|
bool ok = true;
|
|
const size_t sizes[] = {STRESS_SMALL_COUNT, STRESS_MEDIUM_COUNT, STRESS_LARGE_COUNT};
|
|
const char* names[] = {"SMALL", "MEDIUM", "LARGE"};
|
|
|
|
for (int i = 0; i < 3; i++) {
|
|
TEST_INFO("--- Testing %s (%zu elements/rank, %zu bytes/rank) ---",
|
|
names[i], sizes[i], sizes[i] * sizeof(T));
|
|
ok &= runInterleavedStress(ctx, sizes[i], ITERATIONS_PER_SIZE, names[i]);
|
|
}
|
|
ASSERT_TRUE(ok);
|
|
}
|
|
|
|
TEST_F(UBR_AllToAllStress, InterleavedWithSplitComm_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE))
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
CommContext primaryCtx = makeCommContext(getActiveCommunicator(), "Primary");
|
|
|
|
// Split into even/odd groups
|
|
ncclComm_t splitComm = nullptr;
|
|
int color = primaryCtx.rank % 2;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSplit(primaryCtx.comm, color, primaryCtx.rank, &splitComm, nullptr));
|
|
auto splitGuard = makeCommAutoGuard(splitComm);
|
|
|
|
CommContext splitCtx = makeCommContext(splitComm, color == 0 ? "Split-Even" : "Split-Odd");
|
|
if (splitCtx.nRanks < 2) GTEST_SKIP() << "Split comm too small";
|
|
|
|
TEST_INFO("Primary: %d ranks, Split: %d ranks", primaryCtx.nRanks, splitCtx.nRanks);
|
|
|
|
bool ok = true;
|
|
const size_t sizes[] = {STRESS_SMALL_COUNT, STRESS_MEDIUM_COUNT, STRESS_LARGE_COUNT};
|
|
const char* names[] = {"SMALL", "MEDIUM", "LARGE"};
|
|
|
|
for (int i = 0; i < 3; i++) {
|
|
TEST_INFO("--- Phase %d: %s ---", i, names[i]);
|
|
ok &= runInterleavedStress(primaryCtx, sizes[i], ITERATIONS_PER_SIZE / 2, std::string(names[i]) + "-Primary");
|
|
ok &= runInterleavedStress(splitCtx, sizes[i], ITERATIONS_PER_SIZE / 2, std::string(names[i]) + "-Split");
|
|
MPI_Barrier(MPI_COMM_WORLD);
|
|
}
|
|
ASSERT_TRUE(ok);
|
|
}
|
|
|
|
TEST_F(UBR_AllToAllStress, ConcurrentMultiComm_MultiNode)
|
|
{
|
|
if (!setupMultiNode(RegTestConfig::MIN_RANKS_ALLTOALL, RegTestConfig::MIN_NODES_MULTINODE))
|
|
GTEST_SKIP() << "Requires 4+ ranks across 2+ nodes";
|
|
ASSERT_TRUE(isUBREnabled()) << "NCCL_LOCAL_REGISTER must be set to 1";
|
|
|
|
ncclComm_t primaryComm = getActiveCommunicator();
|
|
CommContext ctx1 = makeCommContext(primaryComm, "Comm1");
|
|
|
|
ncclComm_t comm2 = nullptr;
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSplit(primaryComm, 0, ctx1.rank, &comm2, nullptr));
|
|
auto comm2Guard = makeCommAutoGuard(comm2);
|
|
|
|
hipStream_t stream1 = getActiveStream(), stream2;
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamCreate(&stream2));
|
|
auto stream2Guard = makeStreamAutoGuard(stream2);
|
|
|
|
const size_t countPerRank = STRESS_MEDIUM_COUNT;
|
|
const size_t totalSize = countPerRank * ctx1.nRanks * sizeof(T);
|
|
|
|
RegInfo send1, recv1, send2, recv2;
|
|
ASSERT_MPI_TRUE(allocateBufferPair(totalSize, primaryComm, send1, recv1));
|
|
ASSERT_MPI_TRUE(allocateBufferPair(totalSize, comm2, send2, recv2));
|
|
auto bufCleanup = makeScopeGuard([&]() {
|
|
cleanupBufferPair(send1, recv1, primaryComm);
|
|
cleanupBufferPair(send2, recv2, comm2);
|
|
});
|
|
|
|
TEST_INFO("Concurrent multi-comm stress: 2 comms, 2 streams, %d iterations", ITERATIONS_PER_SIZE);
|
|
|
|
int errors = 0;
|
|
for (int iter = 0; iter < ITERATIONS_PER_SIZE; iter++) {
|
|
initBuffer(send1.buffer, countPerRank, ctx1.nRanks, ctx1.rank, iter, 0.0f);
|
|
initBuffer(send2.buffer, countPerRank, ctx1.nRanks, ctx1.rank, iter, 0.5f);
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemsetAsync(recv1.buffer, 0, totalSize, stream1));
|
|
ASSERT_MPI_EQ(hipSuccess, hipMemsetAsync(recv2.buffer, 0, totalSize, stream2));
|
|
|
|
ncclGroupStart();
|
|
ncclAlltoAll(send1.buffer, recv1.buffer, countPerRank, getNcclDataType<T>(), primaryComm, stream1);
|
|
ncclAlltoAll(send2.buffer, recv2.buffer, countPerRank, getNcclDataType<T>(), comm2, stream2);
|
|
ncclGroupEnd();
|
|
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamSynchronize(stream1));
|
|
ASSERT_MPI_EQ(hipSuccess, hipStreamSynchronize(stream2));
|
|
|
|
if (!verifyBuffer(recv1.buffer, countPerRank, ctx1.nRanks, ctx1.rank, iter, 0.0f)) errors++;
|
|
if (!verifyBuffer(recv2.buffer, countPerRank, ctx1.nRanks, ctx1.rank, iter, 0.5f)) errors++;
|
|
|
|
if ((iter + 1) % 50 == 0 && ctx1.rank == 0)
|
|
TEST_INFO("Concurrent: %d/%d iterations, errors=%d", iter + 1, ITERATIONS_PER_SIZE, errors);
|
|
}
|
|
|
|
ASSERT_EQ(0, errors);
|
|
TEST_INFO("Concurrent multi-comm stress test completed successfully");
|
|
}
|
|
|
|
#endif // MPI_TESTS_ENABLED
|