[XPU] avoid malloc redundant memory when creating context in comm manager (#64139)

cqulilujia · web-flow · commit e7208bb6d6b3 · 2024-05-10T17:29:27.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu_context.cc b/paddle/phi/backends/xpu/xpu_context.cc
@@ -263,8 +263,13 @@ XPUContext::XPUContext() : DeviceContext() {
   }
 }
 
-XPUContext::XPUContext(const XPUPlace& place) : DeviceContext() {
-  if (std::getenv("XPU_CDNN_CLUSTER_PARALLEL") != nullptr) {
+XPUContext::XPUContext(const XPUPlace& place, bool is_comm_context)
+    : DeviceContext() {
+  if (is_comm_context) {
+    // for communication context init, with gm_size=1 and l3_size=1
+    impls_.push_back(std::make_unique<Impl>(place));
+    impls_[0]->Init(1, 1);
+  } else if (std::getenv("XPU_CDNN_CLUSTER_PARALLEL") != nullptr) {
     int default_num_stream = 4;
     if (std::getenv("XPU_CDNN_CLUSTER_PARALLEL_STREAM_NUMBER") != nullptr) {
       default_num_stream =
diff --git a/paddle/phi/backends/xpu/xpu_context.h b/paddle/phi/backends/xpu/xpu_context.h
@@ -38,7 +38,8 @@ class XPUContext : public DeviceContext,
  public:
   XPUContext();
 
-  explicit XPUContext(const XPUPlace&);
+  // is_comm_context = 1 for init comm context with gm_size=1 and l3_size=1
+  explicit XPUContext(const XPUPlace&, bool is_comm_context = 0);
 
   virtual ~XPUContext();
 
diff --git a/paddle/phi/core/distributed/comm_context_manager.cc b/paddle/phi/core/distributed/comm_context_manager.cc
@@ -211,8 +211,9 @@ void CommContextManager::CreateBKCLCommContext(
       std::make_unique<BKCLCommContext>(rank, size, bkcl_id);
 
   if (CommContextManager::device_id != -1) {
-    std::unique_ptr<phi::XPUContext> dev_ctx(
-        new phi::XPUContext(phi::XPUPlace(CommContextManager::device_id)));
+    bool is_comm_context = 1;
+    std::unique_ptr<phi::XPUContext> dev_ctx(new phi::XPUContext(
+        phi::XPUPlace(CommContextManager::device_id), is_comm_context));
     dev_ctx->SetAllocator(phi::memory_utils::GetAllocator(
         CommContextManager::device_id, dev_ctx->stream()));
     dev_ctx->SetHostAllocator(phi::memory_utils::GetHostAllocator());
