Fix GpuStreamPool to be aware of the device ID of streams (#590)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: chhwang

include/mscclpp/gpu_utils.hpp

@@ -5,6 +5,7 @@
 #define MSCCLPP_GPU_UTILS_HPP_
 
 #include <memory>
+#include <unordered_map>
 
 #include "env.hpp"
 #include "errors.hpp"
@@ -47,7 +48,7 @@ struct AvoidCudaGraphCaptureGuard {
 struct CudaStreamWithFlags {
   /// Constructor without flags. This will not create any stream. set() can be called later to create a stream with
   /// specified flags.
-  CudaStreamWithFlags() : stream_(nullptr) {}
+  CudaStreamWithFlags();
 
   /// Constructor with flags. This will create a stream with the specified flags on the current device.
   /// @param flags The flags to create the stream with.
@@ -56,8 +57,8 @@ struct CudaStreamWithFlags {
   /// Destructor. This will destroy the stream if it was created.
   ~CudaStreamWithFlags();
 
-  /// Set the stream with the specified flags. If the stream was already created, it will raise an error with
-  /// ErrorCode::InvalidUsage.
+  /// Set the stream with the specified flags. The current device at the time of the construction will be used. If the
+  /// stream was already created, it will raise an error with ErrorCode::InvalidUsage.
   /// @param flags The flags to create the stream with.
   /// @throws Error if the stream was already created.
   void set(unsigned int flags);
@@ -68,7 +69,10 @@ struct CudaStreamWithFlags {
 
   operator cudaStream_t() const { return stream_; }
 
+  int deviceId() const { return deviceId_; }
+
   cudaStream_t stream_;
+  int deviceId_;
 };
 
 class GpuStreamPool;
@@ -89,6 +93,8 @@ class GpuStream {
 
   operator cudaStream_t() const { return stream_->stream_; }
 
+  int deviceId() const { return stream_->deviceId(); }
+
  private:
   friend class GpuStreamPool;
 
@@ -111,7 +117,7 @@ class GpuStreamPool {
 
  protected:
   friend class GpuStream;
-  std::vector<std::shared_ptr<CudaStreamWithFlags>> streams_;
+  std::unordered_map<int, std::vector<std::shared_ptr<CudaStreamWithFlags>>> streams_;
 };
 
 /// Get the singleton instance of GpuStreamPool.

src/gpu_utils.cc

@@ -13,7 +13,10 @@ AvoidCudaGraphCaptureGuard::AvoidCudaGraphCaptureGuard() : mode_(cudaStreamCaptu
 
 AvoidCudaGraphCaptureGuard::~AvoidCudaGraphCaptureGuard() { (void)cudaThreadExchangeStreamCaptureMode(&mode_); }
 
+CudaStreamWithFlags::CudaStreamWithFlags() : stream_(nullptr) { MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId_)); }
+
 CudaStreamWithFlags::CudaStreamWithFlags(unsigned int flags) {
+  MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId_));
   MSCCLPP_CUDATHROW(cudaStreamCreateWithFlags(&stream_, flags));
 }
 
@@ -23,22 +26,29 @@ CudaStreamWithFlags::~CudaStreamWithFlags() {
 
 void CudaStreamWithFlags::set(unsigned int flags) {
   if (!empty()) throw Error("CudaStreamWithFlags already set", ErrorCode::InvalidUsage);
+  int originalDeviceId;
+  MSCCLPP_CUDATHROW(cudaGetDevice(&originalDeviceId)); // Save the current device
+  MSCCLPP_CUDATHROW(cudaSetDevice(deviceId_));
   MSCCLPP_CUDATHROW(cudaStreamCreateWithFlags(&stream_, flags));
+  MSCCLPP_CUDATHROW(cudaSetDevice(originalDeviceId)); // Restore the original device
 }
 
 bool CudaStreamWithFlags::empty() const { return stream_ == nullptr; }
 
 GpuStream::GpuStream(std::shared_ptr<GpuStreamPool> pool, std::shared_ptr<CudaStreamWithFlags> stream)
     : pool_(pool), stream_(stream) {}
 
-GpuStream::~GpuStream() { pool_->streams_.push_back(stream_); }
+GpuStream::~GpuStream() { pool_->streams_[deviceId()].push_back(stream_); }
 
 GpuStreamPool::GpuStreamPool() {}
 
 GpuStream GpuStreamPool::getStream() {
-  if (!streams_.empty()) {
-    auto stream = streams_.back();
-    streams_.pop_back();
+  int deviceId;
+  MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId));
+  auto& streamVec = streams_[deviceId];
+  if (!streamVec.empty()) {
+    auto stream = streamVec.back();
+    streamVec.pop_back();
     return GpuStream(gpuStreamPool(), stream);
   }
   return GpuStream(gpuStreamPool(), std::make_shared<CudaStreamWithFlags>(cudaStreamNonBlocking));
@@ -47,7 +57,7 @@ GpuStream GpuStreamPool::getStream() {
 void GpuStreamPool::clear() { streams_.clear(); }
 
 // A global pool instance
-std::shared_ptr<GpuStreamPool> gGpuStreamPool_;
+static std::shared_ptr<GpuStreamPool> gGpuStreamPool_;
 
 std::shared_ptr<GpuStreamPool> gpuStreamPool() {
   if (!gGpuStreamPool_) {