[Prim] add reduce_as op for paddle

paddle/phi/api/yaml/backward.yaml

@@ -2388,6 +2388,16 @@
   kernel :
     func : stanh_grad
 
+- backward_op : sum_as_grad
+  forward : sum_as(Tensor x, Tensor y, DataType dtype=DataType::UNDEFINED, bool keepdim=false) -> Tensor(out)
+  args : (Tensor x, Tensor y, Tensor out_grad, bool keepdim)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : sum_as_grad
+
 - backward_op : svd_grad
   forward : svd (Tensor x, bool full_matrices = false) -> Tensor(u), Tensor(s), Tensor(vh)
   args : (Tensor x, Tensor u, Tensor vh, Tensor s, Tensor u_grad, Tensor vh_grad, Tensor s_grad, bool full_matrices)

paddle/phi/api/yaml/ops.yaml

@@ -2751,6 +2751,16 @@
     func : stanh
   backward : stanh_grad
 
+- op : sum_as
+  args : (Tensor x, Tensor y, DataType dtype=DataType::UNDEFINED, bool keepdim=false)
+  output : Tensor(out)
+  infer_meta :
+    func : SumAsInferMeta
+  kernel :
+    func : sum_as
+    data_type : x
+  backward : sum_as_grad
+
 - op : svd
   args : (Tensor x, bool full_matrices = false)
   output : Tensor(u), Tensor(s), Tensor(vh)

paddle/phi/infermeta/binary.cc

@@ -3047,6 +3047,26 @@ void SequenceMaskInferMeta(const MetaTensor& x,
   y->set_dtype(out_dtype);
 }
 
+void SumAsInferMeta(const MetaTensor& x,
+                    const MetaTensor& y,
+                    DataType dtype,
+                    bool keep_dim,
+                    MetaTensor* out) {
+  DataType out_dtype;
+  if (dtype != DataType::UNDEFINED) {
+    out_dtype = dtype;
+  } else {
+    if (x.dtype() == DataType::BOOL || x.dtype() == DataType::INT32) {
+      out_dtype = DataType::INT64;
+    } else {
+      out_dtype = x.dtype();
+    }
+  }
+  out->set_dtype(out_dtype);
+  out->set_dims(y.dims());
+  out->set_layout(x.layout());
+}
+
 void SoftmaxMaskFuseInferMeta(const MetaTensor& x,
                               const MetaTensor& mask,
                               MetaTensor* out) {

paddle/phi/infermeta/binary.h

@@ -524,6 +524,12 @@ void ShuffleBatchInferMeta(const MetaTensor& x,
 
 );
 
+void SumAsInferMeta(const MetaTensor& x,
+                    const MetaTensor& y,
+                    DataType dtype,
+                    bool keep_dim,
+                    MetaTensor* out);
+
 void SoftmaxMaskFuseInferMeta(const MetaTensor& x,
                               const MetaTensor& mask,
                               MetaTensor* out);

paddle/phi/kernels/cpu/sum_as_grad_kernel.cc

@@ -0,0 +1,60 @@
+
+// Copyright (c) 2024 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/sum_as_kernel.h"
+
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/reduce_grad.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SumAsGradKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const DenseTensor& y,
+                     const DenseTensor& out_grad,
+                     bool keep_dim,
+                     DenseTensor* x_grad) {
+  auto reduce_dim = phi::funcs::GetReduceDims(x, y);
+  bool reduce_all = recompute_reduce_all(x, reduce_dim);
+  ReduceGradKernel<Context, T, funcs::SumGradFunctor, true>(dev_ctx,
+                                                            x,
+                                                            paddle::none,
+                                                            out_grad,
+                                                            reduce_dim,
+                                                            keep_dim,
+                                                            reduce_all,
+                                                            x_grad);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_as_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SumAsGradKernel,
+                   bool,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int16_t,
+                   int,
+                   int64_t,
+                   uint8_t,
+                   int8_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/cpu/sum_as_kernel.cc

@@ -0,0 +1,57 @@
+
+// Copyright (c) 2024 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/sum_as_kernel.h"
+
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cpu/reduce.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SumAsKernel(const Context& dev_ctx,
+                 const DenseTensor& x,
+                 const DenseTensor& y,
+                 DataType out_dtype,
+                 bool keep_dim,
+                 DenseTensor* out) {
+  auto reduce_dim = phi::funcs::GetReduceDims(x, y);
+  bool reduce_all = recompute_reduce_all(x, reduce_dim);
+  if (out_dtype == DataType::UNDEFINED && out->dtype() != x.dtype()) {
+    out_dtype = out->dtype();
+  }
+  phi::Reduce<CPUContext, T, phi::funcs::SumFunctor>(
+      dev_ctx, x, reduce_all, reduce_dim, keep_dim, out_dtype, out);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_as,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SumAsKernel,
+                   bool,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int16_t,
+                   int,
+                   int64_t,
+                   uint8_t,
+                   int8_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/funcs/common_shape.h

@@ -297,5 +297,23 @@ inline void FCOutputSize(const DDim &in_dims,
   out_dims.push_back(w_dims1);
 }
 
+inline std::vector<int64_t> GetReduceDims(const DenseTensor &in,
+                                          const DenseTensor &out) {
+  std::vector<int64_t> reduce_dims;
+  auto in_dims = in.dims();
+  auto out_dims = out.dims();
+
+  int diff = in_dims.size() - out_dims.size();
+  for (int i = 0; i < diff; ++i) {
+    reduce_dims.push_back(i);
+  }
+  for (int i = diff; i < in_dims.size(); ++i) {
+    if (out_dims[i - diff] != in_dims[i]) {
+      reduce_dims.push_back(i);
+    }
+  }
+  return reduce_dims;
+}
+
 }  // namespace funcs
 }  // namespace phi

paddle/phi/kernels/gpu/sum_as_grad_kernel.cu

@@ -0,0 +1,70 @@
+
+// Copyright (c) 2024 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/sum_as_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/gpu/reduce_grad.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SumAsGradKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const DenseTensor& y,
+                     const DenseTensor& out_grad,
+                     bool keep_dim,
+                     DenseTensor* x_grad) {
+  auto reduce_dim = phi::funcs::GetReduceDims(x, y);
+  bool reduce_all = recompute_reduce_all(x, reduce_dim);
+  auto update_dims = common::vectorize(x.dims());
+  for (auto i : reduce_dim) {
+    update_dims[i] = 1;
+  }
+
+  DenseTensor new_out_grad(out_grad.type());
+  new_out_grad.ShareDataWith(out_grad);
+  new_out_grad.Resize(common::make_ddim(update_dims));
+
+  dev_ctx.Alloc(x_grad, x.dtype());
+  using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
+  phi::ReduceGrad<phi::kps::IdentityFunctor<T, MPType>>(
+      dev_ctx,
+      &new_out_grad,
+      x_grad,
+      x.dtype(),
+      phi::kps::IdentityFunctor<T, MPType>());
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_as_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SumAsGradKernel,
+                   bool,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int16_t,
+                   int,
+                   int64_t,
+                   uint8_t,
+                   int8_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/gpu/sum_as_kernel.cu

@@ -0,0 +1,53 @@
+
+// Copyright (c) 2024 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/sum_as_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SumAsKernel(const Context& dev_ctx,
+                 const DenseTensor& x,
+                 const DenseTensor& y,
+                 DataType out_dtype,
+                 bool keep_dim,
+                 DenseTensor* out) {
+  auto reduce_dim = phi::funcs::GetReduceDims(x, y);
+  dev_ctx.template Alloc<T>(out);
+  phi::SumKernel<T, Context>(dev_ctx, x, reduce_dim, out_dtype, keep_dim, out);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(sum_as,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SumAsKernel,
+                   bool,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int16_t,
+                   int,
+                   int64_t,
+                   uint8_t,
+                   int8_t) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
+}

paddle/phi/kernels/sum_as_grad_kernel.h

@@ -0,0 +1,32 @@
+// Copyright (c) 2024 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 "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/kernels/funcs/common_shape.h"
+#include "paddle/phi/kernels/funcs/reduce_functor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SumAsGradKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const DenseTensor& y,
+                     const DenseTensor& out_grad,
+                     bool keep_dim,
+                     DenseTensor* x_grad);
+
+}  // namespace phi