[XPU] fix several unit tests (#70479)

houj04 · web-flow · commit 2e2b3ded36b3 · 2024-12-27T14:49:05.000+08:00
* [XPU] fix several unit tests

* [XPU] fix several unit tests
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -193,6 +193,7 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::BOOL,
                      phi::DataType::INT8,
                      phi::DataType::UINT8,
+                     phi::DataType::INT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"check_finite_and_unscale",
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -96,6 +96,7 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BOOL})},
       {"assign_value",
        XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT64,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::FLOAT16,
@@ -215,6 +216,7 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BOOL,
                      phi::DataType::INT8,
                      phi::DataType::UINT8,
+                     phi::DataType::INT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"check_finite_and_unscale",
diff --git a/paddle/phi/kernels/cpu/flatten2_grad_kernel.cc b/paddle/phi/kernels/cpu/flatten2_grad_kernel.cc
@@ -25,5 +25,6 @@ PD_REGISTER_KERNEL(flatten2_grad,
                    uint8_t,
                    int,
                    int8_t,
+                   int16_t,
                    int64_t,
                    bool) {}
diff --git a/paddle/phi/kernels/cpu/flatten2_kernel.cc b/paddle/phi/kernels/cpu/flatten2_kernel.cc
@@ -25,5 +25,6 @@ PD_REGISTER_KERNEL(flatten2,
                    uint8_t,
                    int,
                    int8_t,
+                   int16_t,
                    int64_t,
                    bool) {}
diff --git a/paddle/phi/kernels/cpu/pad_grad_kernel.cc b/paddle/phi/kernels/cpu/pad_grad_kernel.cc
@@ -24,6 +24,7 @@ PD_REGISTER_KERNEL(pad_grad,
                    phi::PadGradKernel,
                    float,
                    double,
+                   int16_t,
                    int,
                    int64_t,
                    phi::dtype::complex<float>,
diff --git a/paddle/phi/kernels/cpu/pad_kernel.cc b/paddle/phi/kernels/cpu/pad_kernel.cc
@@ -24,6 +24,7 @@ PD_REGISTER_KERNEL(pad,
                    phi::PadKernel,
                    float,
                    double,
+                   int16_t,
                    int,
                    int64_t,
                    phi::dtype::complex<float>,
diff --git a/paddle/phi/kernels/flatten_grad_kernel.cc b/paddle/phi/kernels/flatten_grad_kernel.cc
@@ -45,6 +45,7 @@ PD_REGISTER_KERNEL(flatten_grad,
                    double,
                    uint8_t,
                    int8_t,
+                   int16_t,
                    int,
                    int64_t,
                    bool) {}
diff --git a/paddle/phi/kernels/reshape_kernel.cc b/paddle/phi/kernels/reshape_kernel.cc
@@ -19,9 +19,6 @@
 #include "paddle/phi/core/tensor_utils.h"
 #include "paddle/phi/infermeta/unary.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
-#ifdef PADDLE_WITH_XPU
-#include "paddle/phi/backends/xpu/enforce_xpu.h"
-#endif
 
 namespace phi {
 
@@ -46,34 +43,6 @@ void ReshapeKernel(const Context& dev_ctx,
   out->ResetLoD(x.lod());
 }
 
-#ifdef PADDLE_WITH_XPU
-template <>
-void ReshapeKernel<phi::XPUContext>(const XPUContext& dev_ctx,
-                                    const DenseTensor& x,
-                                    const IntArray& shape,
-                                    DenseTensor* out) {
-  MetaTensor meta_out(out);
-  InferMetaFromVecValue(x, shape.GetData(), &meta_out);
-
-  if (x.initialized() && x.Holder() == out->Holder()) {
-    dev_ctx.Alloc(out, x.dtype());
-    return;
-  }
-  dev_ctx.Alloc(out, x.dtype());
-  auto dims = out->dims();
-  auto* src_ptr = x.data();
-  auto* dst_ptr = out->data();
-  auto size = x.numel() * phi::SizeOf(x.dtype());
-  int ret = xpu::copy(dev_ctx.x_context(),
-                      reinterpret_cast<const int8_t*>(src_ptr),
-                      reinterpret_cast<int8_t*>(dst_ptr),
-                      size);
-  PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
-  out->Resize(dims);
-  out->ResetLoD(x.lod());
-}
-#endif
-
 template <typename Context>
 void ReshapeWithXShapeKernel(const Context& dev_ctx,
                              const DenseTensor& x,
diff --git a/paddle/phi/kernels/xpu/cast_kernel.cc b/paddle/phi/kernels/xpu/cast_kernel.cc
@@ -124,6 +124,7 @@ PD_REGISTER_KERNEL(cast,
                    XPU,
                    ALL_LAYOUT,
                    phi::CastKernel,
+                   int16_t,
                    int32_t,
                    float,
                    phi::dtype::float16,
diff --git a/test/dygraph_to_static/test_tensor_to.py b/test/dygraph_to_static/test_tensor_to.py
@@ -25,6 +25,7 @@
 )
 
 import paddle
+from paddle import base
 
 # NOTE: only test in PIR mode
 
@@ -38,19 +39,20 @@
     "int32",
     "int64",
     "uint8",
-    "complex64",
-    "complex128",
     "bool",
-]
+] + ([] if base.core.is_compiled_with_xpu() else ["complex64", "complex128"])
 
 _cpu_place = "Place(cpu)"
 _gpu_place = "Place(gpu:0)"
+_xpu_place = "Place(xpu:0)"
 
 
 def place_res():
     def res():
         if paddle.is_compiled_with_cuda():
             return _gpu_place
+        elif paddle.is_compiled_with_xpu():
+            return _xpu_place
         else:
             return _cpu_place
 
@@ -125,6 +127,8 @@ def test_tensor_to_dtype(self):
     def test_tensor_to_device(self):
         if paddle.is_compiled_with_cuda():
             x = paddle.to_tensor([1, 2, 3], place="gpu")
+        elif paddle.is_compiled_with_xpu():
+            x = paddle.to_tensor([1, 2, 3], place="xpu")
         else:
             x = paddle.to_tensor([1, 2, 3])
 
@@ -136,6 +140,8 @@ def test_tensor_to_device(self):
     def test_tensor_to_device2(self):
         if paddle.is_compiled_with_cuda():
             x = paddle.to_tensor([1, 2, 3], place="gpu")
+        elif paddle.is_compiled_with_xpu():
+            x = paddle.to_tensor([1, 2, 3], place="xpu")
         else:
             x = paddle.to_tensor([1, 2, 3])
 
@@ -150,6 +156,8 @@ def test_tensor_to_device_dtype(self):
         places = ["cpu"]
         if paddle.is_compiled_with_cuda():
             places.append("gpu")
+        if paddle.is_compiled_with_xpu():
+            places.append("xpu")
         for dtype in _valid_dtypes:
             for place in places:
                 tensor_x = paddle.jit.to_static(to_device_dtype)(
@@ -158,6 +166,8 @@ def test_tensor_to_device_dtype(self):
                 place_x_str = str(tensor_x.place)
                 if "gpu" == place:
                     self.assertEqual(place_x_str, _gpu_place)
+                elif "xpu" == place:
+                    self.assertEqual(place_x_str, _xpu_place)
                 else:
                     self.assertEqual(place_x_str, _cpu_place)
                 type_x_str = str(tensor_x.dtype)
diff --git a/test/dygraph_to_static/test_to_tensor.py b/test/dygraph_to_static/test_to_tensor.py
@@ -40,6 +40,8 @@ def case1(x):
 def case2(x):
     if core.is_compiled_with_cuda():
         place = paddle.CUDAPlace(0)
+    elif core.is_compiled_with_xpu():
+        place = paddle.XPUPlace(0)
     else:
         place = paddle.CPUPlace()
     a = paddle.to_tensor(
@@ -53,6 +55,8 @@ def case3(x):
     paddle.set_default_dtype("float64")
     if core.is_compiled_with_cuda():
         place = paddle.CUDAPlace(0)
+    elif core.is_compiled_with_xpu():
+        place = paddle.XPUPlace(0)
     else:
         place = paddle.CPUPlace()
     a = paddle.to_tensor([1.0, 2.0, 3.0], place=place)
@@ -64,6 +68,8 @@ def case4(x):
     paddle.set_default_dtype("float64")
     if core.is_compiled_with_cuda():
         place = paddle.CUDAPlace(0)
+    elif core.is_compiled_with_xpu():
+        place = paddle.XPUPlace(0)
     else:
         place = paddle.CPUPlace()
     a = paddle.to_tensor([1], place=place)
@@ -182,6 +188,8 @@ def test_static(self):
         with paddle.static.program_guard(main_prog, startup_prog):
             if core.is_compiled_with_cuda():
                 place = paddle.CUDAPlace(0)
+            elif core.is_compiled_with_xpu():
+                place = paddle.XPUPlace(0)
             else:
                 place = paddle.CPUPlace()
 
diff --git a/test/legacy_test/test_Tensor_to.py b/test/legacy_test/test_Tensor_to.py
@@ -31,10 +31,12 @@ def test_Tensor_to_dtype(self):
             "int32",
             "int64",
             "uint8",
-            "complex64",
-            "complex128",
             "bool",
-        ]
+        ] + (
+            []
+            if base.core.is_compiled_with_xpu()
+            else ["complex64", "complex128"]
+        )
         for dtype in valid_dtypes:
             tensorx = tensorx.to(dtype)
             typex_str = str(tensorx.dtype)
@@ -46,11 +48,14 @@ def test_Tensor_to_device(self):
         if base.core.is_compiled_with_cuda():
             places.append("gpu:0")
             places.append("gpu")
+        if base.core.is_compiled_with_xpu():
+            places.append("xpu:0")
+            places.append("xpu")
 
         for place in places:
             tensorx = tensorx.to(place)
             placex_str = str(tensorx.place)
-            if place == "gpu":
+            if place == "gpu" or place == "xpu":
                 self.assertTrue(placex_str, "Place(" + place + ":0)")
             else:
                 self.assertTrue(placex_str, "Place(" + place + ")")
@@ -68,6 +73,9 @@ def test_Tensor_to_device_dtype(self):
         if base.core.is_compiled_with_cuda():
             places.append("gpu:0")
             places.append("gpu")
+        if base.core.is_compiled_with_xpu():
+            places.append("xpu:0")
+            places.append("xpu")
         valid_dtypes = [
             "bfloat16",
             "float16",
@@ -78,15 +86,17 @@ def test_Tensor_to_device_dtype(self):
             "int32",
             "int64",
             "uint8",
-            "complex64",
-            "complex128",
             "bool",
-        ]
+        ] + (
+            []
+            if base.core.is_compiled_with_xpu()
+            else ["complex64", "complex128"]
+        )
         for dtype in valid_dtypes:
             for place in places:
                 tensorx = tensorx.to(place, dtype)
                 placex_str = str(tensorx.place)
-                if place == "gpu":
+                if place == "gpu" or place == "xpu":
                     self.assertTrue(placex_str, "Place(" + place + ":0)")
                 else:
                     self.assertTrue(placex_str, "Place(" + place + ")")
diff --git a/test/legacy_test/test_dygraph_multi_forward.py b/test/legacy_test/test_dygraph_multi_forward.py
@@ -188,11 +188,12 @@ def test_mnist_forward_float32(self):
                 paddle.framework.random._manual_program_seed(SEED)
             else:
                 paddle.framework.random._manual_program_seed(SEED)
-            exe = base.Executor(
-                base.CPUPlace()
-                if not core.is_compiled_with_cuda()
-                else base.CUDAPlace(0)
-            )
+            if core.is_compiled_with_cuda():
+                exe = base.Executor(base.CUDAPlace(0))
+            elif core.is_compiled_with_xpu():
+                exe = base.Executor(base.XPUPlace(0))
+            else:
+                exe = base.Executor(base.CPUPlace())
 
             mnist = MNIST()
             sgd = paddle.optimizer.SGD(learning_rate=1e-3)
diff --git a/test/legacy_test/test_random_seed.py b/test/legacy_test/test_random_seed.py
diff --git a/test/xpu/test_cast_op_xpu.py b/test/xpu/test_cast_op_xpu.py
