chore: import upstream snapshot with attribution

paddle/phi/kernels/memcpy_kernel.cc

@@ -0,0 +1,174 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/memcpy_kernel.h"
+
+#include <vector>
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/stream.h"
+
+namespace phi {
+
+static constexpr size_t WAIT_THRESHOLD = 64 * 1024;
+
+template <typename Context>
+PADDLE_API void MemcpyH2DKernel(const Context& dev_ctx,
+                                const DenseTensor& x,
+                                int dst_place_type,
+                                DenseTensor* out) {
+  if (!x.initialized()) {
+    out->set_meta(x.meta());
+    return;
+  }
+
+  PADDLE_ENFORCE(
+      (dst_place_type >= 0 && dst_place_type <= 3) || (dst_place_type == 6),
+      errors::OutOfRange("dst_place_type only supports 0-3 or 6, but got: %d",
+                         dst_place_type));
+  Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+}
+
+template <typename Context>
+PADDLE_API void MemcpyD2HKernel(const Context& dev_ctx,
+                                const DenseTensor& x,
+                                int dst_place_type,
+                                DenseTensor* out) {
+  switch (dst_place_type) {
+    case 0:
+      Copy(dev_ctx, x, CPUPlace(), false, out);
+      // NOTE(copy from Aurelius84): host <-> device memory copies of a memory
+      // block of 64 KB or less are asynchronous. See
+      // https://forums.developer.nvidia.com/t/host-device-memory-copies-up-to-64-kb-are-asynchronous/17907
+      if (x.memory_size() <= WAIT_THRESHOLD) {
+        dev_ctx.Wait();
+      }
+      break;
+
+    case 1:
+      Copy(dev_ctx, x, GPUPinnedPlace(), false, out);
+      // Copy use async copy for GPUPinnedPlace
+      dev_ctx.Wait();
+      break;
+
+    default:
+      PADDLE_THROW(errors::InvalidArgument(
+          "Argument 'dst_place_type' only support 0-1, but got: %d",
+          dst_place_type));
+      break;
+  }
+}
+
+template <typename Context>
+void MemcpyD2HMultiIOKernel(const Context& dev_ctx,
+                            const TensorArray& array,
+                            int dst_place_type,
+                            TensorArray* out_array) {
+  PADDLE_ENFORCE_NOT_NULL(
+      out_array,
+      errors::PreconditionNotMet("output tensor_array should not be nullptr"));
+
+  out_array->clear();
+  out_array->resize(array.size());
+  for (size_t i = 0; i < array.size(); i++) {
+    const auto& x = array[i];
+    MemcpyD2HKernel<Context>(dev_ctx, x, dst_place_type, &(out_array->at(i)));
+  }
+}
+
+template <typename Context>
+void MemcpyKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  int dst_place_type,
+                  DenseTensor* out) {
+  if (!x.IsInitialized()) {
+    return;
+  }
+  PADDLE_ENFORCE_GE(
+      dst_place_type,
+      0,
+      errors::OutOfRange("dst_place_type only support 0-4, but got: %d",
+                         dst_place_type));
+  PADDLE_ENFORCE_LE(
+      dst_place_type,
+      4,
+      errors::OutOfRange("dst_place_type only support 0-4, but got: %d",
+                         dst_place_type));
+  switch (dst_place_type) {
+    case 0: /* CPUPlace */
+      dev_ctx.HostAlloc(out, out->dtype());
+      Copy(dev_ctx, x, CPUPlace(), true, out);
+      break;
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    case 1: /* CUDAPlace */
+      dev_ctx.Alloc(out, x.dtype());
+      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+      break;
+    case 2: /* CUDAPinnedPlace */
+      dev_ctx.Alloc(out, x.dtype(), 0, true);
+      Copy(dev_ctx, x, GPUPinnedPlace(), false, out);
+      break;
+#elif defined(PADDLE_WITH_XPU)
+    case 1:  // XPUPlace
+      dev_ctx.Alloc(out, x.dtype());
+      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+      break;
+    case 3:  // XPUPlace
+      dev_ctx.Alloc(out, x.dtype());
+      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+      break;
+    case 5: /* XPUPinnedPlace */
+      dev_ctx.Alloc(out, x.dtype(), 0, true);
+      Copy(dev_ctx, x, XPUPinnedPlace(), false, out);
+      break;
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+    case 1:  // CustomPlace
+      dev_ctx.Alloc(out, x.dtype());
+      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+      break;
+    case 4:  // CustomPlace
+      dev_ctx.Alloc(out, x.dtype());
+      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+      break;
+
+#endif
+    default:
+      PADDLE_THROW(errors::Unimplemented(
+          "memcpy dst_place_type: %d is not supported yet.", dst_place_type));
+      break;
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(memcpy_h2d,
+                                         ALL_LAYOUT,
+                                         phi::MemcpyH2DKernel) {}
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(memcpy_d2h,
+                                         ALL_LAYOUT,
+                                         phi::MemcpyD2HKernel) {
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+}
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(memcpy_d2h_multi_io,
+                                         ALL_LAYOUT,
+                                         phi::MemcpyD2HMultiIOKernel) {
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+}
+PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(memcpy,
+                                         ALL_LAYOUT,
+                                         phi::MemcpyKernel) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}