[XPU] fix memory reserved record when xpu malloc and free (#67828)

lj970926 · web-flow · commit 31186ac361f9 · 2024-08-30T14:46:03.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu_context.cc b/paddle/phi/backends/xpu/xpu_context.cc
@@ -125,11 +125,11 @@ struct XPUContext::Impl {
   void Wait() {
     backends::xpu::XPUDeviceGuard guard(place_.GetDeviceId());
     PD_CHECK(context_ != nullptr, "the xpu context is nullptr.");
-    xpu_wait(context_->xpu_stream);
+    PADDLE_ENFORCE_XRE_SUCCESS(xpu_wait(context_->xpu_stream));
     xpu::Context* ctx_t = GetXdlCtx();
     if (ctx_t) {
       PD_CHECK(ctx_t != nullptr, "the xpu context is nullptr.");
-      xpu_wait(ctx_t->xpu_stream);
+      PADDLE_ENFORCE_XRE_SUCCESS(xpu_wait(ctx_t->xpu_stream));
     }
 
     ClearStashedMemory();
@@ -220,7 +220,7 @@ struct XPUContext::Impl {
   void SetXContext(xpu::Context* context) {
     if (context_ != nullptr) {
       backends::xpu::XPUDeviceGuard guard(place_.GetDeviceId());
-      xpu_wait(context_->xpu_stream);
+      PADDLE_ENFORCE_XRE_SUCCESS(xpu_wait(context_->xpu_stream));
       if (context_->xpu_stream != nullptr && stream_owned_) {
         xpu_stream_destroy(context_->xpu_stream);
         stream_owned_ = false;
diff --git a/paddle/phi/core/platform/device/xpu/xpu_info.cc b/paddle/phi/core/platform/device/xpu/xpu_info.cc
@@ -23,6 +23,7 @@
 #include "paddle/phi/backends/xpu/xpu_header.h"
 #include "paddle/phi/backends/xpu/xpu_info.h"
 #include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/memory.h"
 #include "paddle/phi/core/platform/device/device_wrapper.h"
 #include "paddle/phi/core/platform/device_context.h"
 #include "paddle/phi/core/platform/lock_guard_ptr.h"
@@ -162,6 +163,7 @@ class RecordedXPUMallocHelper {
     auto result = xpu_malloc(ptr, size);
     if (result == XPU_SUCCESS) {
       cur_size_.fetch_add(size);
+      DEVICE_MEMORY_STAT_UPDATE(Reserved, dev_id_, size);
       return result;
     } else {
       RaiseNonOutOfMemoryError(result);
@@ -183,6 +185,7 @@ class RecordedXPUMallocHelper {
     dev_ctx->Wait();
     xpu_free(ptr);
     cur_size_.fetch_sub(size);
+    DEVICE_MEMORY_STAT_UPDATE(Reserved, dev_id_, -size);
   }
 
   inline bool NeedRecord() const { return limit_size_ != 0; }
