PaddlePaddle
diff --git a/‎paddle/fluid/operators/affine_channel_op_xpu.cc
Copy file name to clipboard
+4-9Lines changed: 4 additions & 9 deletions b/‎paddle/fluid/operators/affine_channel_op_xpu.cc
Copy file name to clipboard
+4-9Lines changed: 4 additions & 9 deletions
diff --git a/‎paddle/phi/kernels/funcs/adam_functors.h
Copy file name to clipboard
+17-16Lines changed: 17 additions & 16 deletions b/‎paddle/phi/kernels/funcs/adam_functors.h
Copy file name to clipboard
+17-16Lines changed: 17 additions & 16 deletions
diff --git a/‎paddle/phi/kernels/selected_rows/xpu/adam_kernel.cc
Copy file name to clipboard
+33-21Lines changed: 33 additions & 21 deletions b/‎paddle/phi/kernels/selected_rows/xpu/adam_kernel.cc
Copy file name to clipboard
+33-21Lines changed: 33 additions & 21 deletions
diff --git a/‎paddle/phi/kernels/xpu/adam_kernel.cc
Copy file name to clipboard
+34-22Lines changed: 34 additions & 22 deletions b/‎paddle/phi/kernels/xpu/adam_kernel.cc
Copy file name to clipboard
+34-22Lines changed: 34 additions & 22 deletions
@@ -144,11 +144,10 @@ class AffineChannelGradXPUKernel : public framework::OpKernel<T> {
                             "The reduce_sum XPU OP return wrong value[%d %s]",
                             r,
                             XPUAPIErrorMsg[r]));
-      T* tmp = nullptr;
-      r = xpu_malloc(reinterpret_cast<void**>(&tmp), dy->numel() * sizeof(T));
-      PADDLE_ENFORCE_EQ(r,
-                        xpu::Error_t::SUCCESS,
-                        platform::errors::External("no enough memory in xpu"));
+      xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+      T* tmp = RAII_GUARD.alloc_l3_or_gm<T>(dy->numel());
+      PADDLE_ENFORCE_NOT_NULL(
+          tmp, platform::errors::External("XPU has no enough memory"));
 
       r = xpu::mul<T>(
           dev_ctx.x_context(), dy_d, x->data<T>(), tmp, dy->numel());
@@ -166,10 +165,6 @@ class AffineChannelGradXPUKernel : public framework::OpKernel<T> {
                             "The reduce_sum XPU OP return wrong value[%d %s]",
                             r,
                             XPUAPIErrorMsg[r]));
-      if (dev_ctx.x_context()->xpu_stream) {
-        dev_ctx.Wait();
-      }
-      xpu_free(tmp);
     }
     if (dx_d) {
       r = xpu::broadcast_mul(
 
@@ -32,14 +32,17 @@ using float16 = dtype::float16;
 #ifdef PADDLE_WITH_XPU
 
 template <typename Context, typename T1, typename T2>
-static int ConvertDataByType(
-    const T1* x, T2** y, int len, bool allocateFlag, const Context& dev_ctx) {
+static int ConvertDataByType(const T1* x,
+                             T2** y,
+                             int len,
+                             bool allocateFlag,
+                             const Context& dev_ctx,
+                             xpu::ctx_guard* ctx_guard) {
   if (nullptr == x || nullptr == y || len <= 0)
     return xpu::Error_t::INVALID_PARAM;
-  int r = 0;
   if (allocateFlag) {
-    r = xpu_malloc(reinterpret_cast<void**>(y), sizeof(T2) * len);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
+    *y = ctx_guard->alloc_l3_or_gm<T2>(len);
+    PADDLE_ENFORCE_XDNN_NOT_NULL(*y);
   }
 
   T1* cpu_data = reinterpret_cast<T1*>(malloc(sizeof(T1) * len));
@@ -62,13 +65,14 @@ static int ConvertDataByType(
 template <typename Context, typename T>
 static void GetDataPointer(const phi::DenseTensor& tensorData,
                            T** result,
-                           const Context& dev_ctx) {
+                           const Context& dev_ctx,
+                           xpu::ctx_guard* ctx_guard) {
   if (tensorData.dtype() == DataType::FLOAT16) {
     const float16* real_data = tensorData.template data<float16>();
     int len = tensorData.numel();
 
     int r = ConvertDataByType<Context, float16, T>(
-        real_data, result, len, true, dev_ctx);
+        real_data, result, len, true, dev_ctx, ctx_guard);
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
   }
 }
@@ -88,23 +92,19 @@ static void GetOutDataPointer(DenseTensor* tensorData,
 template <typename Context, typename T>
 static void CopyOutData(const DenseTensor& srcTensor,
                         phi::DenseTensor* dstTensor,
-                        const Context& dev_ctx) {
+                        const Context& dev_ctx,
+                        xpu::ctx_guard* ctx_guard) {
   if (dstTensor->dtype() == DataType::FLOAT16) {
     const T* xpu_out_data = srcTensor.template data<T>();
     float16* out_data = dev_ctx.template Alloc<float16>(dstTensor);
     int len = srcTensor.numel();
 
     int r = ConvertDataByType<Context, T, float16>(
-        xpu_out_data, &out_data, len, false, dev_ctx);
+        xpu_out_data, &out_data, len, false, dev_ctx, ctx_guard);
     PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
   }
 }
 
-template <typename T>
-static void FreeData(const phi::DenseTensor& tensorData, T* dataPtr) {
-  if (tensorData.dtype() == DataType::FLOAT16) xpu_free(dataPtr);
-}
-
 template <typename Context, typename T>
 static void SetBetaData(const phi::DenseTensor& beta_pow,
                         phi::DenseTensor* beta_pow_out,
@@ -125,7 +125,8 @@ static void Scale(phi::DenseTensor* beta_pow_out,
                   const phi::DenseTensor& beta_pow,
                   T* beta_pow_ptr,
                   const T& beta,
-                  const Context& dev_ctx) {
+                  const Context& dev_ctx,
+                  xpu::ctx_guard* ctx_guard) {
   float16* beta_pow_out_p2 = dev_ctx.template Alloc<float16>(beta_pow_out);
 
   DenseTensor xpu_beta_pow_out;
@@ -149,7 +150,7 @@ static void Scale(phi::DenseTensor* beta_pow_out,
   int len = xpu_beta_pow_out.numel();
 
   r = ConvertDataByType<Context, T, float16>(
-      xpu_beta_pow_out_data, &beta_pow_out_p2, len, false, dev_ctx);
+      xpu_beta_pow_out_data, &beta_pow_out_p2, len, false, dev_ctx, ctx_guard);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
 }
 #endif
 
@@ -52,19 +52,23 @@ void AdamDenseParamSparseGradKernel(
     DenseTensor* beta2_pow_out,
     DenseTensor* master_param_outs) {
   using XPUType = typename XPUTypeTrait<T>::Type;
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
   float* param_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(param, &param_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      param, &param_ptr, dev_ctx, &RAII_GUARD);
 
   float* mom1_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(moment1, &mom1_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      moment1, &mom1_ptr, dev_ctx, &RAII_GUARD);
 
   float* mom2_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(moment2, &mom2_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      moment2, &mom2_ptr, dev_ctx, &RAII_GUARD);
 
   float* lr_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(learning_rate, &lr_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      learning_rate, &lr_ptr, dev_ctx, &RAII_GUARD);
 
-  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
   float* beta1_pow_ptr = nullptr;
   const float* beta1_const_pow_ptr = nullptr;
 
@@ -92,7 +96,8 @@ void AdamDenseParamSparseGradKernel(
 
   } else {
     if (beta1_pow.dtype() == DataType::FLOAT16)
-      funcs::GetDataPointer<Context, float>(beta1_pow, &beta1_pow_ptr, dev_ctx);
+      funcs::GetDataPointer<Context, float>(
+          beta1_pow, &beta1_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta1_const_pow_ptr = beta1_pow.template data<float>();
   }
@@ -123,7 +128,8 @@ void AdamDenseParamSparseGradKernel(
     }
   } else {
     if (beta2_pow.dtype() == DataType::FLOAT16)
-      funcs::GetDataPointer<Context, float>(beta2_pow, &beta2_pow_ptr, dev_ctx);
+      funcs::GetDataPointer<Context, float>(
+          beta2_pow, &beta2_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta2_const_pow_ptr = beta2_pow.template data<float>();
   }
@@ -225,7 +231,8 @@ void AdamDenseParamSparseGradKernel(
   auto& grad_merge = *grad_merge_ptr;
   auto& grad_tensor = grad_merge.value();
 
-  funcs::GetDataPointer<Context, float>(grad_tensor, &grad_c, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      grad_tensor, &grad_c, dev_ctx, &RAII_GUARD);
 
   int row_count = grad_merge.rows().size();
   std::vector<int> rows(row_count);
@@ -267,11 +274,12 @@ void AdamDenseParamSparseGradKernel(
 
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
 
-  funcs::FreeData<float>(grad_tensor, grad_c);
-
-  funcs::CopyOutData<Context, float>(xpu_mom1_out, moment1_out, dev_ctx);
-  funcs::CopyOutData<Context, float>(xpu_mom2_out, moment1_out, dev_ctx);
-  funcs::CopyOutData<Context, float>(xpu_param_out, moment1_out, dev_ctx);
+  funcs::CopyOutData<Context, float>(
+      xpu_mom1_out, moment1_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_mom2_out, moment1_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_param_out, moment1_out, dev_ctx, &RAII_GUARD);
 
   if (!use_global_beta_pow) {
     // update in cpu and then copy to xpu
@@ -285,8 +293,12 @@ void AdamDenseParamSparseGradKernel(
       float* beta1_pow_out_p1 = nullptr;
 
       if (beta1_pow_out->dtype() == DataType::FLOAT16) {
-        funcs::Scale<Context, float>(
-            beta1_pow_out, beta1_pow, beta1_pow_ptr, beta1_, dev_ctx);
+        funcs::Scale<Context, float>(beta1_pow_out,
+                                     beta1_pow,
+                                     beta1_pow_ptr,
+                                     beta1_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
       } else {
         const float* beta1_pow_data = beta1_pow.template data<float>();
         beta1_pow_out_p1 = dev_ctx.template Alloc<float>(beta1_pow_out);
@@ -303,8 +315,12 @@ void AdamDenseParamSparseGradKernel(
 
       float* beta2_pow_out_p1 = nullptr;
       if (beta2_pow_out->dtype() == DataType::FLOAT16) {
-        funcs::Scale<Context, float>(
-            beta2_pow_out, beta2_pow, beta2_pow_ptr, beta2_, dev_ctx);
+        funcs::Scale<Context, float>(beta2_pow_out,
+                                     beta2_pow,
+                                     beta2_pow_ptr,
+                                     beta2_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
       } else {
         const float* beta2_pow_data = beta2_pow.template data<float>();
         beta2_pow_out_p1 = dev_ctx.template Alloc<float>(beta2_pow_out);
@@ -320,10 +336,6 @@ void AdamDenseParamSparseGradKernel(
       }
     }
   }
-  funcs::FreeData<float>(param, param_ptr);
-  funcs::FreeData<float>(moment1, mom1_ptr);
-  funcs::FreeData<float>(moment2, mom2_ptr);
-  funcs::FreeData<float>(learning_rate, lr_ptr);
 }
 }  // namespace sr
 }  // namespace phi
 
@@ -49,17 +49,22 @@ void AdamDenseKernel(const Context& dev_ctx,
                      DenseTensor* beta1_pow_out,
                      DenseTensor* beta2_pow_out,
                      DenseTensor* master_param_outs) {
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
   float* param_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(param, &param_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      param, &param_ptr, dev_ctx, &RAII_GUARD);
 
   float* mom1_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(moment1, &mom1_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      moment1, &mom1_ptr, dev_ctx, &RAII_GUARD);
 
   float* mom2_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(moment2, &mom2_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      moment2, &mom2_ptr, dev_ctx, &RAII_GUARD);
 
   float* lr_ptr = nullptr;
-  funcs::GetDataPointer<Context, float>(learning_rate, &lr_ptr, dev_ctx);
+  funcs::GetDataPointer<Context, float>(
+      learning_rate, &lr_ptr, dev_ctx, &RAII_GUARD);
 
   float* beta1_pow_ptr = nullptr;
   const float* beta1_const_pow_ptr = nullptr;
@@ -68,12 +73,13 @@ void AdamDenseKernel(const Context& dev_ctx,
     phi::Copy(dev_ctx, beta1_pow, dev_ctx.GetPlace(), false, &xpu_beta1_pow);
     if (xpu_beta1_pow.dtype() == DataType::FLOAT16)
       funcs::GetDataPointer<Context, float>(
-          xpu_beta1_pow, &beta1_pow_ptr, dev_ctx);
+          xpu_beta1_pow, &beta1_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta1_const_pow_ptr = xpu_beta1_pow.template data<float>();
   } else {
     if (beta1_pow.dtype() == DataType::FLOAT16)
-      funcs::GetDataPointer<Context, float>(beta1_pow, &beta1_pow_ptr, dev_ctx);
+      funcs::GetDataPointer<Context, float>(
+          beta1_pow, &beta1_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta1_const_pow_ptr = beta1_pow.template data<float>();
   }
@@ -85,12 +91,13 @@ void AdamDenseKernel(const Context& dev_ctx,
     phi::Copy(dev_ctx, beta2_pow, dev_ctx.GetPlace(), false, &xpu_beta2_pow);
     if (xpu_beta2_pow.dtype() == DataType::FLOAT16)
       funcs::GetDataPointer<Context, float>(
-          xpu_beta2_pow, &beta2_pow_ptr, dev_ctx);
+          xpu_beta2_pow, &beta2_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta2_const_pow_ptr = xpu_beta2_pow.template data<float>();
   } else {
     if (beta2_pow.dtype() == DataType::FLOAT16)
-      funcs::GetDataPointer<Context, float>(beta2_pow, &beta2_pow_ptr, dev_ctx);
+      funcs::GetDataPointer<Context, float>(
+          beta2_pow, &beta2_pow_ptr, dev_ctx, &RAII_GUARD);
     else
       beta2_const_pow_ptr = beta2_pow.template data<float>();
   }
@@ -163,7 +170,7 @@ void AdamDenseKernel(const Context& dev_ctx,
   auto epsilon_ = epsilon.to<float>();
 
   float* grad_c = nullptr;
-  funcs::GetDataPointer<Context, float>(grad, &grad_c, dev_ctx);
+  funcs::GetDataPointer<Context, float>(grad, &grad_c, dev_ctx, &RAII_GUARD);
 
   int r = xpu::adam(
       dev_ctx.x_context(),
@@ -184,11 +191,12 @@ void AdamDenseKernel(const Context& dev_ctx,
 
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
 
-  funcs::FreeData<float>(grad, grad_c);
-
-  funcs::CopyOutData<Context, float>(xpu_mom1_out, moment1_out, dev_ctx);
-  funcs::CopyOutData<Context, float>(xpu_mom2_out, moment2_out, dev_ctx);
-  funcs::CopyOutData<Context, float>(xpu_param_out, param_out, dev_ctx);
+  funcs::CopyOutData<Context, float>(
+      xpu_mom1_out, moment1_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_mom2_out, moment2_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_param_out, param_out, dev_ctx, &RAII_GUARD);
 
   if (!use_global_beta_pow) {
     // update in cpu and then copy to xpu
@@ -202,8 +210,12 @@ void AdamDenseKernel(const Context& dev_ctx,
       float* beta1_pow_out_p1 = nullptr;
 
       if (beta1_pow_out->dtype() == DataType::FLOAT16) {
-        funcs::Scale<Context, float>(
-            beta1_pow_out, beta1_pow, beta1_pow_ptr, beta1_, dev_ctx);
+        funcs::Scale<Context, float>(beta1_pow_out,
+                                     beta1_pow,
+                                     beta1_pow_ptr,
+                                     beta1_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
       } else {
         const float* beta1_pow_data = beta1_pow.template data<float>();
         beta1_pow_out_p1 = dev_ctx.template Alloc<float>(beta1_pow_out);
@@ -219,8 +231,12 @@ void AdamDenseKernel(const Context& dev_ctx,
 
       float* beta2_pow_out_p1 = nullptr;
       if (beta2_pow_out->dtype() == DataType::FLOAT16) {
-        funcs::Scale<Context, float>(
-            beta2_pow_out, beta2_pow, beta2_pow_ptr, beta2_, dev_ctx);
+        funcs::Scale<Context, float>(beta2_pow_out,
+                                     beta2_pow,
+                                     beta2_pow_ptr,
+                                     beta2_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
       } else {
         const float* beta2_pow_data = beta2_pow.template data<float>();
         beta2_pow_out_p1 = dev_ctx.template Alloc<float>(beta2_pow_out);
@@ -235,10 +251,6 @@ void AdamDenseKernel(const Context& dev_ctx,
       }
     }
   }
-  funcs::FreeData<float>(param, param_ptr);
-  funcs::FreeData<float>(moment1, mom1_ptr);
-  funcs::FreeData<float>(moment2, mom2_ptr);
-  funcs::FreeData<float>(learning_rate, lr_ptr);
 }
 
 template <typename T, typename Context>