support to get custom comm name. (#62556)

wuhuachaocoding · web-flow · commit 7489b0675a2e · 2024-03-11T16:29:01.000+08:00
diff --git a/paddle/fluid/distributed/collective/process_group_custom.cc b/paddle/fluid/distributed/collective/process_group_custom.cc
@@ -161,6 +161,32 @@ phi::ccl::CCLComm ProcessGroupCustom::XCCLComm(const Place& place) const {
   return iter->second->xccl_comm();
 }
 
+std::string ProcessGroupCustom::GetCommName(int rank) {
+  PADDLE_ENFORCE_GE(rank,
+                    0,
+                    phi::errors::PreconditionNotMet(
+                        "The rank must greater or equal than 0!"));
+  auto num_devices = phi::DeviceManager::GetDeviceCount(device_type_);
+  PADDLE_ENFORCE_GT(
+      num_devices,
+      0,
+      phi::errors::InvalidArgument("The num_devices must greater than 0!"));
+
+  auto place_id = rank % num_devices;
+  platform::CustomPlace place(device_type_, place_id);
+  const auto& key = GetKeyFromPlace(place);
+  phi::DeviceGuard guard(place);
+  if (place_to_comm_ctx_.find(key) == place_to_comm_ctx_.end()) {
+    CreateXCCLEnvCache(place, key);
+  }
+
+  char comm_name[128];
+  phi::DeviceManager::CCLCommName(
+      device_type_, this->GetCommContext()->GetXcclComm(), comm_name);
+  std::string name_str(comm_name);
+  return name_str;
+}
+
 std::shared_ptr<ProcessGroup::Task> ProcessGroupCustom::AllGather(
     phi::DenseTensor* out_tensor,
     const phi::DenseTensor& in_tensor,
diff --git a/paddle/fluid/distributed/collective/process_group_custom.h b/paddle/fluid/distributed/collective/process_group_custom.h
@@ -82,6 +82,8 @@ class ProcessGroupCustom final : public ProcessGroupWithStream {
 
   std::string GetBackendName() const override { return "XCCL"; }
 
+  std::string GetCommName(int rank);
+
   phi::DeviceContext* GetDeviceContext(const Place& place) const override;
 
   phi::DeviceContext* GetDeviceContext(const Place& place,
diff --git a/paddle/fluid/pybind/distributed_py.cc b/paddle/fluid/pybind/distributed_py.cc
@@ -1273,7 +1273,11 @@ void BindDistributed(py::module *m) {
                   py::arg("world_size"),
                   py::arg("group_id") = 0,
                   py::return_value_policy::reference_internal,
-                  py::call_guard<py::gil_scoped_release>());
+                  py::call_guard<py::gil_scoped_release>())
+      .def("get_comm_name",
+           &distributed::ProcessGroupCustom::GetCommName,
+           py::arg("rank"),
+           py::call_guard<py::gil_scoped_release>());
 
 #endif
 
diff --git a/paddle/phi/backends/custom/custom_device.cc b/paddle/phi/backends/custom/custom_device.cc
@@ -651,6 +651,12 @@ class CustomDevice : public DeviceInterface {
         pimpl_->xccl_destroy_comm(reinterpret_cast<C_CCLComm>(comm)));
   }
 
+  void CCLCommName(ccl::CCLComm comm, char* comm_name) {
+    CHECK_PTR(pimpl_->xccl_get_comm_name);
+    PADDLE_ENFORCE_CUSTOM_DEVICE_SUCCESS(pimpl_->xccl_get_comm_name(
+        reinterpret_cast<C_CCLComm>(comm), comm_name));
+  }
+
   void CCLAllReduce(void* send_buf,
                     void* recv_buf,
                     size_t count,
diff --git a/paddle/phi/backends/device_base.cc b/paddle/phi/backends/device_base.cc
@@ -267,6 +267,10 @@ size_t DeviceInterface::GetExtraPaddingSize(size_t dev_id) {
   return 0;
 }
 
+void DeviceInterface::CCLCommName(ccl::CCLComm ccl_comm, char* comm_name) {
+  INTERFACE_UNIMPLEMENT;
+}
+
 void DeviceInterface::CCLDestroyComm(ccl::CCLComm ccl_comm) {
   INTERFACE_UNIMPLEMENT;
 }
diff --git a/paddle/phi/backends/device_base.h b/paddle/phi/backends/device_base.h
@@ -169,6 +169,8 @@ class DeviceInterface {  // Driver / Runtime
   virtual size_t GetExtraPaddingSize(size_t dev_id);
 
   // CCL
+  virtual void CCLCommName(ccl::CCLComm ccl_comm, char* comm_name);
+
   virtual void CCLDestroyComm(ccl::CCLComm ccl_comm);
 
   virtual void CCLCommInitRank(size_t num_ranks,
diff --git a/paddle/phi/backends/device_ext.h b/paddle/phi/backends/device_ext.h
@@ -547,6 +547,13 @@ struct C_DeviceInterface {
   // ccl api //
   //////////////
 
+  /**
+   * @brief Get comm name.
+   *
+   * @param[char*]         comm_name
+   */
+  C_Status (*xccl_get_comm_name)(C_CCLComm comm, char* comm_name);
+
   /**
    * @brief Get size of unique id
    *
diff --git a/paddle/phi/backends/device_manager.cc b/paddle/phi/backends/device_manager.cc
@@ -509,6 +509,13 @@ std::vector<size_t> DeviceManager::GetSelectedDeviceList(
   return device_list_map[device_type];
 }
 
+void DeviceManager::CCLCommName(const std::string& device_type,
+                                const ccl::CCLComm& ccl_comm,
+                                char* comm_name) {
+  auto dev_impl = GetDeviceInterfaceWithType(device_type);
+  return dev_impl->CCLCommName(ccl_comm, comm_name);
+}
+
 void DeviceManager::CCLDestroyComm(const std::string& device_type,
                                    ccl::CCLComm ccl_comm) {
   auto dev_impl = GetDeviceInterfaceWithType(device_type);
diff --git a/paddle/phi/backends/device_manager.h b/paddle/phi/backends/device_manager.h
@@ -178,6 +178,9 @@ class DeviceManager {
       const std::string& device_type);
 
   // CCL
+  static void CCLCommName(const std::string& device_type,
+                          const ccl::CCLComm& ccl_comm,
+                          char* comm_name);
   static void CCLDestroyComm(const std::string& device_type,
                              ccl::CCLComm ccl_comm);
   static void CCLCommInitRank(const std::string& device_type,

Original file line number	Diff line number	Diff line change
`@@ -267,6 +267,10 @@ size_t DeviceInterface::GetExtraPaddingSize(size_t dev_id) {`
`267`	`267`	`return 0;`
`268`	`268`	`}`
`269`	`269`
	`270`	`+void DeviceInterface::CCLCommName(ccl::CCLComm ccl_comm, char* comm_name) {`
	`271`	`+ INTERFACE_UNIMPLEMENT;`
	`272`	`+}`
	`273`	`+`
`270`	`274`	`void DeviceInterface::CCLDestroyComm(ccl::CCLComm ccl_comm) {`
`271`	`275`	`INTERFACE_UNIMPLEMENT;`
`272`	`276`	`}`