chore: import upstream snapshot with attribution

paddle/phi/kernels/coalesce_tensor_kernel.cc

@@ -0,0 +1,325 @@
+// 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/coalesce_tensor_kernel.h"
+
+#include <sstream>
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/device_memory_alignment.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/kernels/funcs/math_function_impl.h"
+#endif
+
+namespace phi {
+
+template <typename Context>
+struct FillConstantVisitor {
+  FillConstantVisitor(const Context &dev_ctx,
+                      DenseTensor *tensor,
+                      const float value)
+      : dev_ctx_(dev_ctx), tensor_(tensor), value_(value) {}
+
+  template <typename T>
+  void apply(typename std::enable_if<std::is_same<T, int8_t>::value ||
+                                     std::is_same<T, int16_t>::value>::type * =
+                 nullptr) const {
+    PADDLE_THROW(
+        errors::InvalidArgument("Not support data type for set_constant attr"));
+  }
+
+  template <typename T>
+  void apply(typename std::enable_if<!(std::is_same<T, int8_t>::value ||
+                                       std::is_same<T, int16_t>::value)>::type
+                 * = nullptr) const {
+    phi::funcs::SetConstant<Context, T> set_constant;
+    set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
+  }
+
+  const Context &dev_ctx_;
+  DenseTensor *tensor_;
+  float value_;
+};
+
+void GetMemSizeAndDtype(const std::vector<const DenseTensor *> &lod_tensors,
+                        size_t *numel,
+                        const size_t &size_of_dtype,
+                        const phi::Place &place,
+                        const bool use_align = true,
+                        const int align_size = -1) {
+  *numel = 0;
+  std::stringstream ss;
+  ss << "alloc_space_for_vars: ";
+  for (size_t i = 0; i < lod_tensors.size(); ++i) {
+    auto size = lod_tensors[i]->numel();
+    PADDLE_ENFORCE_GT(size,
+                      0,
+                      errors::InvalidArgument(
+                          "The number of `%d`-th tensor's elements is 0.", i));
+    auto len = use_align
+                   ? phi::Alignment(static_cast<size_t>(size) * size_of_dtype,
+                                    place,
+                                    align_size) /
+                         size_of_dtype
+                   : static_cast<size_t>(size);
+    const void *ptr =
+        lod_tensors[i]->initialized() ? lod_tensors[i]->data() : nullptr;
+    VLOG(4) << size << " " << len;
+    ss << "input(" << i << "-th tensor) dim:(" << lod_tensors[i]->dims() << ") "
+       << " address:" << ptr << " len: " << len << ", ";
+    *numel += len;
+  }
+  VLOG(10) << ss.str();
+}
+
+template <typename T, typename Context>
+void CoalesceTensorKernel(const Context &dev_ctx,
+                          const std::vector<const DenseTensor *> &input,
+                          DataType dtype,
+                          bool copy_data,
+                          bool set_constant,
+                          bool persist_output,
+                          float constant,
+                          bool use_align,
+                          int align_size,
+                          int size_of_dtype,
+                          const std::vector<int64_t> &concated_shapes,
+                          const std::vector<int64_t> &concated_ranks,
+                          std::vector<DenseTensor *> output,
+                          DenseTensor *fused_output) {
+  PADDLE_ENFORCE_GT(
+      input.size(),
+      static_cast<size_t>(0),
+      errors::InvalidArgument("The CoalesceTensor operator has no input."));
+  PADDLE_ENFORCE_EQ(input.size(),
+                    output.size(),
+                    errors::InvalidArgument(
+                        "The number of CoalesceTensor operator's input and "
+                        "output is not match, "
+                        "input number is %u, output number is %u.",
+                        input.size(),
+                        output.size()));
+
+  // Input & Output check: only support DenseTensor
+  bool has_not_init_in_vars = false;
+  for (size_t i = 0; i < input.size(); ++i) {
+    PADDLE_ENFORCE_NOT_NULL(
+        input[i],
+        errors::InvalidArgument("The %d-th input tensor cannot be nullptr.",
+                                i));
+    PADDLE_ENFORCE_NOT_NULL(
+        output[i],
+        errors::InvalidArgument("The %d-th output tensor cannot be nullptr.",
+                                i));
+    if (!input[i]->initialized()) {
+      has_not_init_in_vars = true;
+    }
+  }
+
+  if (has_not_init_in_vars) {
+    PADDLE_ENFORCE_EQ(
+        concated_ranks.size(),
+        output.size(),
+        errors::InvalidArgument("The attribute(concated_ranks) length must be "
+                                "equal to the output tensor number."));
+    int64_t accumulated_ranks = 0;
+    for (size_t i = 0; i < input.size(); ++i) {
+      phi::DDim dims(concated_shapes.data() + accumulated_ranks,
+                     static_cast<int>(concated_ranks[i]));
+      if (!input[i]->initialized()) {
+        PADDLE_ENFORCE_EQ(
+            input[i],
+            output[i],
+            errors::InvalidArgument(
+                "The %d-th output tensor and %d-th input tensor when the "
+                "%d-th input tensor is not initialized.",
+                i,
+                i,
+                i));
+        output[i]->Resize(dims);
+      } else {
+        PADDLE_ENFORCE_EQ(input[i]->dims(),
+                          dims,
+                          errors::InvalidArgument(
+                              "The %d-th input tensor shape does not match the "
+                              "attribute(concated_shapes) and "
+                              "attribute(concated_ranks).",
+                              i));
+      }
+      accumulated_ranks += concated_ranks[i];
+      PADDLE_ENFORCE_LE(
+          accumulated_ranks,
+          concated_shapes.size(),
+          errors::InvalidArgument("The attribute(concated_shapes) and "
+                                  "attribute(concated_ranks) do not match."));
+    }
+    PADDLE_ENFORCE_EQ(
+        accumulated_ranks,
+        concated_shapes.size(),
+        errors::InvalidArgument("The attribute(concated_shapes) and "
+                                "attribute(concated_ranks) do not match."));
+  }
+
+  // Init the output as input
+  for (size_t i = 0; i < input.size(); ++i) {
+    output[i]->Resize(input[i]->dims());
+  }
+
+  // Get numel and dtype
+  size_t numel = 0;
+
+  if (size_of_dtype == -1) {
+    size_of_dtype = static_cast<int>(phi::SizeOf(dtype));
+  }
+  GetMemSizeAndDtype(
+      input, &numel, size_of_dtype, dev_ctx.GetPlace(), use_align, align_size);
+
+  // Alloc the continuous space
+  void *fused_tensor_ptr = dev_ctx.Alloc(
+      &fused_output->Resize({static_cast<int64_t>(numel)}), dtype);
+  VLOG(10) << "Fused tensor addr " << fused_tensor_ptr;
+
+  // Init the continuous space
+  size_t offset = 0;
+  if (copy_data) {
+    for (auto item : input) {
+      size_t len = static_cast<size_t>(item->numel());
+      auto sub_tensor = fused_output->Slice(
+          static_cast<int64_t>(offset),
+          static_cast<int64_t>(offset) + static_cast<int64_t>(len));
+      phi::Copy(dev_ctx, *item, dev_ctx.GetPlace(), false, &sub_tensor);
+
+      offset += use_align
+                    ? phi::Alignment(
+                          len * size_of_dtype, dev_ctx.GetPlace(), align_size) /
+                          size_of_dtype
+                    : len;
+    }
+  } else if (set_constant) {
+    phi::VisitDataType(
+        dtype, FillConstantVisitor<Context>(dev_ctx, fused_output, constant));
+  } else if (persist_output) {
+    for (auto &item : output) {
+      size_t len = static_cast<size_t>(item->numel());
+      auto sub_tensor = fused_output->Slice(
+          static_cast<int64_t>(offset),
+          static_cast<int64_t>(offset) + static_cast<int64_t>(len));
+      // some var may not persistable, or persistable var may not init
+      if (item->initialized()) {
+        phi::Copy(dev_ctx, *item, dev_ctx.GetPlace(), false, &sub_tensor);
+      }
+      offset += use_align
+                    ? phi::Alignment(
+                          len * size_of_dtype, dev_ctx.GetPlace(), align_size) /
+                          size_of_dtype
+                    : len;
+    }
+  }
+
+  // Make the outputs point to the continuous space.
+  offset = 0;
+  std::stringstream ss;
+  ss << "alloc_space_for_vars: ";
+
+  for (size_t i = 0; i < output.size(); ++i) {
+    size_t len = static_cast<size_t>(output[i]->numel());
+    auto dim = output[i]->dims();
+    VLOG(4) << len << " " << dim << " " << offset;
+    output[i]
+        ->ShareDataWith(fused_output->Slice(
+            static_cast<int64_t>(offset),
+            static_cast<int64_t>(offset) + static_cast<int64_t>(len)))
+        .Resize(dim);
+    len = use_align ? phi::Alignment(
+                          len * size_of_dtype, dev_ctx.GetPlace(), align_size) /
+                          size_of_dtype
+                    : len;
+    ss << "output(" << i << "-th tensor) dim:(" << dim << ")"
+       << " address: " << output[i]->data() << " len: " << len << ", ";
+    offset += len;
+  }
+  PADDLE_ENFORCE_EQ((int64_t)offset,
+                    fused_output->numel(),
+                    errors::InvalidArgument(
+                        "The alloc_space_for_vars's offset: %s is unequal with "
+                        "fused_output's numel: %s.",
+                        offset,
+                        fused_output->numel()));
+  VLOG(10) << ss.str();
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(coalesce_tensor,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::CoalesceTensorKernel,
+                   int,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+}
+
+#ifdef PADDLE_WITH_CUDA
+PD_REGISTER_KERNEL(coalesce_tensor,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::CoalesceTensorKernel,
+                   phi::float16,
+                   phi::bfloat16,
+                   int,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+}
+#endif
+
+#ifdef PADDLE_WITH_HIP
+PD_REGISTER_KERNEL(coalesce_tensor,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::CoalesceTensorKernel,
+                   phi::float16,
+                   phi::bfloat16,
+                   int,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+}
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(coalesce_tensor,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::CoalesceTensorKernel,
+                   phi::float16,
+                   phi::bfloat16,
+                   int,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
+}
+#endif