Change Tensor::CopyFrom to a simple double dispatch (pytorch#14268)

Summary:
Pull Request resolved: pytorch#14268

Removes the need for Context in Tensor by doing simple dispatch for CopyBytes. It'd eventually be subsumed by Roy Li's changes of proper copy_ op, but before that is done, let's get a clear logic of how copies are implemented and clean up some craft in CopyFrom implementation.

Note, that with these changes, one can probably can get rid of Context::CopyFromCPU/CopyToCPU, but it's a matter for follow up diffs.

This diff doesn't change the API of Tensor yet, but relies on the fact that passing `Context` to CopyFrom makes copy async if the device is CUDA and doesn't have any effect otherwise (that's how Context methods are implemented).

This doesn't change semantics of copy async implementation - as before it blindly calls cudaMemcpyAsync which probably means that it can be misused if invoked separately outside of operator body. I'll leave it for the follow up copy_ unification.

For Extend() we always do async copy - it makes sense as it's an in-place device-device operation and only any further op would be observable.

Note: there are now three ways of invoking copy in C2 code - templated CopyBytes, virtual CopyFromCPU/etc, and double-dispatch free method here. Hopefully we can get rid of the second one.

Also, please advise whether it's c10-worthy :)

Reviewed By: ezyang

Differential Revision: D13117987

fbshipit-source-id: a6772d6dcf3effaf06717da3a656fc9873b310b5

Author: Dmytro Dzhulgakov

aten/src/ATen/core/TensorImpl.h

@@ -950,48 +950,35 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
       if (data_type_.copy()) {
         AT_ASSERTM(
             device_type() == ::at::DeviceType::CPU,
-            "In CopyFrom source and dest tensors must both be CPU for meta copy, "
-            "but dest tensor was ", device_type());
+            "In CopyFrom source and dest tensors must both be CPU for "
+            "non-POD copy, but dest tensor was ",
+            device_type());
         AT_ASSERTM(
             src.device_type() == ::at::DeviceType::CPU,
-            "In CopyFrom source and dest tensors must both be CPU for meta copy, "
-            "but src tensor was ", src.device_type());
+            "In CopyFrom source and dest tensors must both be CPU for "
+            "non-POD copy, but src tensor was ",
+            src.device_type());
         data_type_.copy()(src.data(), raw_mutable_data(data_type_), numel());
       } else {
         // The following copy uses the current (thread local) stream for copying
-        // and also takes the current GPU id previously set through CUDA API
-        // as we don't invoke SwitchToDevice anywhere
-        // TODO: this logic is overly complex and can be replaced with simple
-        // dispatch based on two device types
+        // and also takes the GPU id from the device() field passed in.
         //
-        // We'll need to use a non-CPU context to perform the copy if
-        // one of the context is not CPU since only non-CPU context
-        // knows how to copy between CPU and that context
-        if (src.device_type() != ::at::DeviceType::CPU || device_type() == ::at::DeviceType::CPU) {
-          if (!context) {
-            CreateContext(src.GetDevice())
-                ->CopyBytesToDevice(
-                    numel() * itemsize(),
-                    src.data(),
-                    raw_mutable_data(data_type_),
-                    device_type());
-          } else {
-            AT_ASSERTM(
-                context->device_type() == src.device_type(),
-                "Type for provided context does not match the type of source");
-            context->CopyBytesToDevice(
-                numel() * itemsize(), src.data(), raw_mutable_data(data_type_), device_type());
-          }
-        } else {
-          // In case source context is CPU, and target context is non-CPU
-          // We'll have to create a Context from target and perform the
-          // copy using that context
-          CreateContext(GetDevice())
-              ->CopyBytesFromCPU(
-                  numel() * itemsize(),
-                  src.data(),
-                  raw_mutable_data(data_type_));
-        }
+        // TODO: Potentially more enforcements are necessary to avoid accidental
+        // switch to sync copy if the currently set device is wrong.
+        //
+        // Specifically, we might need to switch to a different context device
+        // here explicitly to avoid relying on user synchronizing things
+        // properly.
+        //
+        // note: raw_mutable_data initializes device here
+        void* new_data = raw_mutable_data(data_type_);
+        at::CopyBytes(
+            numel() * itemsize(),
+            src.data(),
+            src.device(),
+            new_data,
+            device(),
+            context != nullptr);
       }
     }
   }
@@ -1037,8 +1024,29 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
     auto* newData = raw_mutable_data(data_type_);
     AT_ASSERTM(
         context != nullptr, "Context must be provided to Extend the tensor");
-    context->CopyItemsSameDevice(
-        data_type_, oldSize, oldData.get(), newData);
+    if (data_type_.copy()) {
+      AT_ASSERTM(
+          device_type() == ::at::DeviceType::CPU,
+          "non-POD types work only on CPU");
+      data_type_.copy()(oldData.get(), newData, oldSize);
+    } else {
+      // The following copy uses the current (thread local) stream for copying
+      // and also takes the GPU id from the device() field passed in.
+      //
+      // TODO: Potentially more enforcements are necessary to avoid accidental
+      // switch to sync copy if the currently set device is wrong.
+      //
+      // Specifically, we might need to switch to a different context device
+      // here explicitly to avoid relying on user synchronizing things
+      // properly.
+      at::CopyBytes(
+          oldSize * itemsize(),
+          oldData.get(),
+          device(),
+          newData,
+          device(),
+          true); // non-blocking
+    }
     reserved_ = true;
     sizes_ = newDims;
     numel_ = newNumel;

aten/src/ATen/core/context_base.cpp

@@ -1,5 +1,7 @@
 #include <ATen/core/context_base.h>
 
+#include <c10/util/Logging.h>
+
 namespace at {
 
 C10_DEFINE_TYPED_REGISTRY(
@@ -9,6 +11,49 @@ C10_DEFINE_TYPED_REGISTRY(
     std::unique_ptr,
     at::Device);
 
+// First dimension of the array is `bool async`: 0 is sync,
+// 1 is async (non-blocking)
+static CopyBytesFunction g_copy_bytes[2][COMPILE_TIME_MAX_DEVICE_TYPES]
+                                     [COMPILE_TIME_MAX_DEVICE_TYPES];
+
+_CopyBytesFunctionRegisterer::_CopyBytesFunctionRegisterer(
+    DeviceType fromType,
+    DeviceType toType,
+    CopyBytesFunction func_sync,
+    CopyBytesFunction func_async) {
+  auto from = static_cast<int>(fromType);
+  auto to = static_cast<int>(toType);
+  if (!func_async) {
+    // default to the sync function
+    func_async = func_sync;
+  }
+  CHECK(
+      g_copy_bytes[0][from][to] == nullptr &&
+      g_copy_bytes[1][from][to] == nullptr)
+      << "Duplicate registration for device type pair "
+      << c10::DeviceTypeName(fromType) << ", " << c10::DeviceTypeName(toType);
+  g_copy_bytes[0][from][to] = func_sync;
+  g_copy_bytes[1][from][to] = func_async;
+}
+
+void CopyBytes(
+    size_t nbytes,
+    const void* src,
+    Device src_device,
+    void* dst,
+    Device dst_device,
+    bool async) {
+  auto ptr = g_copy_bytes[async ? 1 : 0][static_cast<int>(src_device.type())]
+                         [static_cast<int>(dst_device.type())];
+  CAFFE_ENFORCE(
+      ptr,
+      "No function found for copying from ",
+      c10::DeviceTypeName(src_device.type()),
+      " to ",
+      c10::DeviceTypeName(dst_device.type()));
+  ptr(nbytes, src, src_device, dst, dst_device);
+}
+
 } // namespace at
 
 namespace caffe2 {

aten/src/ATen/core/context_base.h

@@ -63,25 +63,6 @@ class CAFFE2_API BaseContext {
 
   virtual void CopyBytesToCPU(size_t nbytes, const void* src, void* dst) = 0;
 
-  virtual void CopyBytesToDevice(
-      size_t nbytes,
-      const void* src,
-      void* dst,
-      DeviceType type) {
-    if (type == DeviceType::CPU) {
-      CopyBytesToCPU(nbytes, src, dst);
-    } else if (type == device_type()) {
-      CopyBytesSameDevice(nbytes, src, dst);
-    } else {
-      AT_ERROR(
-          "CopyBytesToDevice can only copy to CPU or between same "
-          "device. Can't copy from: ",
-          device_type(),
-          " to",
-          type);
-    }
-  }
-
   template <typename T>
   inline void CopySameDevice(size_t n, const T* src, T* dst) {
     static_assert(
@@ -175,9 +156,41 @@ inline std::unique_ptr<at::BaseContext> CreateContext(
 
 } // namespace at
 
+// TODO: move it to a separate file in c10 if possible
+namespace at {
+
+using CopyBytesFunction = void (*)(
+    size_t nbytes,
+    const void* src,
+    Device src_device,
+    void* dst,
+    Device dst_device);
+
+struct CAFFE2_API _CopyBytesFunctionRegisterer {
+  _CopyBytesFunctionRegisterer(
+      DeviceType from,
+      DeviceType to,
+      CopyBytesFunction func_sync,
+      CopyBytesFunction func_async = nullptr);
+};
+
+#define REGISTER_COPY_BYTES_FUNCTION(from, to, ...)           \
+  namespace {                                                 \
+  static _CopyBytesFunctionRegisterer C10_ANONYMOUS_VARIABLE( \
+      g_copy_function)(from, to, __VA_ARGS__);                \
+  }
+
+CAFFE2_API void CopyBytes(
+    size_t nbytes,
+    const void* src,
+    Device src_device,
+    void* dst,
+    Device dst_device,
+    bool async);
+} // namespace at
+
 namespace caffe2 {
 
 using at::BaseContext;
 using at::CreateContext;
-
 } // namespace caffe2

c10/DeviceType.h

@@ -29,6 +29,9 @@ enum class DeviceType : int16_t {
   ONLY_FOR_TEST = 20901, // This device type is only for test.
 };
 
+// define explicit int constant
+constexpr int COMPILE_TIME_MAX_DEVICE_TYPES =
+    static_cast<int>(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES);
 C10_API std::string DeviceTypeName(
     DeviceType d,
     bool lower_case = false);

caffe2/core/blob_serialization.cc

@@ -222,6 +222,8 @@ void TensorSerializer::Serialize(
   const TensorProto::DataType data_type = TypeMetaToDataType(input.dtype());
   proto.set_data_type(data_type);
   StoreDeviceDetail(input, &proto);
+  // TODO: use DeviceGuard here instead of context and employ explicit sync
+  // copy
   auto uniq_ptr = CreateContext(input.GetDevice());
   // A lot of copypaste is error prone. Should we create a macro for this?
   switch (data_type) {

caffe2/core/context.cc

@@ -5,10 +5,6 @@
 #include <process.h>
 #endif
 
-namespace at {
-
-REGISTER_CONTEXT(DeviceType::CPU, caffe2::CPUContext);
-} // namespace at
 namespace caffe2 {
 
 uint32_t RandomNumberSeed() {
@@ -28,4 +24,41 @@ uint32_t RandomNumberSeed() {
       kPrime2 * tv_sec + kPrime3 * tv_usec;
 }
 
+namespace {
+inline void CopyBytesImpl(size_t nbytes, const void* src, void* dst) {
+  if (nbytes == 0) {
+    return;
+  }
+  CAFFE_ENFORCE(src);
+  CAFFE_ENFORCE(dst);
+  memcpy(dst, src, nbytes);
+}
+
+void CopyBytesWrapper(
+    size_t nbytes,
+    const void* src,
+    Device src_device,
+    void* dst,
+    Device dst_device) {
+  CopyBytesImpl(nbytes, src, dst);
+}
+} // namespace
+
+void CPUContext::CopyBytesSameDevice(
+    size_t nbytes,
+    const void* src,
+    void* dst) {
+  CopyBytesImpl(nbytes, src, dst);
+}
+
 } // namespace caffe2
+
+namespace at {
+
+REGISTER_CONTEXT(DeviceType::CPU, caffe2::CPUContext);
+
+REGISTER_COPY_BYTES_FUNCTION(
+    DeviceType::CPU,
+    DeviceType::CPU,
+    caffe2::CopyBytesWrapper);
+} // namespace at

caffe2/core/context.h

@@ -79,14 +79,7 @@ class CAFFE2_API CPUContext final : public BaseContext {
     return GetCPUAllocator()->allocate(nbytes);
   }
 
-  void CopyBytesSameDevice(size_t nbytes, const void* src, void* dst) override {
-    if (nbytes == 0) {
-      return;
-    }
-    CAFFE_ENFORCE(src);
-    CAFFE_ENFORCE(dst);
-    memcpy(dst, src, nbytes);
-  }
+  void CopyBytesSameDevice(size_t nbytes, const void* src, void* dst) override;
 
   void CopyBytesFromCPU(size_t nbytes, const void* src, void* dst) override {
     CopyBytesSameDevice(nbytes, src, dst);