From 85907b58a98e06cf7da9f16c294ad2d282dad78d Mon Sep 17 00:00:00 2001
From: jiangfan06 <117341294+MuShangCC@users.noreply.github.com>
Date: Fri, 26 May 2023 14:54:12 +0800
Subject: [PATCH] [XPU] fix l3 malloc (#54114)

---
 paddle/fluid/inference/api/analysis_predictor.cc | 2 +-
 paddle/fluid/inference/api/infer_context.cc      | 4 +++-
 paddle/fluid/inference/api/infer_context.h       | 5 ++++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
index 7abb5ef84e3..18c036a1ebe 100644
--- a/paddle/fluid/inference/api/analysis_predictor.cc
+++ b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -2180,7 +2180,7 @@ bool AnalysisPredictor::ExpRunWithRuntimeConfig(void *config) {
           "l3_autotune_size(%zu) should be less than or equal to l3_size(%zu).",
           l3_autotune_size,
           l3_size));
-  dev_ctx->SetL3Info(l3_size, l3_ptr, l3_autotune_size);
+  dev_ctx->SetL3Info(l3_size, l3_ptr, l3_autotune_size, place_);
 
   bool ret = ZeroCopyRun();
   dev_ctx->L3CacheAutotune();
diff --git a/paddle/fluid/inference/api/infer_context.cc b/paddle/fluid/inference/api/infer_context.cc
index 6c963f49363..bda0eecec0b 100644
--- a/paddle/fluid/inference/api/infer_context.cc
+++ b/paddle/fluid/inference/api/infer_context.cc
@@ -77,7 +77,9 @@ void* InferXPUContext::Alloc(phi::TensorBase* tensor,
 
 void InferXPUContext::SetL3Info(size_t l3_size,
                                 void* l3_ptr,
-                                size_t l3_autotune_size) {
+                                size_t l3_autotune_size,
+                                const phi::Place& place) {
+  phi::backends::xpu::XPUDeviceGuard guard(place.GetDeviceId());
   if (l3_ptr == nullptr) {
     if (l3_size_ != l3_size) {
       if (l3_owned_) {
diff --git a/paddle/fluid/inference/api/infer_context.h b/paddle/fluid/inference/api/infer_context.h
index ebc55098c97..121399dca5a 100644
--- a/paddle/fluid/inference/api/infer_context.h
+++ b/paddle/fluid/inference/api/infer_context.h
@@ -60,7 +60,10 @@ class InferXPUContext : public phi::XPUContext {
               bool pinned = false,
               bool fake_alloc = false) const override;
 
-  void SetL3Info(size_t l3_size, void* l3_ptr, size_t l3_autotune_size);
+  void SetL3Info(size_t l3_size,
+                 void* l3_ptr,
+                 size_t l3_autotune_size,
+                 const phi::Place& place);
 
   void L3CacheAutotune();
 
-- 
GitLab