【Hackathon 5th No.1】 为 Paddle 新增 copysign API -part

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)

paddle/phi/api/yaml/ops.yaml

@@ -558,6 +558,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)

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

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

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) {}

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) {}

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) {}