Move CUDA-related stuff of TP agent to separate file (#59377)

lw · facebook-github-bot · commit 5e5ca0682bb9 · 2021-06-15T03:28:38.000-07:00
Summary: Pull Request resolved: #59377 This PR demonstrates that now the CUDA parts of the TensorPipe agent just "plug on top" of the CPU-only parts. Thus ideally the CPU-only parts could go in libtorch while the CUDA-only parts could go in libtorch_cuda. Unfortunately we can't do that just yet, because the TensorPipe agent depends on c10d (for its Store and its ProcessGroup), which lives in libtorch_python. ghstack-source-id: 131326168 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D28796429 fbshipit-source-id: 41b2eb8400c0da282f3750a4eea21ad83ee4a175
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -367,6 +367,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
           "${TORCH_SRC_DIR}/csrc/distributed/rpc/macros.h"
           "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.cpp"
           "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.h"
+          "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_cuda.cpp"
           "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp"
           "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h"
           )
diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
@@ -718,6 +718,7 @@ libtorch_python_distributed_sources = libtorch_python_distributed_core_sources +
     "torch/csrc/distributed/rpc/python_rpc_handler.cpp",
     "torch/csrc/distributed/rpc/request_callback_impl.cpp",
     "torch/csrc/distributed/rpc/tensorpipe_agent.cpp",
+    "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp",
     "torch/csrc/distributed/rpc/tensorpipe_utils.cpp",
     "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp",
     "torch/csrc/distributed/rpc/testing/init.cpp",
diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp
@@ -136,9 +136,9 @@ std::vector<c10::Device> getDevicesOfTensors(
   }
   std::vector<c10::Device> devices;
   devices.reserve(deviceCount);
-  for (c10::DeviceIndex idx = 0; idx < indexBitset.size(); idx++) {
+  for (const auto idx : c10::irange(indexBitset.size())) {
     if (indexBitset[idx]) {
-      devices.emplace_back(impl->type(), idx);
+      devices.emplace_back(impl->type(), static_cast<c10::DeviceIndex>(idx));
     }
   }
   return devices;
@@ -195,36 +195,6 @@ const std::string& TensorPipeAgent::guessAddress() {
 
 namespace {
 
-// These priorities instruct TensorPipe on which transport/channel to pick
-// during handshake. Higher priorities will take precedence over lower ones.
-// The transport with lowest priority will be the one used to bootstrap pipes.
-
-constexpr int64_t kShmTransportPriority = 200;
-constexpr int64_t kIbvTransportPriority = 100;
-// The UV transport just uses TCP and should work everywhere, thus keep it last.
-constexpr int64_t kUvTransportPriority = 0;
-
-constexpr int64_t kCmaChannelPriority = 1200;
-constexpr int64_t kMultiplexedUvChannelPriority = 1100;
-// The basic channel reuses a transport as a channel, and is thus our fallback.
-constexpr int64_t kBasicChannelPriority = 1000;
-
-// CPU channel have higher priority than CUDA channels, since the latter might
-// handle CPU-to-CPU transfers, but will always be less efficient than their
-// CPU-only counterparts.
-#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL && defined(USE_CUDA_NOT_ROCM)
-constexpr int64_t kCudaIpcChannelPriority = 300;
-#endif
-
-#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM)
-constexpr int64_t kCudaGdrChannelPriority = 200;
-#endif
-
-#ifdef USE_CUDA_NOT_ROCM
-constexpr int64_t kCudaXthChannelPriority = 400;
-constexpr int64_t kCudaBasicChannelPriority = 0;
-#endif
-
 std::unique_ptr<TransportRegistration> makeUvTransport() {
   auto context = tensorpipe::transport::uv::create();
   std::string address = TensorPipeAgent::guessAddress();
@@ -319,7 +289,7 @@ constexpr static int kNumUvThreads = 16;
 std::unique_ptr<ChannelRegistration> makeMultiplexedUvChannel() {
   std::vector<std::shared_ptr<tensorpipe::transport::Context>> contexts;
   std::vector<std::shared_ptr<tensorpipe::transport::Listener>> listeners;
-  for (const auto laneIdx : c10::irange(kNumUvThreads)) {
+  for (const auto laneIdx C10_UNUSED : c10::irange(kNumUvThreads)) {
     auto context = tensorpipe::transport::uv::create();
     std::string address = TensorPipeAgent::guessAddress();
     contexts.push_back(std::move(context));
@@ -343,69 +313,6 @@ C10_REGISTER_CREATOR(
     mpt_uv,
     makeMultiplexedUvChannel);
 
-#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL && defined(USE_CUDA_NOT_ROCM)
-
-std::unique_ptr<ChannelRegistration> makeCudaIpcChannel() {
-  auto context = tensorpipe::channel::cuda_ipc::create();
-  return std::make_unique<ChannelRegistration>(
-      ChannelRegistration{std::move(context), kCudaIpcChannelPriority});
-}
-
-// The cuda_ipc channels use cudaMemcpy to transmit CUDA tensor across processes
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_ipc, makeCudaIpcChannel);
-
-#endif
-
-#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM)
-
-std::unique_ptr<ChannelRegistration> makeCudaGdrChannel() {
-  auto context = tensorpipe::channel::cuda_gdr::create();
-  return std::make_unique<ChannelRegistration>(
-      ChannelRegistration{std::move(context), kCudaGdrChannelPriority});
-}
-
-// The cuda_gdr channel sends CUDA memory over InfiniBand using GPUDirect RDMA.
-// It directly registers the user-provided tensor with libibverbs, an operation
-// which is expensive the first time, but it then caches the registration in
-// order to amortize the cost and get low latency for subsequent transfers. A
-// ready-to-send/ready-to-receive handshake is still needed before the transfer
-// in order to ensure readiness and to agree on the device indices and thus the
-// queue pair to use. It automatically pairs each GPU to the "closest" NIC if
-// there are multiple of them (closest = longest prefix match in PCI tree).
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_gdr, makeCudaGdrChannel);
-
-#endif
-
-#ifdef USE_CUDA_NOT_ROCM
-
-std::unique_ptr<ChannelRegistration> makeCudaXthChannel() {
-  auto context = tensorpipe::channel::cuda_xth::create();
-  return std::make_unique<ChannelRegistration>(
-      ChannelRegistration{std::move(context), kCudaXthChannelPriority});
-}
-
-// The cuda_xth channel supports same-process GPU-to-GPU comm
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_xth, makeCudaXthChannel);
-
-std::unique_ptr<ChannelRegistration> makeCudaBasicChannel() {
-  auto context = tensorpipe::channel::cuda_basic::create(
-      tensorpipe::channel::basic::create());
-  return std::make_unique<ChannelRegistration>(
-      ChannelRegistration{std::move(context), kCudaBasicChannelPriority});
-}
-
-// The cuda_basic is the fallback channel for GPU-to-GPU comm
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-C10_REGISTER_CREATOR(
-    TensorPipeChannelRegistry,
-    cuda_basic,
-    makeCudaBasicChannel);
-
-#endif
-
 } // namespace
 
 //////////////////////////  MetricsTracker  /////////////////////////////////
diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h
@@ -17,12 +17,6 @@
 
 namespace tensorpipe {
 
-struct CpuBuffer;
-
-#ifdef USE_CUDA_NOT_ROCM
-struct CudaBuffer;
-#endif
-
 class Context;
 class Error;
 class Listener;
@@ -43,6 +37,28 @@ namespace torch {
 namespace distributed {
 namespace rpc {
 
+// These priorities instruct TensorPipe on which transport/channel to pick
+// during handshake. Higher priorities will take precedence over lower ones.
+// The transport with lowest priority will be the one used to bootstrap pipes.
+
+constexpr int64_t kShmTransportPriority = 200;
+constexpr int64_t kIbvTransportPriority = 100;
+// The UV transport just uses TCP and should work everywhere, thus keep it last.
+constexpr int64_t kUvTransportPriority = 0;
+
+constexpr int64_t kCmaChannelPriority = 1200;
+constexpr int64_t kMultiplexedUvChannelPriority = 1100;
+// The basic channel reuses a transport as a channel, and is thus our fallback.
+constexpr int64_t kBasicChannelPriority = 1000;
+
+// CPU channel have higher priority than CUDA channels, since the latter might
+// handle CPU-to-CPU transfers, but will always be less efficient than their
+// CPU-only counterparts.
+constexpr int64_t kCudaIpcChannelPriority = 300;
+constexpr int64_t kCudaGdrChannelPriority = 200;
+constexpr int64_t kCudaXthChannelPriority = 400;
+constexpr int64_t kCudaBasicChannelPriority = 0;
+
 using steady_clock_time_point =
     std::chrono::time_point<std::chrono::steady_clock>;
 
diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp
@@ -0,0 +1,136 @@
+#include <torch/csrc/distributed/rpc/macros.h>
+#include <torch/csrc/distributed/rpc/tensorpipe_agent.h>
+#include <torch/csrc/distributed/rpc/tensorpipe_utils.h>
+
+#if defined(USE_TENSORPIPE) && defined(USE_CUDA_NOT_ROCM)
+
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#include <tensorpipe/tensorpipe.h>
+
+namespace torch {
+namespace distributed {
+namespace rpc {
+namespace {
+
+#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL
+
+std::unique_ptr<ChannelRegistration> makeCudaIpcChannel() {
+  auto context = tensorpipe::channel::cuda_ipc::create();
+  return std::make_unique<ChannelRegistration>(
+      ChannelRegistration{std::move(context), kCudaIpcChannelPriority});
+}
+
+// The cuda_ipc channels use cudaMemcpy to transmit CUDA tensor across processes
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_ipc, makeCudaIpcChannel);
+
+#endif
+
+#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL
+
+std::unique_ptr<ChannelRegistration> makeCudaGdrChannel() {
+  auto context = tensorpipe::channel::cuda_gdr::create();
+  return std::make_unique<ChannelRegistration>(
+      ChannelRegistration{std::move(context), kCudaGdrChannelPriority});
+}
+
+// The cuda_gdr channel sends CUDA memory over InfiniBand using GPUDirect RDMA.
+// It directly registers the user-provided tensor with libibverbs, an operation
+// which is expensive the first time, but it then caches the registration in
+// order to amortize the cost and get low latency for subsequent transfers. A
+// ready-to-send/ready-to-receive handshake is still needed before the transfer
+// in order to ensure readiness and to agree on the device indices and thus the
+// queue pair to use. It automatically pairs each GPU to the "closest" NIC if
+// there are multiple of them (closest = longest prefix match in PCI tree).
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_gdr, makeCudaGdrChannel);
+
+#endif
+
+std::unique_ptr<ChannelRegistration> makeCudaXthChannel() {
+  auto context = tensorpipe::channel::cuda_xth::create();
+  return std::make_unique<ChannelRegistration>(
+      ChannelRegistration{std::move(context), kCudaXthChannelPriority});
+}
+
+// The cuda_xth channel supports same-process GPU-to-GPU comm
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_xth, makeCudaXthChannel);
+
+std::unique_ptr<ChannelRegistration> makeCudaBasicChannel() {
+  auto context = tensorpipe::channel::cuda_basic::create(
+      tensorpipe::channel::basic::create());
+  return std::make_unique<ChannelRegistration>(
+      ChannelRegistration{std::move(context), kCudaBasicChannelPriority});
+}
+
+// The cuda_basic is the fallback channel for GPU-to-GPU comm
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+C10_REGISTER_CREATOR(
+    TensorPipeChannelRegistry,
+    cuda_basic,
+    makeCudaBasicChannel);
+
+class TensorpipeCudaConverter : public TensorpipeDeviceTypeConverter {
+ public:
+  c10::optional<std::vector<char>> prepareTensorForSending(
+      const c10::Storage& storage,
+      const std::vector<c10::Stream>& streams,
+      tensorpipe::Message& message) const override {
+    auto stream =
+        at::cuda::CUDAStream(getStreamForDevice(streams, storage.device()));
+    // record tensor data ptrs on TensorPipe streams, so that the tensors
+    // won't be destructed before TensorPipe finishing sending them.
+    c10::cuda::CUDACachingAllocator::recordStream(storage.data_ptr(), stream);
+
+    tensorpipe::CudaBuffer buffer;
+    buffer.ptr = storage.data<char>();
+    buffer.stream = stream.stream();
+
+    tensorpipe::Message::Tensor tensor;
+    tensor.buffer = buffer;
+    tensor.length = storage.nbytes();
+
+    message.tensors.push_back(std::move(tensor));
+
+    return c10::nullopt;
+  }
+
+  at::DataPtr allocateTensorForReceiving(
+      int deviceIndex,
+      size_t length,
+      const std::vector<c10::Stream>& streams,
+      tensorpipe::Allocation& allocation) const override {
+    c10::Device device(c10::kCUDA, deviceIndex);
+    at::cuda::CUDAStream stream(getStreamForDevice(streams, device));
+    // CUDACachingAllocator will call recordStream accordingly on the current
+    // stream.
+    at::cuda::CUDAStreamGuard guard(stream);
+    at::DataPtr dataPtr =
+        c10::cuda::CUDACachingAllocator::get()->allocate(length);
+
+    tensorpipe::CudaBuffer buffer;
+    buffer.ptr = dataPtr.get();
+    buffer.stream = stream.stream();
+
+    tensorpipe::Allocation::Tensor tensor;
+    tensor.buffer = buffer;
+
+    allocation.tensors.push_back(tensor);
+
+    return dataPtr;
+  }
+};
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CUDA, TensorpipeCudaConverter);
+
+} // namespace
+} // namespace rpc
+} // namespace distributed
+} // namespace torch
+
+#endif
diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp
@@ -3,12 +3,7 @@
 
 #ifdef USE_TENSORPIPE
 
-#ifdef USE_CUDA_NOT_ROCM
-#include <c10/core/DeviceGuard.h>
-#include <c10/cuda/CUDACachingAllocator.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <c10/util/irange.h>
-#endif
 
 #include <tensorpipe/tensorpipe.h>
 
@@ -98,62 +93,6 @@ class TensorpipeCpuConverter : public TensorpipeDeviceTypeConverter {
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CPU, TensorpipeCpuConverter);
 
-#ifdef USE_CUDA_NOT_ROCM
-class TensorpipeCudaConverter : public TensorpipeDeviceTypeConverter {
- public:
-  c10::optional<std::vector<char>> prepareTensorForSending(
-      const c10::Storage& storage,
-      const std::vector<c10::Stream>& streams,
-      tensorpipe::Message& message) const override {
-    auto stream =
-        at::cuda::CUDAStream(getStreamForDevice(streams, storage.device()));
-    // record tensor data ptrs on TensorPipe streams, so that the tensors
-    // won't be destructed before TensorPipe finishing sending them.
-    c10::cuda::CUDACachingAllocator::recordStream(storage.data_ptr(), stream);
-
-    tensorpipe::CudaBuffer buffer;
-    buffer.ptr = storage.data<char>();
-    buffer.stream = stream.stream();
-
-    tensorpipe::Message::Tensor tensor;
-    tensor.buffer = buffer;
-    tensor.length = storage.nbytes();
-
-    message.tensors.push_back(std::move(tensor));
-
-    return c10::nullopt;
-  }
-
-  at::DataPtr allocateTensorForReceiving(
-      int deviceIndex,
-      size_t length,
-      const std::vector<c10::Stream>& streams,
-      tensorpipe::Allocation& allocation) const override {
-    c10::Device device(c10::kCUDA, deviceIndex);
-    at::cuda::CUDAStream stream(getStreamForDevice(streams, device));
-    // CUDACachingAllocator will call recordStream accordingly on the current
-    // stream.
-    at::cuda::CUDAStreamGuard guard(stream);
-    at::DataPtr dataPtr =
-        c10::cuda::CUDACachingAllocator::get()->allocate(length);
-
-    tensorpipe::CudaBuffer buffer;
-    buffer.ptr = dataPtr.get();
-    buffer.stream = stream.stream();
-
-    tensorpipe::Allocation::Tensor tensor;
-    tensor.buffer = buffer;
-
-    allocation.tensors.push_back(tensor);
-
-    return dataPtr;
-  }
-};
-
-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CUDA, TensorpipeCudaConverter);
-#endif
-
 c10::DeviceType convertDeviceType(const std::string& tpDeviceType) {
   if (tpDeviceType == tensorpipe::kCpuDeviceType) {
     return c10::kCPU;
diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h
