[XPU] add types and fix ut for stack op (#59375)

houj04 · web-flow · commit 68436e5c24dd · 2023-11-27T10:26:22.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -881,9 +881,14 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"stack_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::INT32})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
+                     phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"strided_slice",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/xpu/adamw_kernel.cc b/paddle/phi/kernels/xpu/adamw_kernel.cc
@@ -32,17 +32,16 @@ float GetAbsMax(const Context& dev_ctx,
                 float* buffer_xpu,
                 int64_t numel) {
   int max_ptr_size = phi::backends::xpu::get_xpu_max_ptr_size(-1);
-  float buffer_cpu[12];  // 12 is enough even for XPU3
+  std::vector<float> buffer_cpu(max_ptr_size);
   // int findmax(Context* ctx, const T* x, float* maxptr, int64_t len);
   int r = xpu::findmax<float>(dev_ctx.x_context(), input, buffer_xpu, numel);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax");
   memory_utils::Copy(CPUPlace(),
-                     static_cast<void*>(buffer_cpu),
+                     static_cast<void*>(buffer_cpu.data()),
                      dev_ctx.GetPlace(),
                      static_cast<void*>(buffer_xpu),
                      sizeof(float) * max_ptr_size);
-  float* max_value = std::max_element(buffer_cpu, buffer_cpu + max_ptr_size);
-  return *max_value;
+  return *std::max_element(buffer_cpu.begin(), buffer_cpu.end());
 }
 
 template <typename T, typename Context>
@@ -259,7 +258,8 @@ void AdamwDenseKernel(const Context& dev_ctx,
     using XPUType16 = typename XPUTypeTrait<phi::dtype::float16>::Type;
 
     // findmax and calculate scale_value for moment1 and moment2
-    float* buffer_for_findmax = RAII_GUARD.alloc_l3_or_gm<float>(6);
+    int max_ptr_size = phi::backends::xpu::get_xpu_max_ptr_size(-1);
+    float* buffer_for_findmax = RAII_GUARD.alloc_l3_or_gm<float>(max_ptr_size);
 
     // for moment1
     float moment1_max = GetAbsMax<Context>(dev_ctx,
diff --git a/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc
@@ -22,7 +22,6 @@
 namespace phi {
 
 template <typename T, typename Context>
-
 void SquaredL2NormGradKernel(const Context& dev_ctx,
                              const DenseTensor& x,
                              const DenseTensor& dout,
diff --git a/paddle/phi/kernels/xpu/stack_grad_kernel.cc b/paddle/phi/kernels/xpu/stack_grad_kernel.cc
@@ -49,5 +49,12 @@ void StackGradKernel(const Context& dev_ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    stack_grad, XPU, ALL_LAYOUT, phi::StackGradKernel, float, int) {}
+PD_REGISTER_KERNEL(stack_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::StackGradKernel,
+                   float,
+                   int,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/stack_kernel.cc b/paddle/phi/kernels/xpu/stack_kernel.cc
@@ -62,4 +62,5 @@ PD_REGISTER_KERNEL(stack,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_stack_op_xpu.py b/test/xpu/test_stack_op_xpu.py
@@ -20,7 +20,6 @@
     create_test_class,
     get_xpu_op_support_types,
 )
-from op_test import skip_check_grad_ci
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -33,19 +32,18 @@ def __init__(self):
         self.op_name = 'stack'
         self.use_dynamic_create_class = False
 
-    @skip_check_grad_ci(reason="There is no grad kernel for stack_xpu op.")
     class TestStackOp(XPUOpTest):
         def initDefaultParameters(self):
             self.num_inputs = 4
             self.input_dim = (5, 6, 7)
             self.axis = 0
-            self.dtype = np.float32
 
         def setUp(self):
             self.initDefaultParameters()
             self.initParameters()
             self.__class__.use_xpu = True
             self.__class__.op_type = 'stack'
+            self.dtype = self.in_type
             self.x = []
             for i in range(self.num_inputs):
                 self.x.append(
@@ -61,9 +59,6 @@ def setUp(self):
             self.outputs = {'Y': np.stack(self.x, axis=self.axis)}
             self.attrs = {'axis': self.axis}
 
-        def init_dtype(self):
-            self.dtype = self.in_type
-
         def initParameters(self):
             pass
 
@@ -77,12 +72,9 @@ def test_check_output(self):
             self.check_output_with_place(paddle.XPUPlace(0))
 
         def test_check_grad(self):
-            if self.dtype == np.int32 or self.dtype == np.int64:
-                pass
-            else:
-                self.check_grad_with_place(
-                    paddle.XPUPlace(0), self.get_x_names(), 'Y'
-                )
+            self.check_grad_with_place(
+                paddle.XPUPlace(0), self.get_x_names(), 'Y'
+            )
 
     class TestStackOp1(TestStackOp):
         def initParameters(self):
@@ -96,16 +88,10 @@ class TestStackOp3(TestStackOp):
         def initParameters(self):
             self.axis = -1
 
-        def test_check_grad(self):
-            pass
-
     class TestStackOp4(TestStackOp):
         def initParameters(self):
             self.axis = -4
 
-        def test_check_grad(self):
-            pass
-
     class TestStackOp5(TestStackOp):
         def initParameters(self):
             self.axis = 1
@@ -121,19 +107,13 @@ def initParameters(self):
             self.axis = 0
             self.dtype = np.int64
 
-        def test_check_grad(self):
-            pass
-
     class TestStackOp8(TestStackOp):
         def initParameters(self):
             self.num_inputs = 4
             self.input_dim = (5, 6, 7)
             self.axis = 0
             self.dtype = np.int32
 
-        def test_check_grad(self):
-            pass
-
 
 support_types = get_xpu_op_support_types('stack')
 for stype in support_types:
