[XPU] bind mean_all_grad xpu kernel (#61148)

runzhech · web-flow · commit 77c660fe6ee7 · 2024-02-01T10:24:40.000+08:00
* bind mean_all_grad xpu kernel

* Update xpu.cmake

* bind mean_all_grad xpu kernel
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -29,7 +29,7 @@ if(NOT DEFINED XPU_BASE_DATE)
   set(XPU_BASE_DATE "20240104")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240125")
+  set(XPU_XHPC_BASE_DATE "20240129")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.1.8.1")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/kernels/xpu/mean_all_grad_kernel.cc b/paddle/phi/kernels/xpu/mean_all_grad_kernel.cc
@@ -26,17 +26,26 @@ void MeanAllGradKernel(const Context& dev_ctx,
                        const DenseTensor& out_grad,
                        DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
-
+  auto IG = x_grad;
+  dev_ctx.template Alloc<T>(IG);
   auto OG = &out_grad;
   PADDLE_ENFORCE_EQ(
       OG->numel(),
       1,
       phi::errors::InvalidArgument("Mean Gradient should be scalar"));
-  auto IG = x_grad;
-  dev_ctx.template Alloc<T>(IG);
+  auto dev_version =
+      phi::backends::xpu::get_xpu_version(dev_ctx.GetPlace().GetDeviceId());
+  if (dev_version == phi::backends::xpu::XPUVersion::XPU3) {
+    int r =
+        xpu::mean_all_grad(dev_ctx.x_context(),
+                           reinterpret_cast<const XPUType*>(out_grad.data<T>()),
+                           reinterpret_cast<XPUType*>(IG->data<T>()),
+                           IG->numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "mean_all_grad");
+    return;
+  }
 
   XPUType* dx = reinterpret_cast<XPUType*>(IG->data<T>());
-
   const T* dy = OG->data<T>();
   T dy0_value;
   xpu_wait(dev_ctx.x_context()->xpu_stream);
