Skip to content

Add registered-source send to IBGDA transport (#3411) #3411

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
snarayankh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add registered-source send to IBGDA transport (#3411) #3411

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions comms/prims/P2pIbTransportBuildContractTest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,14 @@

_PROGRESS_ONLY_FUNCTIONS = (
"init_send_progress",
"init_registered_send_progress",
"init_recv_progress",
"progress_send_once",
"progress_registered_send_once",
"progress_registered_send_drain_once",
"poll_recv_data_ready",
"progress_recv_once",
"send_registered",
"store_progress_state",
"make_progress_geometry",
"active_payload_offset",
Expand Down
175 changes: 174 additions & 1 deletion comms/prims/tests/MultipeerIbgdaTransportTest.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
#include <folly/logging/xlog.h>
#include <array>

#include <chrono>
#include <memory>
#include <string>
#include <tuple>
#include <vector>

#ifdef __HIP_PLATFORM_AMD__
Expand Down Expand Up @@ -1791,6 +1791,179 @@ TEST_P(
}
}

TEST_F(
MultipeerIbgdaTransportTestFixture,
RegisteredSendTailsDrainAndStagingIsolation) {
if (numRanks != 2) {
GTEST_SKIP() << "Skipping test: requires exactly 2 ranks, got " << numRanks;
}
if (!test::supportsProgressSendRecv()) {
GTEST_SKIP() << "registered-source send is not supported for this build";
}

constexpr std::size_t perChannelSize = 64 * 1024;
constexpr int pipelineDepth = 2;
constexpr std::size_t pipelineChunk = perChannelSize / pipelineDepth;
constexpr std::size_t maxSignalBytes = 4 * 1024;
constexpr std::size_t simpleProtocolDataBytes = 16;
constexpr int numBlocks = 1;
constexpr int blockSize = 128;
constexpr uint8_t stagingPoison = 0xA5;
constexpr std::size_t zeroByteAfterPostedIndex = 2;
const int peerRank = globalRank == 0 ? 1 : 0;
const std::array<std::size_t, 8> sizes{
0,
1,
7,
simpleProtocolDataBytes - 1,
simpleProtocolDataBytes,
simpleProtocolDataBytes + 1,
pipelineChunk - 1,
pipelineChunk + 1,
};

std::unique_ptr<MultipeerIbgdaTransport> transport;
try {
MultipeerIbgdaTransportConfig config{
.cudaDevice = localRank,
.perChannelSize = perChannelSize,
.max_num_channels = numBlocks,
.pipelineDepth = pipelineDepth,
};
auto bootstrap = std::make_shared<meta::comms::MpiBootstrap>();
transport = std::make_unique<MultipeerIbgdaTransport>(
globalRank, numRanks, bootstrap, config);
transport->exchange();
} catch (const std::exception& e) {
GTEST_SKIP() << "IBGDA transport not available: " << e.what();
}

auto* peerTransport = transport->getP2pTransportDevice(peerRank);
DeviceBuffer errorCountBuffer(sizeof(int));
DeviceBuffer observationBuffer(sizeof(test::RegisteredSendObservation));
auto* errorCount = static_cast<int*>(errorCountBuffer.get());
auto* observation =
static_cast<test::RegisteredSendObservation*>(observationBuffer.get());

test::testFillTransportStaging(
peerTransport,
globalRank == 0,
0,
perChannelSize,
stagingPoison,
numBlocks,
blockSize);
CUDACHECK_TEST(cudaDeviceSynchronize());

for (std::size_t index = 0; index < sizes.size(); ++index) {
const std::size_t nbytes = sizes[index];
const std::size_t allocationBytes = nbytes == 0 ? 1 : nbytes;
const uint8_t pattern = static_cast<uint8_t>(0x20 + index * 13);
DeviceBuffer sendBuffer(allocationBytes);
DeviceBuffer recvBuffer(allocationBytes);
IbgdaLocalBuffer registeredSource{};
if (globalRank == 0) {
if (nbytes > 0) {
registeredSource = transport->registerBuffer(sendBuffer.get(), nbytes);
}
test::fillBufferWithPattern(
sendBuffer.get(), nbytes, pattern, numBlocks, blockSize);
CUDACHECK_TEST(
cudaMemset(observation, 0, sizeof(test::RegisteredSendObservation)));
} else {
CUDACHECK_TEST(cudaMemset(recvBuffer.get(), 0, allocationBytes));
}
CUDACHECK_TEST(cudaDeviceSynchronize());
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));

test::testRegisteredSendRecv(
peerTransport,
registeredSource,
recvBuffer.get(),
nbytes,
maxSignalBytes,
globalRank == 0,
numBlocks,
blockSize,
globalRank == 0 ? observation : nullptr,
index == 4,
globalRank == 0 && nbytes > 0,
static_cast<uint8_t>(pattern ^ 0xFF),
index == zeroByteAfterPostedIndex);
CUDACHECK_TEST(cudaDeviceSynchronize());
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));

if (globalRank == 0) {
test::RegisteredSendObservation hostObservation{};
CUDACHECK_TEST(cudaMemcpy(
&hostObservation,
observation,
sizeof(hostObservation),
cudaMemcpyDeviceToHost));
if (index == 4) {
EXPECT_EQ(hostObservation.postedCount, 0);
} else if (index == zeroByteAfterPostedIndex) {
EXPECT_EQ(hostObservation.postedCount, 2);
} else {
EXPECT_EQ(hostObservation.postedCount, 1);
}
EXPECT_EQ(hostObservation.drainedCount, 1);
} else {
CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
test::verifyBufferPattern(
recvBuffer.get(), nbytes, pattern, errorCount, numBlocks, blockSize);
CUDACHECK_TEST(cudaDeviceSynchronize());
int hostErrors = 0;
CUDACHECK_TEST(cudaMemcpy(
&hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
EXPECT_EQ(hostErrors, 0) << "registered send corrupted size " << nbytes;
}

if (index == 1 && globalRank == 1) {
CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
test::testVerifyTransportStaging(
peerTransport,
false,
1,
simpleProtocolDataBytes - 1,
stagingPoison,
errorCount,
numBlocks,
blockSize);
CUDACHECK_TEST(cudaDeviceSynchronize());
int hostErrors = 0;
CUDACHECK_TEST(cudaMemcpy(
&hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
EXPECT_EQ(hostErrors, 0)
<< "registered send wrote rounded tail bytes into recv staging";
}
if (globalRank == 0 && nbytes > 0) {
transport->deregisterBuffer(sendBuffer.get());
}
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
}

if (globalRank == 0) {
CUDACHECK_TEST(cudaMemset(errorCount, 0, sizeof(int)));
test::testVerifyTransportStaging(
peerTransport,
true,
0,
perChannelSize,
stagingPoison,
errorCount,
numBlocks,
blockSize);
CUDACHECK_TEST(cudaDeviceSynchronize());
int hostErrors = 0;
CUDACHECK_TEST(cudaMemcpy(
&hostErrors, errorCount, sizeof(hostErrors), cudaMemcpyDeviceToHost));
EXPECT_EQ(hostErrors, 0)
<< "registered send modified transport send staging";
}
MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
}

// =============================================================================
// Sustained chunked send/recv - repro for the GB200 per-channel DATA_READY
// deadlock (two NICs atomic-FA the same flag at maxGroups>=8). Streams a large
Expand Down
Loading
Loading