connect() APIs changed to return an instance instead of a shared_ptr (#680)

The key purpose is handling all mscclpp objects' memory internally by
hiding shared pointers from user APIs.
* `Connection` class is now a wrapper of `BaseConnection` class that is
equivalent to the previous `Connection` class
* `connect()` methods now return `Connection` instead of
`std::shared_ptr<Connection>`
* Removed `connectOnSetup()` method

Author: chhwang

apps/nccl/src/allgather.hpp

@@ -216,7 +216,7 @@ class AllgatherAlgo6 : public mscclpp::AlgorithmBuilder {
 
  private:
   bool disableChannelCache_;
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
   std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memorySemaphores_;
   const int nChannelsPerConnection_ = 35;
 
@@ -235,7 +235,7 @@ class AllgatherAlgo8 : public mscclpp::AlgorithmBuilder {
   mscclpp::Algorithm build() override;
 
  private:
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
 
   void initialize(std::shared_ptr<mscclpp::Communicator> comm,
                   std::unordered_map<std::string, std::shared_ptr<void>>& extras);

apps/nccl/src/allreduce.hpp

@@ -1137,7 +1137,7 @@ class AllreducePacket : public mscclpp::AlgorithmBuilder {
   size_t scratchBufferSize_;
   std::shared_ptr<char> scratchBuffer_;
   const int nSegmentsForScratchBuffer_ = 2;
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
 
   std::shared_ptr<uint32_t> deviceFlag7_;
   std::shared_ptr<uint32_t> deviceFlag28_;
@@ -1164,7 +1164,7 @@ class AllreduceNvls : public mscclpp::AlgorithmBuilder {
   uint32_t nSwitchChannels_;
   std::shared_ptr<mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel>> memoryChannelsDeviceHandle_;
   std::vector<mscclpp::BaseMemoryChannel> baseChannels_;
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
 };
 
 class AllreduceNvlsWithCopy : public mscclpp::AlgorithmBuilder {
@@ -1188,7 +1188,7 @@ class AllreduceNvlsWithCopy : public mscclpp::AlgorithmBuilder {
   uint32_t nSwitchChannels_;
   std::shared_ptr<mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel>> memoryChannelsDeviceHandle_;
   std::vector<mscclpp::BaseMemoryChannel> baseChannels_;
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
 };
 
 class Allreduce8 : public mscclpp::AlgorithmBuilder {
@@ -1209,7 +1209,7 @@ class Allreduce8 : public mscclpp::AlgorithmBuilder {
   size_t scratchBufferSize_;
   std::shared_ptr<mscclpp::Communicator> comm_;
   int nChannelsPerConnection_;
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
   std::shared_ptr<char> scratchBuffer_;
   std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> outputSemaphores_;
   std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> inputScratchSemaphores_;

apps/nccl/src/broadcast.hpp

@@ -171,7 +171,7 @@ class BroadcastAlgo6 : public mscclpp::AlgorithmBuilder {
                                                               void* output, size_t, ncclDataType_t);
   mscclpp::AlgorithmCtxKey generateBroadcastContextKey(const void*, void*, size_t, ncclDataType_t);
 
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
   size_t scratchMemSize_;
   std::shared_ptr<char> scratchBuffer_;
 };

apps/nccl/src/common.cu

@@ -20,41 +20,41 @@ std::vector<mscclpp::RegisteredMemory> setupRemoteMemories(std::shared_ptr<msccl
 }
 
 std::vector<mscclpp::MemoryChannel> setupMemoryChannels(
-    const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    const std::vector<mscclpp::Connection>& connections,
     const std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memorySemaphores,
     const std::vector<mscclpp::RegisteredMemory>& remoteMemories, mscclpp::RegisteredMemory localMemory,
     int nChannelsPerConnection) {
   std::vector<mscclpp::MemoryChannel> channels;
   size_t nConnections = connections.size();
   for (int idx = 0; idx < nChannelsPerConnection; ++idx) {
     for (size_t cid = 0; cid < nConnections; ++cid) {
-      if (connections[cid]->transport() == mscclpp::Transport::CudaIpc) {
+      if (connections[cid].transport() == mscclpp::Transport::CudaIpc) {
         channels.emplace_back(memorySemaphores[idx * nConnections + cid], remoteMemories[cid], localMemory, nullptr);
       }
     }
   }
   return channels;
 }
 
-std::vector<std::shared_ptr<mscclpp::Connection>> setupConnections(std::shared_ptr<mscclpp::Communicator> comm) {
-  std::vector<std::shared_future<std::shared_ptr<mscclpp::Connection>>> connectionFutures;
+std::vector<mscclpp::Connection> setupConnections(std::shared_ptr<mscclpp::Communicator> comm) {
+  std::vector<std::shared_future<mscclpp::Connection>> connectionFutures;
   for (int i = 0; i < comm->bootstrap()->getNranks(); i++) {
     if (i == comm->bootstrap()->getRank()) continue;
     connectionFutures.push_back(comm->connect(mscclpp::Transport::CudaIpc, i));
   }
-  std::vector<std::shared_ptr<mscclpp::Connection>> connections;
+  std::vector<mscclpp::Connection> connections;
   std::transform(connectionFutures.begin(), connectionFutures.end(), std::back_inserter(connections),
                  [](const auto& future) { return future.get(); });
   return connections;
 }
 
 std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> setupMemorySemaphores(
-    std::shared_ptr<mscclpp::Communicator> comm, const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    std::shared_ptr<mscclpp::Communicator> comm, const std::vector<mscclpp::Connection>& connections,
     int nChannelsPerConnection) {
   std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memorySemaphores;
   for (int idx = 0; idx < nChannelsPerConnection; ++idx) {
     for (size_t cid = 0; cid < connections.size(); ++cid) {
-      if (connections[cid]->transport() == mscclpp::Transport::CudaIpc) {
+      if (connections[cid].transport() == mscclpp::Transport::CudaIpc) {
         memorySemaphores.emplace_back(
             std::make_shared<mscclpp::MemoryDevice2DeviceSemaphore>(*(comm), connections[cid]));
       }
@@ -117,14 +117,14 @@ std::shared_ptr<mscclpp::DeviceHandle<mscclpp::SwitchChannel>> setupNvlsChannelD
 }
 
 std::vector<mscclpp::BaseMemoryChannel> setupBaseMemoryChannels(
-    const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    const std::vector<mscclpp::Connection>& connections,
     const std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memorySemaphores,
     int nChannelsPerConnection) {
   std::vector<mscclpp::BaseMemoryChannel> channels;
   size_t nConnections = connections.size();
   for (int idx = 0; idx < nChannelsPerConnection; ++idx) {
     for (size_t cid = 0; cid < nConnections; ++cid) {
-      if (connections[cid]->transport() == mscclpp::Transport::CudaIpc) {
+      if (connections[cid].transport() == mscclpp::Transport::CudaIpc) {
         channels.emplace_back(memorySemaphores[idx * nConnections + cid]);
       }
     }

apps/nccl/src/common.hpp

@@ -33,15 +33,15 @@ std::vector<mscclpp::RegisteredMemory> setupRemoteMemories(std::shared_ptr<msccl
                                                            mscclpp::RegisteredMemory localMemory);
 
 std::vector<mscclpp::MemoryChannel> setupMemoryChannels(
-    const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    const std::vector<mscclpp::Connection>& connections,
     const std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memorySemaphores,
     const std::vector<mscclpp::RegisteredMemory>& remoteMemories, mscclpp::RegisteredMemory localMemory,
     int nChannelsPerConnection);
 
-std::vector<std::shared_ptr<mscclpp::Connection>> setupConnections(std::shared_ptr<mscclpp::Communicator> comm);
+std::vector<mscclpp::Connection> setupConnections(std::shared_ptr<mscclpp::Communicator> comm);
 
 std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> setupMemorySemaphores(
-    std::shared_ptr<mscclpp::Communicator> comm, const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    std::shared_ptr<mscclpp::Communicator> comm, const std::vector<mscclpp::Connection>& connections,
     int nChannelsPerConnection);
 
 std::shared_ptr<mscclpp::DeviceHandle<mscclpp::MemoryChannel>> setupMemoryChannelDeviceHandles(
@@ -57,7 +57,7 @@ std::shared_ptr<mscclpp::DeviceHandle<mscclpp::SwitchChannel>> setupNvlsChannelD
     const std::vector<mscclpp::SwitchChannel>& nvlsChannels);
 
 std::vector<mscclpp::BaseMemoryChannel> setupBaseMemoryChannels(
-    const std::vector<std::shared_ptr<mscclpp::Connection>>& connections,
+    const std::vector<mscclpp::Connection>& connections,
     const std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>>& memorySemaphores,
     int nChannelsPerConnection);
 

docs/tutorials/01-basic-concepts.md

@@ -83,8 +83,8 @@ The connection is created by calling `connect` on the context object:
 
 ```cpp
 // From gpu_ping_pong.cu, lines 76 and 82
-std::shared_ptr<mscclpp::Connection> conn0 = ctx->connect(/*localEndpoint*/ ep0, /*remoteEndpoint*/ ep1);
-std::shared_ptr<mscclpp::Connection> conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
+mscclpp::Connection conn0 = ctx->connect(/*localEndpoint*/ ep0, /*remoteEndpoint*/ ep1);
+mscclpp::Connection conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
 ```
 
 The `localEndpoint` and `remoteEndpoint` parameters specify which endpoints are used for the connection. A connection is asymmetric by nature, meaning that we need to create one connection for each endpoint. In this case, `conn0` is created for `ep0` to communicate with `ep1`, and `conn1` is created for `ep1` to communicate with `ep0`.
@@ -101,7 +101,7 @@ sendToProcessB(serializedEp0);  // send serializedEp0 to Process B using any IPC
 mscclpp::Endpoint ep1 = ctx->createEndpoint({transport, {mscclpp::DeviceType::GPU, 1}});
 std::vector<char> serializedEp0 = recvFromProcessA();  // receive serializedEp0 from Process A
 mscclpp::Endpoint ep0 = mscclpp::Endpoint::deserialize(serializedEp0);
-std::shared_ptr<mscclpp::Connection> conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
+mscclpp::Connection conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
 ```
 
 ## SemaphoreStub and Semaphore

examples/customized-collective-algorithm/customized_allgather.cu

@@ -107,18 +107,18 @@ class AllgatherAlgoBuilder : public mscclpp::AlgorithmBuilder {
   }
 
  private:
-  std::vector<std::shared_ptr<mscclpp::Connection>> conns_;
+  std::vector<mscclpp::Connection> conns_;
   std::shared_ptr<mscclpp::ProxyService> proxyService_;
   int worldSize_;
 
   void initialize(std::shared_ptr<mscclpp::Communicator> comm) {
-    std::vector<std::shared_future<std::shared_ptr<mscclpp::Connection>>> connectionFutures;
+    std::vector<std::shared_future<mscclpp::Connection>> connectionFutures;
     worldSize_ = comm->bootstrap()->getNranks();
     for (int i = 0; i < worldSize_; i++) {
       if (i == comm->bootstrap()->getRank()) continue;
       connectionFutures.push_back(comm->connect(mscclpp::Transport::CudaIpc, i));
     }
-    std::vector<std::shared_ptr<mscclpp::Connection>> connections;
+    std::vector<mscclpp::Connection> connections;
     std::transform(connectionFutures.begin(), connectionFutures.end(), std::back_inserter(connections),
                    [](const auto& future) { return future.get(); });
     this->conns_ = std::move(connections);

examples/tutorials/01-basic-concepts/gpu_ping_pong.cu

@@ -73,13 +73,13 @@ int main() {
   log("GPU 0: Creating a connection and a semaphore stub ...");
 
   MSCCLPP_CUDATHROW(cudaSetDevice(0));
-  std::shared_ptr<mscclpp::Connection> conn0 = ctx->connect(/*localEndpoint*/ ep0, /*remoteEndpoint*/ ep1);
+  mscclpp::Connection conn0 = ctx->connect(/*localEndpoint*/ ep0, /*remoteEndpoint*/ ep1);
   mscclpp::SemaphoreStub semaStub0(conn0);
 
   log("GPU 1: Creating a connection and a semaphore stub ...");
 
   MSCCLPP_CUDATHROW(cudaSetDevice(1));
-  std::shared_ptr<mscclpp::Connection> conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
+  mscclpp::Connection conn1 = ctx->connect(/*localEndpoint*/ ep1, /*remoteEndpoint*/ ep0);
   mscclpp::SemaphoreStub semaStub1(conn1);
 
   log("GPU 0: Creating a semaphore and a memory channel ...");

include/mscclpp/core.hpp

@@ -425,6 +425,7 @@ struct EndpointConfig {
 
 class Context;
 class Connection;
+class BaseConnection;
 class RegisteredMemory;
 class SemaphoreStub;
 class Semaphore;
@@ -474,7 +475,7 @@ class Endpoint {
   std::shared_ptr<Impl> pimpl_;
 
   friend class Context;
-  friend class Connection;
+  friend class BaseConnection;
 };
 
 /// Context for communication. This provides a low-level interface for forming connections in use-cases
@@ -521,8 +522,8 @@ class Context : public std::enable_shared_from_this<Context> {
   ///
   /// @param localEndpoint The local endpoint.
   /// @param remoteEndpoint The remote endpoint.
-  /// @return A shared pointer to the connection.
-  std::shared_ptr<Connection> connect(const Endpoint& localEndpoint, const Endpoint& remoteEndpoint);
+  /// @return A connection object.
+  Connection connect(const Endpoint& localEndpoint, const Endpoint& remoteEndpoint);
 
  private:
   Context();
@@ -531,7 +532,7 @@ class Context : public std::enable_shared_from_this<Context> {
   std::unique_ptr<Impl> pimpl_;
 
   friend class Endpoint;
-  friend class Connection;
+  friend class BaseConnection;
   friend class RegisteredMemory;
   friend class SemaphoreStub;
 };
@@ -578,7 +579,7 @@ class RegisteredMemory {
   std::shared_ptr<Impl> pimpl_;
 
   friend class Context;
-  friend class Connection;
+  friend class BaseConnection;
   friend class SemaphoreStub;
   friend class Semaphore;
 };
@@ -587,12 +588,7 @@ class RegisteredMemory {
 class Connection {
  public:
   /// Constructor.
-  /// @param context The context associated with the connection.
-  /// @param localEndpoint The local endpoint of the connection.
-  Connection(std::shared_ptr<Context> context, const Endpoint& localEndpoint);
-
-  /// Destructor.
-  virtual ~Connection() = default;
+  Connection() = default;
 
   /// Write data from a source RegisteredMemory to a destination RegisteredMemory.
   ///
@@ -601,28 +597,27 @@ class Connection {
   /// @param src The source RegisteredMemory.
   /// @param srcOffset The offset in bytes from the start of the source RegisteredMemory.
   /// @param size The number of bytes to write.
-  virtual void write(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset,
-                     uint64_t size) = 0;
+  void write(RegisteredMemory dst, uint64_t dstOffset, RegisteredMemory src, uint64_t srcOffset, uint64_t size);
 
   /// Update an 8-byte value in a destination RegisteredMemory and synchronize the change with the remote process.
   ///
   /// @param dst The destination RegisteredMemory.
   /// @param dstOffset The offset in bytes from the start of the destination RegisteredMemory.
   /// @param src A pointer to the value to update.
   /// @param newValue The new value to write.
-  virtual void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue) = 0;
+  void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue);
 
   /// Flush any pending writes to the remote process.
   /// @param timeoutUsec Timeout in microseconds. Default: -1 (no timeout)
-  virtual void flush(int64_t timeoutUsec = -1) = 0;
+  void flush(int64_t timeoutUsec = -1);
 
   /// Get the transport used by the local process.
   /// @return The transport used by the local process.
-  virtual Transport transport() const = 0;
+  Transport transport() const;
 
   /// Get the transport used by the remote process.
   /// @return The transport used by the remote process.
-  virtual Transport remoteTransport() const = 0;
+  Transport remoteTransport() const;
 
   /// Get the context associated with this connection.
   /// @return A shared pointer to the context associated with this connection.
@@ -636,22 +631,23 @@ class Connection {
   /// @return The maximum number of write requests that can be queued.
   int getMaxWriteQueueSize() const;
 
- protected:
-  static const Endpoint::Impl& getImpl(const Endpoint& endpoint);
-  static const RegisteredMemory::Impl& getImpl(const RegisteredMemory& memory);
-  static Context::Impl& getImpl(Context& context);
+ private:
+  Connection(std::shared_ptr<BaseConnection> impl);
+  std::shared_ptr<BaseConnection> impl_;
 
-  std::shared_ptr<Context> context_;
-  Endpoint localEndpoint_;
-  int maxWriteQueueSize_;
+  friend class Context;
+  friend class Communicator;
+  friend class SemaphoreStub;
+  friend class Semaphore;
+  friend class ProxyService;
 };
 
 /// SemaphoreStub object only used for constructing Semaphore, not for direct use by the user.
 class SemaphoreStub {
  public:
   /// Constructor.
-  /// @param connection A shared pointer to the connection associated with this semaphore.
-  SemaphoreStub(std::shared_ptr<Connection> connection);
+  /// @param connection The connection associated with this semaphore.
+  SemaphoreStub(const Connection& connection);
 
   /// Get the memory associated with this semaphore.
   /// @return A reference to the registered memory for this semaphore.
@@ -686,8 +682,8 @@ class Semaphore {
   Semaphore(const SemaphoreStub& localStub, const SemaphoreStub& remoteStub);
 
   /// Get the connection associated with this semaphore.
-  /// @return A shared pointer to the connection.
-  std::shared_ptr<Connection> connection() const;
+  /// @return The connection.
+  Connection& connection();
 
   /// Get the local memory associated with this semaphore.
   /// @return A reference to the local registered memory.
@@ -873,34 +869,23 @@ class Communicator {
   /// @param localEndpoint The local endpoint.
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag to use for identifying the send and receive.
-  /// @return A future of shared pointer to the connection.
+  /// @return A future of the connection.
   ///
-  std::shared_future<std::shared_ptr<Connection>> connect(const Endpoint& localEndpoint, int remoteRank, int tag = 0);
+  std::shared_future<Connection> connect(const Endpoint& localEndpoint, int remoteRank, int tag = 0);
 
   /// Connect to a remote rank. Wrapper of `connect(localEndpoint, remoteRank, tag)`.
   /// @param localConfig The configuration for the local endpoint.
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag to use for identifying the send and receive.
-  /// @return A future of shared pointer to the connection.
-  std::shared_future<std::shared_ptr<Connection>> connect(const EndpointConfig& localConfig, int remoteRank,
-                                                          int tag = 0);
-
-  [[deprecated("Use connect(localConfig, remoteRank, tag) instead. This will be removed in a future release.")]] std::
-      shared_future<std::shared_ptr<Connection>>
-      connect(int remoteRank, int tag, EndpointConfig localConfig);
-
-  [[deprecated("Use connect() instead. This will be removed in a future release.")]] NonblockingFuture<
-      std::shared_ptr<Connection>>
-  connectOnSetup(int remoteRank, int tag, EndpointConfig localConfig) {
-    return connect(localConfig, remoteRank, tag);
-  }
+  /// @return A future of the connection.
+  std::shared_future<Connection> connect(const EndpointConfig& localConfig, int remoteRank, int tag = 0);
 
   /// Build a semaphore for cross-process synchronization.
   /// @param connection The connection associated with this semaphore.
   /// @param remoteRank The rank of the remote process.
   /// @param tag The tag to use for identifying the operation.
   /// @return A future of the built semaphore.
-  std::shared_future<Semaphore> buildSemaphore(std::shared_ptr<Connection> connection, int remoteRank, int tag = 0);
+  std::shared_future<Semaphore> buildSemaphore(const Connection& connection, int remoteRank, int tag = 0);
 
   /// Get the remote rank a connection is connected to.
   ///