【Hackathon 5th No.1】为 Paddle 新增 copysign API (#57785)

* add copysign op * fix codestyle * codestyle * fix test * fix std bug * merge init * merge init * merge init * add static cast * add std * static cast * static cast * copysignf * static cast to float input * float input * static cast to double input * fix * add inplace test * fix api * fix cast when grad * modify paddle.cast_ to cast_ * remove cast in python api * support fp16 && bf16 * set grad y to zero * fix en doc * support number input * add hostdevice * refactor kernel * fix nan when backward * add broadcast unit test * modify .cu * Update __init__.py * Update __init__.py * for ci test * static float * codestyle * static double * fix broadcast, try coverage * Delete paddle/phi/kernels/funcs/broadcast_function.h * remove unused * Update math.py * Update math.py * fix en doc * add test for output dtype, integer unsupported for now * update * update * fix * fix * add cast for input * fix * add pir test * fix doc * fix doc * fix doc * detail doc * adjust for MSVC * fix * Update python/paddle/tensor/math.py Co-authored-by: zachary sun <70642955+sunzhongkai588@users.noreply.github.com> * Update python/paddle/tensor/math.py Co-authored-by: zachary sun <70642955+sunzhongkai588@users.noreply.github.com> * fix doc output dtype, fix Equation * codestyle * codestyle * Update math.py --------- Co-authored-by: zachary sun <70642955+sunzhongkai588@users.noreply.github.com>
PaddlePaddle · Jan 11, 2024 · 3184c3c · 3184c3c
1 parent f178fb4
commit 3184c3c
Show file tree

Hide file tree

Showing 13 changed files with 852 additions and 0 deletions.
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -413,6 +413,17 @@
     func : conv3d_transpose_grad
     data_type : x
 
+- backward_op : copysign_grad
+  forward : copysign (Tensor x, Tensor y) -> Tensor(out)
+  args : (Tensor x, Tensor y, Tensor out_grad)
+  output : Tensor(x_grad), Tensor(y_grad)
+  infer_meta :
+    func : GeneralBinaryGradInferMeta
+    param : [x, y]
+  kernel :
+    func : copysign_grad
+  inplace : (out_grad -> x_grad)
+
 - backward_op : cos_double_grad
   forward : cos_grad (Tensor x, Tensor grad_out) -> Tensor(grad_x)
   args : (Tensor x, Tensor grad_out, Tensor grad_x_grad)

diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -578,6 +578,16 @@
     data_type : x
   backward : conv3d_transpose_grad
 
+- op : copysign
+  args : (Tensor x, Tensor y)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+  kernel :
+    func : copysign
+  inplace: (x -> out)
+  backward : copysign_grad
+
 - op : cos
   args : (Tensor x)
   output : Tensor(out)

diff --git a/paddle/phi/kernels/copysign_grad_kernel.h b/paddle/phi/kernels/copysign_grad_kernel.h
@@ -0,0 +1,53 @@
+// Copyright (c) 2023 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/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/funcs/elementwise_functor.h"
+
+namespace phi {
+
+using float16 = phi::dtype::float16;
+using bfloat16 = phi::dtype::bfloat16;
+
+template <typename T>
+inline HOSTDEVICE auto copysign_func(const T& a, const T& b) {
+#ifdef WIN32
+  using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
+  return static_cast<T>(std::copysign(static_cast<U>(a), static_cast<U>(b)));
+#else
+  return static_cast<T>(std::copysign(a, b));
+#endif
+}
+
+inline HOSTDEVICE phi::dtype::float16 copysign_func(phi::dtype::float16 a,
+                                                    phi::dtype::float16 b) {
+  return phi::dtype::raw_uint16_to_float16((a.x & 0x7fff) | (b.x & 0x8000));
+}
+
+inline HOSTDEVICE phi::dtype::bfloat16 copysign_func(phi::dtype::bfloat16 a,
+                                                     phi::dtype::bfloat16 b) {
+  return phi::dtype::raw_uint16_to_bfloat16((a.x & 0x7fff) | (b.x & 0x8000));
+}
+
+template <typename T, typename Context>
+void CopySignGradKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y,
+                        const DenseTensor& out_grad,
+                        DenseTensor* x_grad,
+                        DenseTensor* y_grad);
+}  // namespace  phi
diff --git a/paddle/phi/kernels/copysign_kernel.h b/paddle/phi/kernels/copysign_kernel.h
@@ -0,0 +1,63 @@
+// Copyright (c) 2023 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/kernels/funcs/broadcast_function.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+namespace phi {
+
+using float16 = phi::dtype::float16;
+using bfloat16 = phi::dtype::bfloat16;
+
+template <typename T>
+inline HOSTDEVICE auto copysign_func(const T& a, const T& b) {
+#ifdef WIN32
+  using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
+  return static_cast<T>(std::copysign(static_cast<U>(a), static_cast<U>(b)));
+#else
+  return static_cast<T>(std::copysign(a, b));
+#endif
+}
+
+inline HOSTDEVICE phi::dtype::float16 copysign_func(phi::dtype::float16 a,
+                                                    phi::dtype::float16 b) {
+  return phi::dtype::raw_uint16_to_float16((a.x & 0x7fff) | (b.x & 0x8000));
+}
+
+inline HOSTDEVICE phi::dtype::bfloat16 copysign_func(phi::dtype::bfloat16 a,
+                                                     phi::dtype::bfloat16 b) {
+  return phi::dtype::raw_uint16_to_bfloat16((a.x & 0x7fff) | (b.x & 0x8000));
+}
+
+template <typename T>
+struct CopySignFunctor {
+  inline HOSTDEVICE T operator()(const T a, const T b) const {
+    return copysign_func(a, b);
+  }
+};
+template <typename T>
+struct InverseCopySignFunctor {
+  inline HOSTDEVICE T operator()(const T a, const T b) const {
+    return copysign_func(b, a);
+  }
+};
+
+template <typename T, typename Context>
+void CopySignKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const DenseTensor& y,
+                    DenseTensor* out);
+}  // namespace  phi
diff --git a/paddle/phi/kernels/cpu/copysign_grad_kernel.cc b/paddle/phi/kernels/cpu/copysign_grad_kernel.cc
@@ -0,0 +1,82 @@
+// Copyright (c) 2023 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/copysign_grad_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/cpu/elementwise_grad.h"
+
+namespace phi {
+
+template <typename T>
+HOSTDEVICE T compute_copysign_grad_dx(T x, T y, T out, T dout) {
+  if (x == static_cast<T>(0))
+    return x;
+  else
+    return static_cast<T>(dout * (phi::copysign_func(x, y) / x));
+}
+
+template <typename T>
+struct CopySignGradDX {
+  HOSTDEVICE T operator()(T x, T y, T out, T dout) const {
+    return compute_copysign_grad_dx<T>(x, y, out, dout);
+  }
+};
+
+template <typename T>
+struct CopySignGradDY {
+  HOSTDEVICE T operator()(T x, T y, T out, T dout) const {
+    return static_cast<T>(0);
+  }
+};
+
+template <typename T, typename Context>
+void CopySignGradKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y,
+                        const DenseTensor& out_grad,
+                        DenseTensor* x_grad,
+                        DenseTensor* y_grad) {
+  funcs::ElementwiseGradPreProcess(out_grad, x_grad);
+  int axis = -1;
+  phi::funcs::
+      ElemwiseGradCompute<Context, T, CopySignGradDX<T>, CopySignGradDY<T>>(
+          dev_ctx,
+          x,
+          y,
+          out_grad,
+          out_grad,
+          axis,
+          x_grad,
+          y_grad,
+          CopySignGradDX<T>(),
+          CopySignGradDY<T>());
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(copysign_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::CopySignGradKernel,
+                   bool,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/cpu/copysign_kernel.cc b/paddle/phi/kernels/cpu/copysign_kernel.cc
@@ -0,0 +1,51 @@
+// Copyright (c) 2023 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/copysign_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+namespace phi {
+template <typename T, typename Context>
+void CopySignKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const DenseTensor& y,
+                    DenseTensor* out) {
+  dev_ctx.template Alloc<T>(out);
+  auto x_dims = x.dims();
+  auto y_dims = y.dims();
+  if (x_dims.size() >= y_dims.size()) {
+    funcs::ElementwiseCompute<phi::CopySignFunctor<T>, T>(
+        dev_ctx, x, y, phi::CopySignFunctor<T>(), out);
+  } else {
+    funcs::ElementwiseCompute<phi::InverseCopySignFunctor<T>, T>(
+        dev_ctx, x, y, phi::InverseCopySignFunctor<T>(), out);
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(copysign,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::CopySignKernel,
+                   bool,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/gpu/copysign_grad_kernel.cu b/paddle/phi/kernels/gpu/copysign_grad_kernel.cu
@@ -0,0 +1,99 @@
+// Copyright (c) 2023 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/copysign_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/gpu/elementwise_grad.h"
+
+namespace phi {
+
+template <typename T>
+struct CopySignGradXFunctor {
+  inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
+    if (x == static_cast<T>(0)) return x;
+    return dout * (phi::copysign_func(x, y) / x);
+  }
+};
+
+template <typename T>
+struct CopySignGradYFunctor {
+  inline HOSTDEVICE T operator()(const T x, const T y, const T dout) const {
+    return static_cast<T>(0);
+  }
+};
+
+template <typename InT, typename OutT>
+struct CopySignGradXYFunctor {
+  inline HOSTDEVICE phi::Array<OutT, 2> operator()(const InT x,
+                                                   const InT y,
+                                                   const InT dout) {
+    phi::Array<OutT, 2> outs;
+    // dx
+    if (x == static_cast<InT>(0))
+      outs[0] = static_cast<OutT>(0);
+    else
+      outs[0] = static_cast<OutT>(dout * (phi::copysign_func(x, y)) / x);
+    // dy = 0
+    outs[1] = static_cast<OutT>(0);
+    return outs;
+  }
+};
+
+template <typename T, typename Context>
+void CopySignGradKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y,
+                        const DenseTensor& out_grad,
+                        DenseTensor* x_grad,
+                        DenseTensor* y_grad) {
+  const auto place = dev_ctx.GetPlace();
+  int axis = -1;
+  if (x_grad != nullptr && y_grad != nullptr) {
+    std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
+    GetGradXAndYOut<T>(dev_ctx,
+                       place,
+                       axis,
+                       ins,
+                       out_grad,
+                       x_grad,
+                       y_grad,
+                       CopySignGradXYFunctor<T, T>());
+  } else if (x_grad != nullptr && y_grad == nullptr) {
+    std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
+    GetGradXOrYOut<T>(
+        dev_ctx, place, axis, ins, out_grad, x_grad, CopySignGradXFunctor<T>());
+  } else if (y_grad != nullptr && x_grad == nullptr) {
+    std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
+    GetGradXOrYOut<T>(
+        dev_ctx, place, axis, ins, out_grad, y_grad, CopySignGradYFunctor<T>());
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(copysign_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::CopySignGradKernel,
+                   bool,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}