[XPU][PHI Kernels] bind bitwise_add kernel & add int32/int64 support to...

[XPU][PHI Kernels] bind bitwise_add kernel & add int32/int64 support to scatter_nd_add kernel for xpu (#54066)

* bind new kernels to xpu

* refine code

* fix bugs in unittest

paddle/phi/backends/xpu/xpu2_op_list.cc

@@ -81,6 +81,7 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"bilinear_interp_v2_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"bitwise_not", XPUKernelSet({phi::DataType::BOOL})},
+      {"bitwise_and", XPUKernelSet({phi::DataType::BOOL})},
       {"broadcast", XPUKernelSet({phi::DataType::FLOAT32})},
       {"c_allgather",
        XPUKernelSet({phi::DataType::FLOAT16,
@@ -644,7 +645,10 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::FLOAT32})},
       {"scatter_grad",
        XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
-      {"scatter_nd_add", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"scatter_nd_add",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"sampling_id",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT64})},
       {"set_value",

paddle/phi/kernels/xpu/bitwise.cc

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/logical_kernel.h"
 
 namespace phi {
 
@@ -29,8 +30,21 @@ void BitwiseNotKernel(const Context& ctx,
                            reinterpret_cast<const XPUDataType*>(x.data<T>()),
                            reinterpret_cast<XPUDataType*>(out->data<T>()),
                            x.numel());
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "bitwise not");
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "logical_not");
+}
+
+template <typename T, typename Context>
+void BitwiseAndKernel(const Context& ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& y,
+                      DenseTensor* out) {
+  // XPU api do not support bitwise operation now.
+  // However, because biwise and logical operation is identical for bool type,
+  // we can implement bitwise_and_bool kernel by calling their logical
+  // counterpart. Need to be changed when adding support to other types.
+  LogicalAndKernel<T, Context>(ctx, x, y, out);
 }
 }  // namespace phi
 
 PD_REGISTER_KERNEL(bitwise_not, XPU, ALL_LAYOUT, phi::BitwiseNotKernel, bool) {}
+PD_REGISTER_KERNEL(bitwise_and, XPU, ALL_LAYOUT, phi::BitwiseAndKernel, bool) {}

paddle/phi/kernels/xpu/scatter_nd_add_kernel.cc

@@ -38,13 +38,12 @@ void ScatterNdAddKernel(const Context &ctx,
         static_cast<int>(index.dims().size() == 0 ? 1 : index.dims()[0]);
 
     for (int i = 0; i < loop_time; i++) {
-      // xpu::add only support float or float16 template typename
-      // now, register this op only with float type
-      r = xpu::add<T>(ctx.x_context(),
+      r = xpu::broadcast_add<T>(ctx.x_context(),
                                 updates_ptr + out->numel() * i,
                                 out_ptr,
                                 out_ptr,
-                      out->numel());
+                                {out->numel()},
+                                {out->numel()});
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
     }
     return;
@@ -100,5 +99,10 @@ void ScatterNdAddKernel(const Context &ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    scatter_nd_add, XPU, ALL_LAYOUT, phi::ScatterNdAddKernel, float) {}
+PD_REGISTER_KERNEL(scatter_nd_add,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ScatterNdAddKernel,
+                   float,
+                   int64_t,
+                   int) {}

test/xpu/test_bitwise_op_xpu.py

@@ -39,12 +39,18 @@ class XPUTestBitwiseAnd(XPUOpTestWrapper):
     class XPUTestBitwiseAndBase(XPUOpTest):
         def setUp(self):
             self.place = paddle.XPUPlace(0)
+            self.dtype = self.in_type
             self.init_case()
             self.set_case()
 
         def set_case(self):
             self.op_type = 'bitwise_and'
 
+            # special range for bool dtype
+            if self.dtype == np.bool_:
+                self.low = 0
+                self.high = 2
+
             x = np.random.randint(
                 self.low, self.high, self.x_shape, dtype=self.dtype
             )
@@ -61,7 +67,6 @@ class XPUTestBitwiseAnd(XPUOpTestWrapper):
             self.outputs = {'Out': out}
 
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -75,7 +80,6 @@ class XPUTestBitwiseAnd(XPUOpTestWrapper):
 
     class XPUTestBitwiseAndCase1(XPUTestBitwiseAndBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -83,7 +87,6 @@ class XPUTestBitwiseAnd(XPUOpTestWrapper):
 
     class XPUTestBitwiseAndCase2(XPUTestBitwiseAndBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [4, 1]
             self.low = -100
@@ -91,7 +94,6 @@ class XPUTestBitwiseAnd(XPUOpTestWrapper):
 
     class XPUTestBitwiseAndCase3(XPUTestBitwiseAndBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = 0
@@ -111,12 +113,18 @@ class XPUTestBitwiseOr(XPUOpTestWrapper):
     class XPUTestBitwiseOrBase(XPUOpTest):
         def setUp(self):
             self.place = paddle.XPUPlace(0)
+            self.dtype = self.in_type
             self.init_case()
             self.set_case()
 
         def set_case(self):
             self.op_type = 'bitwise_or'
 
+            # special range for bool dtype
+            if self.dtype == np.bool_:
+                self.low = 0
+                self.high = 2
+
             x = np.random.randint(
                 self.low, self.high, self.x_shape, dtype=self.dtype
             )
@@ -133,7 +141,6 @@ class XPUTestBitwiseOr(XPUOpTestWrapper):
             self.outputs = {'Out': out}
 
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -147,7 +154,6 @@ class XPUTestBitwiseOr(XPUOpTestWrapper):
 
     class XPUTestBitwiseOrCase1(XPUTestBitwiseOrBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -155,7 +161,6 @@ class XPUTestBitwiseOr(XPUOpTestWrapper):
 
     class XPUTestBitwiseOrCase2(XPUTestBitwiseOrBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [4, 1]
             self.low = -100
@@ -163,7 +168,6 @@ class XPUTestBitwiseOr(XPUOpTestWrapper):
 
     class XPUTestBitwiseOrCase3(XPUTestBitwiseOrBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = 0
@@ -183,11 +187,16 @@ class XPUTestBitwiseXor(XPUOpTestWrapper):
     class XPUTestBitwiseXorBase(XPUOpTest):
         def setUp(self):
             self.place = paddle.XPUPlace(0)
+            self.dtype = self.in_type
             self.init_case()
             self.set_case()
 
         def set_case(self):
             self.op_type = 'bitwise_xor'
+            # special case for bool dtype
+            if self.dtype == np.bool_:
+                self.low = 0
+                self.high = 2
 
             x = np.random.randint(
                 self.low, self.high, self.x_shape, dtype=self.dtype
@@ -205,7 +214,6 @@ class XPUTestBitwiseXor(XPUOpTestWrapper):
             self.outputs = {'Out': out}
 
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -219,7 +227,6 @@ class XPUTestBitwiseXor(XPUOpTestWrapper):
 
     class XPUTestBitwiseXorCase1(XPUTestBitwiseXorBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = -100
@@ -227,7 +234,6 @@ class XPUTestBitwiseXor(XPUOpTestWrapper):
 
     class XPUTestBitwiseXorCase2(XPUTestBitwiseXorBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [4, 1]
             self.low = -100
@@ -235,7 +241,6 @@ class XPUTestBitwiseXor(XPUOpTestWrapper):
 
     class XPUTestBitwiseXorCase3(XPUTestBitwiseXorBase):
         def init_case(self):
-            self.dtype = np.int32
             self.x_shape = [2, 3, 4, 5]
             self.y_shape = [2, 3, 4, 5]
             self.low = 0