【Hackathon 5 No.5】为 Paddle 增强 scatter API

paddle/fluid/prim/api/composite_backward/composite_backward_api.h

@@ -1191,10 +1191,15 @@ void cos_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
 }
 
 template <typename T>
-void scatter_grad(const Tensor& index,
+void scatter_grad(const Tensor& x,
+                  const Tensor& index,
                   const Tensor& updates,
+                  const Tensor& out,
                   const Tensor& out_grad,
                   bool overwrite,
+                  int axis,
+                  const std::string& reduce,
+                  bool include_self,
                   Tensor* x_grad,
                   Tensor* updates_grad) {
   if (x_grad) {

paddle/phi/api/include/tensor.h

@@ -700,7 +700,10 @@ class PADDLE_API Tensor final {
               const std::vector<int64_t>& axis = {}) const;
   Tensor scatter(const Tensor& index,
                  const Tensor& updates,
-                 bool overwrite = true) const;
+                 bool overwrite = true,
+                 int axis = 0,
+                 const std::string& reduce = "add",
+                 bool include_self = false) const;
   Tensor scatter_nd_add(const Tensor& index, const Tensor& updates) const;
   Tensor abs() const;
   Tensor assign() const;

paddle/phi/api/yaml/backward.yaml

@@ -1918,16 +1918,15 @@
   invoke : scale(out_grad, scale, 0.0f, true)
 
 - backward_op : scatter_grad
-  forward : scatter (Tensor x, Tensor index, Tensor updates, bool overwrite=true) -> Tensor(out)
-  args : (Tensor index, Tensor updates, Tensor out_grad, bool overwrite)
+  forward : scatter (Tensor x, Tensor index, Tensor updates, bool overwrite=true, int axis=0, str reduce="add", bool include_self=false) -> Tensor(out)
+  args : (Tensor x, Tensor index, Tensor updates, Tensor out, Tensor out_grad, bool overwrite, int axis, str reduce, bool include_self)
   output : Tensor(x_grad), Tensor(updates_grad)
   infer_meta :
     func : ScatterGradInferMeta
-    param : [index, updates, out_grad, overwrite]
+    param : [index, updates, out_grad]
   kernel :
     func : scatter_grad
-  no_need_buffer : updates
-  composite: scatter_grad(index, updates, out_grad, overwrite, x_grad, updates_grad)
+  composite: scatter_grad(x, index, updates, out, out_grad, overwrite, axis, reduce, include_self, x_grad, updates_grad)
 
 - backward_op : scatter_nd_add_grad
   forward : scatter_nd_add (Tensor x, Tensor index, Tensor updates) -> Tensor(out)

paddle/phi/api/yaml/ops.yaml

@@ -2198,7 +2198,7 @@
   backward : scale_grad
 
 - op : scatter
-  args : (Tensor x, Tensor index, Tensor updates, bool overwrite=true)
+  args : (Tensor x, Tensor index, Tensor updates, bool overwrite=true, int axis=0, str reduce="add", bool include_self=false)
   output : Tensor(out)
   infer_meta :
     func : ScatterInferMeta

paddle/phi/backends/gpu/gpu_primitives.h

@@ -310,6 +310,123 @@ CUDA_ATOMIC_WRAPPER(Add, complex<double>) {
                          CudaAtomicAdd(imag, val.imag));
 }
 
+// Atomic multiplication implementation.
+CUDA_ATOMIC_WRAPPER(Mul, int64_t) {
+  // Here, we check long long int must be int64_t.
+  static_assert(sizeof(int64_t) == sizeof(long long int),  // NOLINT
+                "long long should be int64");
+  unsigned long long int *address_as_ull =       // NOLINT
+      (unsigned long long int *)address;         // NOLINT
+  unsigned long long int old = *address_as_ull;  // NOLINT
+  unsigned long long int assumed;                // NOLINT
+
+  do {
+    assumed = old;
+    old = atomicCAS(address_as_ull,
+                    assumed,
+                    static_cast<unsigned long long int>(  // NOLINT
+                        val * static_cast<int64_t>(assumed)));
+    // Note: uses integer comparison to avoid hang in case of NaN (since NaN
+    // != NaN)
+  } while (assumed != old);
+
+  return static_cast<int64_t>(old);
+}
+
+CUDA_ATOMIC_WRAPPER(Mul, int) {
+  int old = *address;
+  int assumed;
+
+  do {
+    assumed = old;
+    old = atomicCAS(address, assumed, val * assumed);
+
+    // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=
+    // NaN)
+  } while (assumed != old);
+
+  return old;
+}
+
+#ifdef PADDLE_CUDA_FP16
+CUDA_ATOMIC_WRAPPER(Mul, phi::dtype::float16) {
+  unsigned int *address_as_ui =
+      (unsigned int *)((char *)address - ((size_t)address & 2));  // NOLINT
+  unsigned int old = *address_as_ui;
+  unsigned int assumed;
+
+  phi::dtype::float16 hsum;
+  do {
+    assumed = old;
+    hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);  // NOLINT
+
+    hsum = hsum * val;
+    old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16)  // NOLINT
+                              : (old & 0xffff0000) | hsum.x;     // NOLINT
+    old = atomicCAS(address_as_ui, assumed, old);
+  } while (assumed != old);
+  hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);  // NOLINT
+  return hsum;
+}
+#endif
+
+CUDA_ATOMIC_WRAPPER(Mul, phi::dtype::bfloat16) {
+  unsigned int *address_as_ui =
+      (unsigned int *)((char *)address - ((size_t)address & 2));  // NOLINT
+  unsigned int old = *address_as_ui;
+  unsigned int assumed;
+
+  phi::dtype::bfloat16 bsum;
+  do {
+    assumed = old;
+    bsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);  // NOLINT
+    bsum = bsum * val;
+    old = (size_t)address & 2 ? (old & 0xffff) | (bsum.x << 16)  // NOLINT
+                              : (old & 0xffff0000) | bsum.x;     // NOLINT
+    old = atomicCAS(address_as_ui, assumed, old);
+  } while (assumed != old);
+  bsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);  // NOLINT
+  return bsum;
+}
+
+CUDA_ATOMIC_WRAPPER(Mul, double) {
+  unsigned long long int *address_as_ull =       // NOLINT
+      (unsigned long long int *)address;         // NOLINT
+  unsigned long long int old = *address_as_ull;  // NOLINT
+  unsigned long long int assumed;                // NOLINT
+
+  do {
+    assumed = old;
+    old = atomicCAS(
+        address_as_ull,
+        assumed,
+        __double_as_longlong(val * __longlong_as_double(assumed)));  // NOLINT
+    // Note: uses integer comparison to avoid hang in case of NaN (since NaN
+    // != NaN)
+  } while (assumed != old);
+
+  return __longlong_as_double(old);
+}
+
+// Dont use a templated function for this since the addition function defaults
+// to the CUDA built-in.
+CUDA_ATOMIC_WRAPPER(Mul, float) {
+  unsigned int *address_as_ull = (unsigned int *)address;  // NOLINT
+  unsigned int old = *address_as_ull;
+  unsigned int assumed;
+
+  do {
+    assumed = old;
+    old = atomicCAS(
+        address_as_ull, assumed, __float_as_int(val * __int_as_float(assumed)));
+
+    // Note: uses integer comparison to avoid hang in case of NaN (since NaN !=
+    // NaN)
+  } while (assumed != old);
+
+  return __int_as_float(old);
+}
+
 // For atomicMax
 USE_CUDA_ATOMIC(Max, int);
 USE_CUDA_ATOMIC(Max, unsigned int);

paddle/phi/infermeta/backward.cc

@@ -979,7 +979,6 @@ void RnnGradInferMeta(const MetaTensor& x,
 void ScatterGradInferMeta(const MetaTensor& index,
                           const MetaTensor& updates,
                           const MetaTensor& out_grad,
-                          bool overwrite,
                           MetaTensor* x_grad,
                           MetaTensor* updates_grad) {
   const auto& dtype = out_grad.dtype();

paddle/phi/infermeta/backward.h

@@ -401,7 +401,6 @@ void RnnGradInferMeta(const MetaTensor& x,
 void ScatterGradInferMeta(const MetaTensor& index,
                           const MetaTensor& updates,
                           const MetaTensor& out_grad,
-                          bool overwrite,
                           MetaTensor* x_grad,
                           MetaTensor* updates_grad);
 

paddle/phi/infermeta/ternary.cc

@@ -1035,6 +1035,9 @@ void ScatterInferMeta(const MetaTensor& x,
                       const MetaTensor& index,
                       const MetaTensor& updates,
                       bool overwrite,
+                      int axis,
+                      const std::string& reduce,
+                      bool include_self,
                       MetaTensor* out) {
   const auto& updates_dims = updates.dims();
   const auto& ref_dims = x.dims();

paddle/phi/infermeta/ternary.h

@@ -178,6 +178,9 @@ void ScatterInferMeta(const MetaTensor& x,
                       const MetaTensor& index,
                       const MetaTensor& updates,
                       bool overwrite,
+                      int axis,
+                      const std::string& reduce,
+                      bool include_self,
                       MetaTensor* out);
 
 void ScatterNdAddInferMeta(const MetaTensor& x,

paddle/phi/kernels/bitwise_kernel.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/infermeta/binary.h"
 
 namespace phi {
 
@@ -41,4 +42,17 @@ void BitwiseNotKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       DenseTensor* out);
 
+template <typename T, typename Context>
+DenseTensor BitwiseAnd(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y) {
+  DenseTensor dense_out;
+  MetaTensor meta_out(&dense_out);
+  MetaTensor meta_x(&x);
+  MetaTensor meta_y(&y);
+  ElementwiseInferMeta(meta_x, meta_y, &meta_out);
+  BitwiseAndKernel<T, Context>(dev_ctx, x, y, &dense_out);
+  return dense_out;
+}
+
 }  // namespace phi

paddle/phi/kernels/compare_kernel.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/infermeta/binary.h"
 
 namespace phi {
 
@@ -43,4 +44,30 @@ DECALRE_COMPARE_KERNEL(NotEqual)
 DECALRE_COMPARE_ALL_KERNEL(EqualAll)
 #undef DECALRE_COMPARE_KERNEL
 
+template <typename T, typename Context>
+DenseTensor Equal(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& y) {
+  DenseTensor dense_out;
+  MetaTensor meta_out(&dense_out);
+  MetaTensor meta_x(&x);
+  MetaTensor meta_y(&y);
+  CompareInferMeta(meta_x, meta_y, &meta_out);
+  EqualKernel<T, Context>(dev_ctx, x, y, &dense_out);
+  return dense_out;
+}
+
+template <typename T, typename Context>
+DenseTensor GreaterThan(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y) {
+  DenseTensor dense_out;
+  MetaTensor meta_out(&dense_out);
+  MetaTensor meta_x(&x);
+  MetaTensor meta_y(&y);
+  CompareInferMeta(meta_x, meta_y, &meta_out);
+  GreaterThanKernel<T, Context>(dev_ctx, x, y, &dense_out);
+  return dense_out;
+}
+
 }  // namespace phi