From c35a9880f918853d1a888220a89acaab700d4b1e Mon Sep 17 00:00:00 2001
From: QingshuChen <qingshu.chen714@gmail.com>
Date: Mon, 1 Feb 2021 10:23:59 +0800
Subject: [PATCH] fix malloc L3 failed bug for kunlun (#30745)

* fix malloc L3 failed bug for kunlun

* minor
---
 paddle/fluid/platform/device_context.cc | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/paddle/fluid/platform/device_context.cc b/paddle/fluid/platform/device_context.cc
index 23690cb8791..b9a8dd98456 100644
--- a/paddle/fluid/platform/device_context.cc
+++ b/paddle/fluid/platform/device_context.cc
@@ -189,13 +189,25 @@ XPUDeviceContext::XPUDeviceContext(XPUPlace place) : place_(place) {
                         "Baidu Kunlun Card is properly installed.",
                         ret));
   context_ = xpu::create_context();
-  void* l3ptr = nullptr;
-  int l3_size = 13.5 * 1024 * 1024;
-  xpu_malloc(static_cast<void**>(&l3ptr), l3_size, XPU_MEM_L3);
-  if (l3ptr != nullptr) {
-    context_->_l3_mgr.set(l3ptr, l3_size);
-    std::cout << "set l3 size " << l3_size << std::endl;
+  const int MAX_XPU_NUM = 16;
+  const int l3_size = 13.5 * 1024 * 1024;
+  static void* l3ptrs[MAX_XPU_NUM] = {nullptr};
+
+  auto selected_xpus = GetXPUSelectedDevices();
+  for (unsigned int i = 0; i < selected_xpus.size(); i++) {
+    if (place.device == selected_xpus[i]) {
+      if (l3ptrs[place.device] == nullptr) {
+        xpu_malloc(static_cast<void**>(&l3ptrs[place.device]), l3_size,
+                   XPU_MEM_L3);
+      }
+      if (l3ptrs[place.device] != nullptr) {
+        context_->_l3_mgr.set(l3ptrs[place.device], l3_size);
+        VLOG(3) << "xpu place " << place.device << " set l3 size " << l3_size;
+      }
+      break;
+    }
   }
+
   ret = xpu_set_device(dev_id);
   PADDLE_ENFORCE_EQ(ret, XPU_SUCCESS,
                     platform::errors::External(
-- 
GitLab