Skip to content

Commit 79bd43c

Browse files
Santosh Narayankhedkarfacebook-github-bot
Santosh Narayankhedkar
authored and
facebook-github-bot
committed
Add registered-send protocol regression coverage (#3412)
Summary: Cover two protocol interactions for registered-source IBGDA sends: sharing the normal send cursor with staged sends, and maintaining backpressure and data integrity across repeated staging-slot wrap. This is test-only hardening for the registered-source transport primitive. Differential Revision: D114552006
1 parent a112817 commit 79bd43c

4 files changed

Lines changed: 316 additions & 0 deletions

File tree

‎comms/prims/tests/MultipeerIbgdaTransportTest.cc‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1963,6 +1963,205 @@ TEST_F(
19631963
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
19641964
}
19651965

1966+
TEST_F(
1967+
MultipeerIbgdaTransportTestFixture,
1968+
RegisteredAndStagedSendShareProtocolCursor) {
1969+
if (numRanks != 2) {
1970+
GTEST_SKIP() << "Skipping test: requires exactly 2 ranks, got " << numRanks;
1971+
}
1972+
if (!test::supportsProgressSendRecv()) {
1973+
GTEST_SKIP() << "registered-source send is not supported for this build";
1974+
}
1975+
1976+
constexpr std::size_t perChannelSize = 64 * 1024;
1977+
constexpr int pipelineDepth = 2;
1978+
constexpr std::size_t firstBytes = 17 * 1024 + 1;
1979+
constexpr std::size_t secondBytes = 9 * 1024 + 7;
1980+
constexpr std::size_t thirdBytes = 40 * 1024 + 17;
1981+
constexpr std::size_t totalBytes = firstBytes + secondBytes + thirdBytes;
1982+
constexpr std::size_t maxSignalBytes = 4 * 1024;
1983+
constexpr int numBlocks = 1;
1984+
constexpr int blockSize = 128;
1985+
const int peerRank = globalRank == 0 ? 1 : 0;
1986+
1987+
std::unique_ptr<MultipeerIbgdaTransport> transport;
1988+
try {
1989+
MultipeerIbgdaTransportConfig config{
1990+
.cudaDevice = localRank,
1991+
.perChannelSize = perChannelSize,
1992+
.max_num_channels = numBlocks,
1993+
.pipelineDepth = pipelineDepth,
1994+
};
1995+
auto bootstrap = std::make_shared<meta::comms::MpiBootstrap>();
1996+
transport = std::make_unique<MultipeerIbgdaTransport>(
1997+
globalRank, numRanks, bootstrap, config);
1998+
transport->exchange();
1999+
} catch (const std::exception& e) {
2000+
GTEST_SKIP() << "IBGDA transport not available: " << e.what();
2001+
}
2002+
2003+
auto* peerTransport = transport->getP2pTransportDevice(peerRank);
2004+
DeviceBuffer sendBuffer(totalBytes);
2005+
DeviceBuffer recvBuffer(totalBytes);
2006+
DeviceBuffer errorCountBuffer(sizeof(int));
2007+
auto* errorCount = static_cast<int*>(errorCountBuffer.get());
2008+
IbgdaLocalBuffer registeredSource{};
2009+
if (globalRank == 0) {
2010+
registeredSource = transport->registerBuffer(sendBuffer.get(), totalBytes);
2011+
test::fillBufferWithPattern(
2012+
sendBuffer.get(), firstBytes, 0x31, numBlocks, blockSize);
2013+
test::fillBufferWithPattern(
2014+
static_cast<char*>(sendBuffer.get()) + firstBytes,
2015+
secondBytes,
2016+
0x72,
2017+
numBlocks,
2018+
blockSize);
2019+
test::fillBufferWithPattern(
2020+
static_cast<char*>(sendBuffer.get()) + firstBytes + secondBytes,
2021+
thirdBytes,
2022+
0xB4,
2023+
numBlocks,
2024+
blockSize);
2025+
} else {
2026+
CUDACHECK_TEST(cudaMemset(recvBuffer.get(), 0, totalBytes));
2027+
}
2028+
CUDACHECK_TEST(cudaDeviceSynchronize());
2029+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2030+
2031+
test::testMixedRegisteredAndStagedSendRecv(
2032+
peerTransport,
2033+
registeredSource,
2034+
recvBuffer.get(),
2035+
firstBytes,
2036+
secondBytes,
2037+
thirdBytes,
2038+
maxSignalBytes,
2039+
globalRank == 0,
2040+
numBlocks,
2041+
blockSize);
2042+
CUDACHECK_TEST(cudaDeviceSynchronize());
2043+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2044+
2045+
if (globalRank == 1) {
2046+
const std::array<std::tuple<std::size_t, std::size_t, uint8_t>, 3> ranges{{
2047+
{0, firstBytes, 0x31},
2048+
{firstBytes, secondBytes, 0x72},
2049+
{firstBytes + secondBytes, thirdBytes, 0xB4},
2050+
}};
2051+
for (const auto& [offset, nbytes, pattern] : ranges) {
2052+
CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
2053+
test::verifyBufferPattern(
2054+
static_cast<char*>(recvBuffer.get()) + offset,
2055+
nbytes,
2056+
pattern,
2057+
errorCount,
2058+
numBlocks,
2059+
blockSize);
2060+
CUDACHECK_TEST(cudaDeviceSynchronize());
2061+
int hostErrors = 0;
2062+
CUDACHECK_TEST(cudaMemcpy(
2063+
&hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
2064+
EXPECT_EQ(hostErrors, 0)
2065+
<< "mixed registered/staged send corrupted range at " << offset;
2066+
}
2067+
} else {
2068+
transport->deregisterBuffer(sendBuffer.get());
2069+
}
2070+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2071+
}
2072+
2073+
TEST_F(
2074+
MultipeerIbgdaTransportTestFixture,
2075+
RegisteredSendBackpressureAcrossStagingWrap) {
2076+
if (numRanks != 2) {
2077+
GTEST_SKIP() << "Skipping test: requires exactly 2 ranks, got " << numRanks;
2078+
}
2079+
if (!test::supportsProgressSendRecv()) {
2080+
GTEST_SKIP() << "registered-source send is not supported for this build";
2081+
}
2082+
2083+
constexpr std::size_t perChannelSize = 64 * 1024;
2084+
constexpr int pipelineDepth = 2;
2085+
constexpr std::size_t nbytes = 8 * 1024 * 1024 + 17;
2086+
constexpr std::size_t maxSignalBytes = 4 * 1024;
2087+
constexpr int numBlocks = 1;
2088+
constexpr int blockSize = 128;
2089+
const int peerRank = globalRank == 0 ? 1 : 0;
2090+
std::vector<uint8_t> expected(nbytes);
2091+
for (std::size_t i = 0; i < nbytes; ++i) {
2092+
const std::size_t generation = i / perChannelSize;
2093+
expected[i] = static_cast<uint8_t>(generation * 37 + (i % 251));
2094+
}
2095+
2096+
std::unique_ptr<MultipeerIbgdaTransport> transport;
2097+
try {
2098+
MultipeerIbgdaTransportConfig config{
2099+
.cudaDevice = localRank,
2100+
.perChannelSize = perChannelSize,
2101+
.max_num_channels = numBlocks,
2102+
.pipelineDepth = pipelineDepth,
2103+
};
2104+
auto bootstrap = std::make_shared<meta::comms::MpiBootstrap>();
2105+
transport = std::make_unique<MultipeerIbgdaTransport>(
2106+
globalRank, numRanks, bootstrap, config);
2107+
transport->exchange();
2108+
} catch (const std::exception& e) {
2109+
GTEST_SKIP() << "IBGDA transport not available: " << e.what();
2110+
}
2111+
2112+
auto* peerTransport = transport->getP2pTransportDevice(peerRank);
2113+
DeviceBuffer sendBuffer(nbytes);
2114+
DeviceBuffer recvBuffer(nbytes);
2115+
DeviceBuffer observationBuffer(sizeof(test::RegisteredSendObservation));
2116+
auto* observation =
2117+
static_cast<test::RegisteredSendObservation*>(observationBuffer.get());
2118+
IbgdaLocalBuffer registeredSource{};
2119+
if (globalRank == 0) {
2120+
registeredSource = transport->registerBuffer(sendBuffer.get(), nbytes);
2121+
CUDACHECK_TEST(cudaMemcpy(
2122+
sendBuffer.get(), expected.data(), nbytes, cudaMemcpyHostToDevice));
2123+
CUDACHECK_TEST(
2124+
cudaMemset(observation, 0, sizeof(test::RegisteredSendObservation)));
2125+
} else {
2126+
CUDACHECK_TEST(cudaMemset(recvBuffer.get(), 0, nbytes));
2127+
}
2128+
CUDACHECK_TEST(cudaDeviceSynchronize());
2129+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2130+
2131+
test::testRegisteredSendRecv(
2132+
peerTransport,
2133+
registeredSource,
2134+
recvBuffer.get(),
2135+
nbytes,
2136+
maxSignalBytes,
2137+
globalRank == 0,
2138+
numBlocks,
2139+
blockSize,
2140+
globalRank == 0 ? observation : nullptr);
2141+
CUDACHECK_TEST(cudaDeviceSynchronize());
2142+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2143+
2144+
if (globalRank == 0) {
2145+
test::RegisteredSendObservation hostObservation{};
2146+
CUDACHECK_TEST(cudaMemcpy(
2147+
&hostObservation,
2148+
observation,
2149+
sizeof(hostObservation),
2150+
cudaMemcpyDeviceToHost));
2151+
EXPECT_GT(hostObservation.waitingCount, 0);
2152+
EXPECT_EQ(hostObservation.postedCount, 1);
2153+
EXPECT_EQ(hostObservation.drainedCount, 1);
2154+
transport->deregisterBuffer(sendBuffer.get());
2155+
} else {
2156+
std::vector<uint8_t> received(nbytes);
2157+
CUDACHECK_TEST(cudaMemcpy(
2158+
received.data(), recvBuffer.get(), nbytes, cudaMemcpyDeviceToHost));
2159+
EXPECT_EQ(received, expected)
2160+
<< "registered send corrupted data across slot wrap";
2161+
}
2162+
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
2163+
}
2164+
19662165
// =============================================================================
19672166
// Sustained chunked send/recv - repro for the GB200 per-channel DATA_READY
19682167
// deadlock (two NICs atomic-FA the same flag at maxGroups>=8). Streams a large

‎comms/prims/tests/MultipeerIbgdaTransportTest.cu‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,57 @@ __global__ void registeredSendRecvKernel(
568568
}
569569
}
570570

571+
__global__ void mixedRegisteredAndStagedSendRecvKernel(
572+
P2pIbgdaTransportDevice* transport,
573+
IbgdaLocalBuffer sendBuffer,
574+
void* recvBuffer,
575+
std::size_t firstBytes,
576+
std::size_t secondBytes,
577+
std::size_t thirdBytes,
578+
std::size_t maxSignalBytes,
579+
bool send) {
580+
auto group = make_block_group();
581+
Timeout timeout(kDefaultDeviceTimeoutCycles);
582+
timeout.start();
583+
if (send) {
584+
(void)postRegisteredSend(
585+
*transport,
586+
group,
587+
sendBuffer,
588+
firstBytes,
589+
maxSignalBytes,
590+
timeout,
591+
nullptr);
592+
transport->send(
593+
group,
594+
static_cast<const char*>(sendBuffer.ptr) + firstBytes,
595+
secondBytes,
596+
maxSignalBytes,
597+
timeout);
598+
(void)postRegisteredSend(
599+
*transport,
600+
group,
601+
sendBuffer.subBuffer(firstBytes + secondBytes),
602+
thirdBytes,
603+
maxSignalBytes,
604+
timeout,
605+
nullptr);
606+
drainRegisteredSends(*transport, group, timeout, nullptr);
607+
return;
608+
}
609+
610+
auto* output = static_cast<char*>(recvBuffer);
611+
transport->recv(group, output, firstBytes, maxSignalBytes, timeout);
612+
transport->recv(
613+
group, output + firstBytes, secondBytes, maxSignalBytes, timeout);
614+
transport->recv(
615+
group,
616+
output + firstBytes + secondBytes,
617+
thirdBytes,
618+
maxSignalBytes,
619+
timeout);
620+
}
621+
571622
__global__ void fillTransportStagingKernel(
572623
P2pIbgdaTransportDevice* transport,
573624
bool sendStaging,
@@ -719,6 +770,47 @@ void testRegisteredSendRecv(
719770
#endif
720771
}
721772

773+
void testMixedRegisteredAndStagedSendRecv(
774+
P2pIbgdaTransportDevice* transport,
775+
const IbgdaLocalBuffer& sendBuffer,
776+
void* recvBuffer,
777+
std::size_t firstBytes,
778+
std::size_t secondBytes,
779+
std::size_t thirdBytes,
780+
std::size_t maxSignalBytes,
781+
bool send,
782+
int numBlocks,
783+
int blockSize) {
784+
#ifdef __HIP_PLATFORM_AMD__
785+
(void)transport;
786+
(void)sendBuffer;
787+
(void)recvBuffer;
788+
(void)firstBytes;
789+
(void)secondBytes;
790+
(void)thirdBytes;
791+
(void)maxSignalBytes;
792+
(void)send;
793+
(void)numBlocks;
794+
(void)blockSize;
795+
throw std::runtime_error("registered-source send is NVIDIA-only");
796+
#else
797+
mixedRegisteredAndStagedSendRecvKernel<<<numBlocks, blockSize>>>(
798+
transport,
799+
sendBuffer,
800+
recvBuffer,
801+
firstBytes,
802+
secondBytes,
803+
thirdBytes,
804+
maxSignalBytes,
805+
send);
806+
const cudaError_t err = cudaGetLastError();
807+
if (err != cudaSuccess) {
808+
throw std::runtime_error(
809+
std::string("Kernel launch failed: ") + cudaGetErrorString(err));
810+
}
811+
#endif
812+
}
813+
722814
void testFillTransportStaging(
723815
P2pIbgdaTransportDevice* transport,
724816
bool sendStaging,

‎comms/prims/tests/MultipeerIbgdaTransportTest.cuh‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,16 @@ __global__ void registeredSendRecvKernel(
118118
uint8_t overwriteValue,
119119
bool zeroByteAfterPosted);
120120

121+
__global__ void mixedRegisteredAndStagedSendRecvKernel(
122+
P2pIbgdaTransportDevice* transport,
123+
IbgdaLocalBuffer sendBuffer,
124+
void* recvBuffer,
125+
std::size_t firstBytes,
126+
std::size_t secondBytes,
127+
std::size_t thirdBytes,
128+
std::size_t maxSignalBytes,
129+
bool send);
130+
121131
__global__ void fillTransportStagingKernel(
122132
P2pIbgdaTransportDevice* transport,
123133
bool sendStaging,

‎comms/prims/tests/MultipeerIbgdaTransportTest.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,21 @@ void testRegisteredSendRecv(
243243
uint8_t overwriteValue = 0,
244244
bool zeroByteAfterPosted = false);
245245

246+
/**
247+
* Test kernel: registered A, staged B, registered C on one send cursor.
248+
*/
249+
void testMixedRegisteredAndStagedSendRecv(
250+
P2pIbgdaTransportDevice* transport,
251+
const IbgdaLocalBuffer& sendBuffer,
252+
void* recvBuffer,
253+
std::size_t firstBytes,
254+
std::size_t secondBytes,
255+
std::size_t thirdBytes,
256+
std::size_t maxSignalBytes,
257+
bool send,
258+
int numBlocks,
259+
int blockSize);
260+
246261
/** Fill or verify a byte range in this channel's transport staging. */
247262
void testFillTransportStaging(
248263
P2pIbgdaTransportDevice* transport,

0 commit comments

Comments
 (0)