chore: import upstream snapshot with attribution

paddle/phi/kernels/impl/broadcast_tensors_kernel_impl.h

@@ -0,0 +1,129 @@
+// 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.
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/broadcast_tensors_kernel.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+#define SWITCH_OUT_RANK_CASE(n)                                            \
+  case n: {                                                                \
+    ApplyBroadcast<T, Context, n>(dev_ctx, in_tensors[i], out_tensors[i]); \
+    break;                                                                 \
+  }
+
+namespace phi {
+
+template <typename T, typename Context, int OutRank>
+void ApplyBroadcast(const Context& dev_ctx,
+                    const DenseTensor* input_tensor,
+                    DenseTensor* output_tensor) {
+  const auto& input_dims = input_tensor->dims();
+  const auto& output_dims = output_tensor->dims();
+
+  int in_rank = input_dims.size();
+  int out_rank = output_dims.size();
+
+  // 1. Collect bcast_dims, each element of which indicates how many
+  // times we need to replicate along the corresponding dimension
+  // 2. Collect new_input_dims_vec. Eigen::broadcast requires same rank for
+  // both input and output tensors, so we need to initialize input X with
+  // expanded dims: "new_input_dims_vec"
+  Eigen::DSizes<int64_t, OutRank> bcast_dims;
+  std::vector<int64_t> new_input_dims_vec(out_rank);
+  for (int i = 0; i < out_rank; i++) {
+    int in_axis = in_rank - i - 1;
+    int out_axis = out_rank - i - 1;
+
+    bcast_dims[out_axis] = output_dims[out_axis];
+    new_input_dims_vec[out_axis] = 1;
+    if (in_axis >= 0 && input_dims[in_axis] == output_dims[out_axis]) {
+      bcast_dims[out_axis] = 1;
+      new_input_dims_vec[out_axis] = input_dims[in_axis];
+    }
+  }
+  auto new_input_dims = make_ddim(new_input_dims_vec);
+
+  // Initialize input X with new_input_dims_vec, so it's rank-aligned with the
+  // output
+  auto x = EigenTensor<T, OutRank>::From(*input_tensor, new_input_dims);
+
+  dev_ctx.template Alloc<T>(output_tensor);
+  auto y = EigenTensor<T, OutRank>::From(*output_tensor, output_dims);
+
+  auto& place = *dev_ctx.eigen_device();
+  funcs::EigenBroadcast<std::decay_t<decltype(place)>, T, OutRank>::Eval(
+      place, y, x, bcast_dims);
+}
+
+template <typename T, typename Context>
+void BroadcastTensorsKernel(const Context& dev_ctx,
+                            const std::vector<const DenseTensor*>& x,
+                            std::vector<DenseTensor*> out) {
+  const auto& in_tensors = x;
+  auto out_tensors = out;
+  size_t num_ins = in_tensors.size();
+
+  PADDLE_ENFORCE_GE(
+      num_ins,
+      1,
+      errors::InvalidArgument(
+          "Expected at least 1 input tensor, but only received %d.",
+          in_tensors.size()));
+
+  PADDLE_ENFORCE_EQ(num_ins,
+                    out_tensors.size(),
+                    errors::InvalidArgument(
+                        "BroadcastTensorsOp expects equal number of inputs and "
+                        "outputs,but received: %d inputs v.s %d outputs",
+                        num_ins,
+                        out_tensors.size()));
+
+  // Eigen has no support for dynamic ranked tensor
+  // Thus we perform static expansion for each possible ranks
+  for (size_t i = 0; i < num_ins; i++) {
+    int out_rank = out_tensors[i]->dims().size();
+    switch (out_rank) {
+      case 0: {
+        const DenseTensor* src = in_tensors[i];
+        DenseTensor* dst = out_tensors[i];
+        Copy(dev_ctx, *src, src->place(), false, dst);
+        break;
+      }
+        SWITCH_OUT_RANK_CASE(1)
+        SWITCH_OUT_RANK_CASE(2)
+        SWITCH_OUT_RANK_CASE(3)
+        SWITCH_OUT_RANK_CASE(4)
+        SWITCH_OUT_RANK_CASE(5)
+        SWITCH_OUT_RANK_CASE(6)
+        SWITCH_OUT_RANK_CASE(7)
+        SWITCH_OUT_RANK_CASE(8)
+        SWITCH_OUT_RANK_CASE(9)
+      default: {
+        PADDLE_THROW(common::errors::InvalidArgument(
+            "Target tensor rank out of range. "
+            "Maximum supported rank for broadcast is: 9"));
+      }
+    }
+  }
+}
+
+}  // namespace phi