Skip to content

Commit 59e15c8

Browse files
Binyang2014Binyang2014
committed
address comments
1 parent 4b9e709 commit 59e15c8

11 files changed

Lines changed: 34 additions & 71 deletions

‎include/mscclpp/core.hpp‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -331,11 +331,6 @@ class RegisteredMemory {
331331
/// @return The size of the memory block.
332332
size_t size();
333333

334-
/// Get the pitch of the memory block.
335-
///
336-
/// @return The pitch of the memory block.
337-
size_t pitch();
338-
339334
/// Get the rank of the process that owns the memory block.
340335
///
341336
/// @return The rank of the process that owns the memory block.
@@ -384,12 +379,14 @@ class Connection {
384379
///
385380
/// @param dst The destination @ref RegisteredMemory.
386381
/// @param dstOffset The offset in bytes from the start of the destination @ref RegisteredMemory.
382+
/// @param dstPitch The pitch of the destination @ref RegisteredMemory in bytes.
387383
/// @param src The source @ref RegisteredMemory.
388384
/// @param srcOffset The offset in bytes from the start of the source @ref RegisteredMemory.
385+
/// @param srcPitch The pitch of the source @ref RegisteredMemory in bytes.
389386
/// @param width The width of the 2D region to write in bytes.
390387
/// @param height The height of the 2D region.
391-
virtual void write2D(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset,
392-
uint64_t width, uint64_t height) = 0;
388+
virtual void write2D(RegisteredMemory dst, uint64_t dstOffset, uint64_t dstPitch, RegisteredMemory src,
389+
uint64_t srcOffset, uint64_t srcPitch, uint64_t width, uint64_t height) = 0;
393390
/// Update a 8-byte value in a destination @ref RegisteredMemory and synchronize the change with the remote process.
394391
///
395392
/// @param dst The destination @ref RegisteredMemory.

‎include/mscclpp/proxy_channel.hpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <mscclpp/fifo.hpp>
99
#include <mscclpp/proxy.hpp>
1010
#include <mscclpp/semaphore.hpp>
11+
#include <unordered_map>
1112

1213
namespace mscclpp {
1314

@@ -40,6 +41,8 @@ class ProxyService : public BaseProxyService {
4041
/// @return The ID of the semaphore.
4142
SemaphoreId addSemaphore(std::shared_ptr<Connection> connection);
4243

44+
void addPitch(SemaphoreId id, std::pair<uint64_t, uint64_t> pitch);
45+
4346
/// Register a memory region with the proxy service.
4447
/// @param memory The memory region to register.
4548
/// @return The ID of the memory region.
@@ -65,6 +68,7 @@ class ProxyService : public BaseProxyService {
6568
Communicator& communicator_;
6669
std::vector<std::shared_ptr<Host2DeviceSemaphore>> semaphores_;
6770
std::vector<RegisteredMemory> memories_;
71+
std::unordered_map<SemaphoreId, std::pair<uint64_t, uint64_t>> pitches_;
6872
Proxy proxy_;
6973
int deviceNumaNode;
7074

‎src/communicator.cc‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,12 +57,6 @@ MSCCLPP_API_CPP RegisteredMemory Communicator::registerMemory(void* ptr, size_t
5757
std::make_shared<RegisteredMemory::Impl>(ptr, size, pimpl->bootstrap_->getRank(), transports, *pimpl));
5858
}
5959

60-
MSCCLPP_API_CPP RegisteredMemory Communicator::registerMemory(void* ptr, size_t size, size_t pitchSize,
61-
TransportFlags transports) {
62-
return RegisteredMemory(
63-
std::make_shared<RegisteredMemory::Impl>(ptr, size, pitchSize, pimpl->bootstrap_->getRank(), transports, *pimpl));
64-
}
65-
6660
struct MemorySender : public Setuppable {
6761
MemorySender(RegisteredMemory memory, int remoteRank, int tag)
6862
: memory_(memory), remoteRank_(remoteRank), tag_(tag) {}

‎src/connection.cc‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,18 +58,18 @@ void CudaIpcConnection::write(RegisteredMemory dst, uint64_t dstOffset, Register
5858
// npkitCollectEntryEvent(conn, NPKIT_EVENT_DMA_SEND_DATA_ENTRY, (uint32_t)size);
5959
}
6060

61-
void CudaIpcConnection::write2D(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset,
62-
uint64_t width, uint64_t height) {
61+
void CudaIpcConnection::write2D(RegisteredMemory dst, uint64_t dstOffset, uint64_t dstPitch, RegisteredMemory src,
62+
uint64_t srcOffset, uint64_t srcPitch, uint64_t width, uint64_t height) {
6363
validateTransport(dst, remoteTransport());
6464
validateTransport(src, transport());
6565

6666
char* dstPtr = (char*)dst.data();
6767
char* srcPtr = (char*)src.data();
6868

69-
MSCCLPP_CUDATHROW(cudaMemcpy2DAsync(dstPtr + dstOffset, dst.pitch(), srcPtr + srcOffset, src.pitch(), width, height,
69+
MSCCLPP_CUDATHROW(cudaMemcpy2DAsync(dstPtr + dstOffset, dstPitch, srcPtr + srcOffset, srcPitch, width, height,
7070
cudaMemcpyDeviceToDevice, stream_));
7171
INFO(MSCCLPP_P2P, "CudaIpcConnection write: from %p to %p, width %lu height %lu dstPitch %lu srcPitch %lu",
72-
srcPtr + srcOffset, dstPtr + dstOffset, width, height, dst.pitch(), src.pitch());
72+
srcPtr + srcOffset, dstPtr + dstOffset, width, height, dstPitch, srcPitch);
7373
}
7474

7575
void CudaIpcConnection::updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue) {
@@ -141,7 +141,8 @@ void IBConnection::write(RegisteredMemory dst, uint64_t dstOffset, RegisteredMem
141141
// npkitCollectEntryEvent(conn, NPKIT_EVENT_IB_SEND_DATA_ENTRY, (uint32_t)size);
142142
}
143143

144-
void IBConnection::write2D(RegisteredMemory, uint64_t, RegisteredMemory, uint64_t, uint64_t, uint64_t) {
144+
void IBConnection::write2D(RegisteredMemory, uint64_t, uint64_t, RegisteredMemory, uint64_t, uint64_t, uint64_t,
145+
uint64_t) {
145146
throw Error("write2D is not supported", ErrorCode::InvalidUsage);
146147
}
147148

‎src/include/connection.hpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,8 @@ class CudaIpcConnection : public ConnectionBase {
4444

4545
void write(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset,
4646
uint64_t size) override;
47-
void write2D(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset, uint64_t width,
48-
uint64_t height) override;
47+
void write2D(RegisteredMemory dst, uint64_t dstOffset, uint64_t dstPitch, RegisteredMemory src, uint64_t srcOffset,
48+
uint64_t srcPitch, uint64_t width, uint64_t height) override;
4949
void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue) override;
5050

5151
void flush() override;
@@ -69,8 +69,8 @@ class IBConnection : public ConnectionBase {
6969

7070
void write(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset,
7171
uint64_t size) override;
72-
void write2D(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset, uint64_t width,
73-
uint64_t height) override;
72+
void write2D(RegisteredMemory dst, uint64_t dstOffset, uint64_t dstPitch, RegisteredMemory src, uint64_t srcOffset,
73+
uint64_t srcPitch, uint64_t width, uint64_t height) override;
7474
void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue) override;
7575

7676
void flush() override;

‎src/include/registered_memory.hpp‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,11 @@ struct TransportInfo {
3434
struct RegisteredMemory::Impl {
3535
void* data;
3636
size_t size;
37-
size_t pitch; // for 2D
3837
int rank;
3938
uint64_t hostHash;
4039
TransportFlags transports;
4140
std::vector<TransportInfo> transportInfos;
4241

43-
Impl(void* data, size_t size, size_t pitch, int rank, TransportFlags transports, Communicator::Impl& commImpl);
4442
Impl(void* data, size_t size, int rank, TransportFlags transports, Communicator::Impl& commImpl);
4543
Impl(const std::vector<char>& data);
4644

‎src/proxy_channel.cc‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ MSCCLPP_API_CPP SemaphoreId ProxyService::addSemaphore(std::shared_ptr<Connectio
2929
return semaphores_.size() - 1;
3030
}
3131

32+
MSCCLPP_API_CPP void ProxyService::addPitch(SemaphoreId id, std::pair<uint64_t, uint64_t> pitch) {
33+
pitches_[id] = pitch;
34+
}
35+
3236
MSCCLPP_API_CPP MemoryId ProxyService::addMemory(RegisteredMemory memory) {
3337
memories_.push_back(memory);
3438
return memories_.size() - 1;
@@ -63,8 +67,9 @@ ProxyHandlerResult ProxyService::handleTrigger(ProxyTrigger triggerRaw) {
6367
RegisteredMemory& dst = memories_[trigger->fields.dstMemoryId];
6468
RegisteredMemory& src = memories_[trigger->fields.srcMemoryId];
6569
if (trigger->fields2D.multiDimensionFlag) {
66-
semaphore->connection()->write2D(dst, trigger->fields.dstOffset, src, trigger->fields.srcOffset,
67-
trigger->fields2D.width, trigger->fields2D.height);
70+
std::pair<uint64_t, uint64_t>& pitch = pitches_[trigger->fields.chanId];
71+
semaphore->connection()->write2D(dst, trigger->fields.dstOffset, pitch.first, src, trigger->fields.srcOffset,
72+
pitch.second, trigger->fields2D.width, trigger->fields2D.height);
6873
} else {
6974
semaphore->connection()->write(dst, trigger->fields.dstOffset, src, trigger->fields.srcOffset,
7075
trigger->fields.size);

‎src/registered_memory.cc‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,7 @@
1515
namespace mscclpp {
1616

1717
RegisteredMemory::Impl::Impl(void* data, size_t size, int rank, TransportFlags transports, Communicator::Impl& commImpl)
18-
: Impl(data, size, size, rank, transports, commImpl) {}
19-
20-
RegisteredMemory::Impl::Impl(void* data, size_t size, size_t pitch, int rank, TransportFlags transports,
21-
Communicator::Impl& commImpl)
22-
: data(data),
23-
size(size),
24-
pitch(pitch),
25-
rank(rank),
26-
hostHash(commImpl.rankToHash_.at(rank)),
27-
transports(transports) {
18+
: data(data), size(size), rank(rank), hostHash(commImpl.rankToHash_.at(rank)), transports(transports) {
2819
if (transports.has(Transport::CudaIpc)) {
2920
TransportInfo transportInfo;
3021
transportInfo.transport = Transport::CudaIpc;
@@ -69,16 +60,13 @@ MSCCLPP_API_CPP void* RegisteredMemory::data() { return pimpl->data; }
6960

7061
MSCCLPP_API_CPP size_t RegisteredMemory::size() { return pimpl->size; }
7162

72-
MSCCLPP_API_CPP size_t RegisteredMemory::pitch() { return pimpl->pitch; }
73-
7463
MSCCLPP_API_CPP int RegisteredMemory::rank() { return pimpl->rank; }
7564

7665
MSCCLPP_API_CPP TransportFlags RegisteredMemory::transports() { return pimpl->transports; }
7766

7867
MSCCLPP_API_CPP std::vector<char> RegisteredMemory::serialize() {
7968
std::vector<char> result;
8069
std::copy_n(reinterpret_cast<char*>(&pimpl->size), sizeof(pimpl->size), std::back_inserter(result));
81-
std::copy_n(reinterpret_cast<char*>(&pimpl->pitch), sizeof(pimpl->pitch), std::back_inserter(result));
8270
std::copy_n(reinterpret_cast<char*>(&pimpl->rank), sizeof(pimpl->rank), std::back_inserter(result));
8371
std::copy_n(reinterpret_cast<char*>(&pimpl->hostHash), sizeof(pimpl->hostHash), std::back_inserter(result));
8472
std::copy_n(reinterpret_cast<char*>(&pimpl->transports), sizeof(pimpl->transports), std::back_inserter(result));
@@ -111,8 +99,6 @@ RegisteredMemory::Impl::Impl(const std::vector<char>& serialization) {
11199
auto it = serialization.begin();
112100
std::copy_n(it, sizeof(this->size), reinterpret_cast<char*>(&this->size));
113101
it += sizeof(this->size);
114-
std::copy_n(it, sizeof(this->pitch), reinterpret_cast<char*>(&this->pitch));
115-
it += sizeof(this->pitch);
116102
std::copy_n(it, sizeof(this->rank), reinterpret_cast<char*>(&this->rank));
117103
it += sizeof(this->rank);
118104
std::copy_n(it, sizeof(this->hostHash), reinterpret_cast<char*>(&this->hostHash));

‎test/mp_unit/communicator_tests.cu‎

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -61,16 +61,7 @@ void CommunicatorTestBase::registerMemoryPairs(void* buff, size_t buffSize, mscc
6161
const std::vector<int>& remoteRanks,
6262
mscclpp::RegisteredMemory& localMemory,
6363
std::unordered_map<int, mscclpp::RegisteredMemory>& remoteMemories) {
64-
registerMemoryPairs(buff, buffSize, buffSize, transport, tag, remoteRanks, localMemory, remoteMemories);
65-
}
66-
67-
// Register a local memory with pitch and receive corresponding remote memories
68-
void CommunicatorTestBase::registerMemoryPairs(void* buff, size_t buffSize, size_t pitchSize,
69-
mscclpp::TransportFlags transport, int tag,
70-
const std::vector<int>& remoteRanks,
71-
mscclpp::RegisteredMemory& localMemory,
72-
std::unordered_map<int, mscclpp::RegisteredMemory>& remoteMemories) {
73-
localMemory = communicator->registerMemory(buff, buffSize, pitchSize, transport);
64+
localMemory = communicator->registerMemory(buff, buffSize, transport);
7465
std::unordered_map<int, mscclpp::NonblockingFuture<mscclpp::RegisteredMemory>> futureRemoteMemories;
7566
for (int remoteRank : remoteRanks) {
7667
if (remoteRank != communicator->bootstrap()->getRank()) {
@@ -105,9 +96,7 @@ void CommunicatorTest::SetUp() {
10596

10697
devicePtr.resize(numBuffers);
10798
localMemory.resize(numBuffers);
108-
local2DMemory.resize(numBuffers);
10999
remoteMemory.resize(numBuffers);
110-
remote2DMemory.resize(numBuffers);
111100

112101
std::vector<int> remoteRanks;
113102
for (int i = 0; i < gEnv->worldSize; i++) {
@@ -121,18 +110,11 @@ void CommunicatorTest::SetUp() {
121110
registerMemoryPairs(devicePtr[n].get(), deviceBufferSize, mscclpp::Transport::CudaIpc | ibTransport, 0, remoteRanks,
122111
localMemory[n], remoteMemory[n]);
123112
}
124-
125-
for (size_t n = 0; n < numBuffers; n++) {
126-
registerMemoryPairs(devicePtr[n].get(), deviceBufferSize, deviceBufferPitchSize, mscclpp::Transport::CudaIpc, 0,
127-
remoteRanks, local2DMemory[n], remote2DMemory[n]);
128-
}
129113
}
130114

131115
void CommunicatorTest::TearDown() {
132116
remoteMemory.clear();
133-
remote2DMemory.clear();
134117
localMemory.clear();
135-
local2DMemory.clear();
136118
devicePtr.clear();
137119
CommunicatorTestBase::TearDown();
138120
}
@@ -168,8 +150,9 @@ void CommunicatorTest::writeTileToRemote(size_t rowIndex, size_t colIndex, size_
168150
for (int i = 0; i < gEnv->worldSize; i++) {
169151
if (i != gEnv->rank) {
170152
auto& conn = connections.at(i);
171-
auto& peerMemory = remote2DMemory[n].at(i);
172-
conn->write2D(peerMemory, offset, local2DMemory[n], offset, width * sizeof(int), height);
153+
auto& peerMemory = remoteMemory[n].at(i);
154+
conn->write2D(peerMemory, offset, deviceBufferPitchSize, localMemory[n], offset, deviceBufferPitchSize,
155+
width * sizeof(int), height);
173156
conn->flush();
174157
}
175158
}

‎test/mp_unit/mp_unit_tests.hpp‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -99,10 +99,6 @@ class CommunicatorTestBase : public MultiProcessTest {
9999
void registerMemoryPairs(void* buff, size_t buffSize, mscclpp::TransportFlags transport, int tag,
100100
const std::vector<int>& remoteRanks, mscclpp::RegisteredMemory& localMemory,
101101
std::unordered_map<int, mscclpp::RegisteredMemory>& remoteMemories);
102-
// Register a local memory with pitch and receive corresponding remote memories
103-
void registerMemoryPairs(void* buff, size_t buffSize, size_t pitch, mscclpp::TransportFlags transport, int tag,
104-
const std::vector<int>& remoteRanks, mscclpp::RegisteredMemory& localMemory,
105-
std::unordered_map<int, mscclpp::RegisteredMemory>& remoteMemories);
106102
// Register a local memory an receive one corresponding remote memory
107103
void registerMemoryPair(void* buff, size_t buffSize, mscclpp::TransportFlags transport, int tag, int remoteRank,
108104
mscclpp::RegisteredMemory& localMemory, mscclpp::RegisteredMemory& remoteMemory);
@@ -128,9 +124,7 @@ class CommunicatorTest : public CommunicatorTestBase {
128124
const int deviceBufferPitchSize = 512;
129125
std::vector<std::shared_ptr<int>> devicePtr;
130126
std::vector<mscclpp::RegisteredMemory> localMemory;
131-
std::vector<mscclpp::RegisteredMemory> local2DMemory;
132127
std::vector<std::unordered_map<int, mscclpp::RegisteredMemory>> remoteMemory;
133-
std::vector<std::unordered_map<int, mscclpp::RegisteredMemory>> remote2DMemory;
134128
};
135129

136130
class ProxyChannelOneToOneTest : public CommunicatorTestBase {

0 commit comments

Comments
 (0)