From 350cd82a56d2159dcfdc005b7ef6bb671bef8ee4 Mon Sep 17 00:00:00 2001
From: Roc <30228238+sljlp@users.noreply.github.com>
Date: Thu, 9 Feb 2023 17:03:44 +0800
Subject: [PATCH] [kunlun] support async send/recv via group (#50329)

Co-authored-by: zhangxiaoci <zhangxiaoci@baidu.com>
---
 paddle/fluid/pybind/distributed_py.cc                         | 4 +++-
 .../fleet/meta_parallel/pp_utils/p2p_communication.py         | 4 ++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/pybind/distributed_py.cc b/paddle/fluid/pybind/distributed_py.cc
index cc06861d8d0..0b4a14d51b0 100644
--- a/paddle/fluid/pybind/distributed_py.cc
+++ b/paddle/fluid/pybind/distributed_py.cc
@@ -1251,7 +1251,9 @@ void BindDistributed(py::module *m) {
                       py::arg("rank"),
                       py::arg("world_size"),
                       py::arg("group_id") = 0,
-                      py::call_guard<py::gil_scoped_release>());
+                      py::call_guard<py::gil_scoped_release>())
+          .def_static("group_start", distributed::ProcessGroupBKCL::GroupStart)
+          .def_static("group_end", distributed::ProcessGroupBKCL::GroupEnd);
 #endif
 
   py::class_<distributed::ProcessGroup::Task,
diff --git a/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py b/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py
index 8586515b413..30563d67aee 100644
--- a/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py
+++ b/python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py
@@ -350,6 +350,8 @@ def _p2p_helper(
 
     # TODO(Yuang Liu): use batch_isend_irecv replace all these comm ops
     tasks = []
+    if paddle.is_compiled_with_xpu():
+        framework.core.ProcessGroupBKCL.group_start()
     # start to p2p communicate
     if tensor_send_prev is not None:
         if isinstance(tensor_send_prev, tuple):
@@ -479,6 +481,8 @@ def _p2p_helper(
                 )
             else:
                 tasks.append(task)
+    if paddle.is_compiled_with_xpu():
+        framework.core.ProcessGroupBKCL.group_end()
 
     if not sync_recv:
         if framework.in_dygraph_mode():
-- 
GitLab