【Comm】pack flagcx library into paddle whl (#71981)

Author: mikethegoblin

paddle/fluid/pybind/pybind.cc

@@ -308,6 +308,14 @@ bool IsCompiledWithNCCL() {
 #endif
 }
 
+bool IsCompiledWithFlagcx() {
+#ifdef PADDLE_WITH_FLAGCX
+  return true;
+#else
+  return false;
+#endif
+}
+
 bool IsCompiledWithMPI() {
 #ifdef PADDLE_WITH_MPI
   return true;
@@ -2582,6 +2590,7 @@ All parameter, weight, gradient are variables in Paddle.
         std::make_unique<paddle::operants::PhiTensorOperants>();
     VLOG(4) << "Initialize tensor operants successfully";
   });
+  m.def("is_compiled_with_flagcx", IsCompiledWithFlagcx);
   m.def("is_compiled_with_avx", IsCompiledWithAVX);
   m.def("is_compiled_with_cuda", IsCompiledWithCUDA);
   m.def("is_compiled_with_cudnn_frontend", IsCompiledWithCudnnFrontend);

python/env_dict.py.in

@@ -110,4 +110,6 @@ env_dict={
     'WITH_TENSORRT':'@WITH_TENSORRT@',
     'TR_INFER_RT':'@TR_INFER_RT@',
     'TENSORRT_LIBRARY_DIR':'@TENSORRT_LIBRARY_DIR@',
+    'WITH_FLAGCX':'@WITH_FLAGCX@',
+    'FLAGCX_LIB':'@FLAGCX_LIB@',
 }

python/paddle/distributed/fleet/launch_utils.py

@@ -1937,6 +1937,7 @@ def check_backend(backend):
         'auto',
         'heter',
         'xccl',
+        'flagcx',
     ]:
         raise ValueError(
             "paddle.distributed initialize error, "
@@ -1957,6 +1958,12 @@ def check_backend(backend):
             "your paddle is not compiled with xpu but you assign 'bkcl' as backend."
         )
 
+    if backend == 'flagcx' and not framework.core.is_compiled_with_flagcx():
+        raise ValueError(
+            "paddle.distributed initialize error, "
+            "your paddle is not compiled with flagcx but you assign 'flagcx' as backend."
+        )
+
 
 def block_windows_and_macos(backend):
     if backend != 'gloo':

python/paddle/distributed/parallel.py

@@ -932,7 +932,7 @@ def _start_kv_server(port, http_server_d, size):
 def _is_cpuonly(backend):
     check_backend(backend)
     if (
-        backend in ['auto', 'nccl', 'bkcl', 'heter']
+        backend in ['auto', 'nccl', 'bkcl', 'heter', 'flagcx']
         and (core.is_compiled_with_cuda() or core.is_compiled_with_xpu())
     ) or backend == 'xccl':
         # passes 'auto' and can use cuda or xpu, use the default logics. so return False
@@ -1134,7 +1134,7 @@ def init_parallel_env() -> Group:
         default_store = core.create_or_get_global_tcp_store()
         _set_default_store(default_store)
 
-        if backend in ["nccl", 'xccl', 'bkcl']:
+        if backend in ["nccl", 'xccl', 'bkcl', 'flagcx']:
             core.CommContextManager.set_device_id(parallel_env.device_id)
 
         pg = _new_process_group_impl(

python/paddle/jit/sot/utils/paddle_api_config.py

@@ -163,5 +163,6 @@ def is_directly_run_api(api):
         paddle.base.libpaddle.is_compiled_with_distribute,
         paddle.base.libpaddle.is_compiled_with_brpc,
         paddle.base.libpaddle.is_compiled_with_dist,
+        paddle.base.libpaddle.is_compiled_with_flagcx,
     }
     return api in NATIVE_CODE_PURE_FUNCTIONS

python/setup.py.in

@@ -960,6 +960,9 @@ package_dir={
 libs_path='${PADDLE_BINARY_DIR}/python/paddle/libs'
 
 package_data['paddle.libs']= []
+if('${WITH_FLAGCX}' == 'ON'):
+    package_data['paddle.libs'] += [('libflagcx' if os.name != 'nt' else 'flagcx') + ext_name]
+    shutil.copy('${FLAGCX_LIB}', libs_path)
 if('${WITH_SHARED_PHI}' == 'ON'):
     package_data['paddle.libs'] += [('libphi' if os.name != 'nt' else 'phi') + ext_name]
     shutil.copy('${PHI_LIB}', libs_path)

setup.py

@@ -1339,6 +1339,11 @@ def get_package_data_and_package_dir():
     # put all thirdparty libraries in paddle.libs
     libs_path = paddle_binary_dir + '/python/paddle/libs'
     package_data['paddle.libs'] = []
+    if env_dict.get("WITH_FLAGCX") == 'ON':
+        package_data['paddle.libs'] += [
+            ('libflagcx' if os.name != 'nt' else 'flagcx') + ext_suffix
+        ]
+        shutil.copy(env_dict.get("FLAGCX_LIB"), libs_path)
     if env_dict.get("WITH_SHARED_PHI") == "ON":
         package_data['paddle.libs'] += [
             ('libphi' if os.name != 'nt' else 'phi') + ext_suffix