[PIR] Adapt distributed API all_reduce (#62694)

SigureMo · web-flow · commit 55aba83fa183 · 2024-03-19T10:30:35.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -69,6 +69,9 @@
 {{"{name}", (PyCFunction)(void (*)(void)){name}, METH_VARARGS | METH_KEYWORDS, "C++ interface function for {name}."}},"""
 
 NEED_GEN_STATIC_ONLY_APIS = [
+    'c_allreduce_avg_',
+    'c_allreduce_min_',
+    'c_allreduce_prod_',
     'distributed_fused_lamb_init',
     'distributed_fused_lamb_init_',
     'fetch',
@@ -124,13 +127,10 @@
     'all_reduce_',
     'c_allgather',
     'c_allreduce_avg',
-    'c_allreduce_avg_',
     'c_allreduce_max',
     'c_allreduce_min',
-    'c_allreduce_min_',
     'c_allreduce_sum',
     'c_allreduce_prod',
-    'c_allreduce_prod_',
     'c_embedding',
     'c_identity',
     'c_reduce_sum',
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -315,7 +315,9 @@ std::set<std::string> GetRegisterDataType(const std::string& op_name) {
       data_type.insert(phi::DataTypeToString(info_pair.first.dtype()));
     }
   }
-
+  if (data_type.empty()) {
+    VLOG(6) << "No data type is registered for " << op_name;
+  }
   return data_type;
 }
 
diff --git a/python/paddle/distributed/communication/reduce.py b/python/paddle/distributed/communication/reduce.py
@@ -82,6 +82,10 @@ def _get_reduce_op(reduce_op, func_name):
     raise ValueError(f"Unknown reduce_op type for {func_name}.")
 
 
+def _to_inplace_op(op_name):
+    return f"{op_name}_"
+
+
 def reduce(tensor, dst, op=ReduceOp.SUM, group=None, sync_op=True):
     """
 
diff --git a/python/paddle/distributed/communication/stream/all_reduce.py b/python/paddle/distributed/communication/stream/all_reduce.py
@@ -12,13 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from paddle import framework
+from paddle import _C_ops, framework
 from paddle.base import data_feeder
 from paddle.distributed.communication.group import (
     _get_global_group,
     _warn_cur_rank_not_in_group,
 )
-from paddle.distributed.communication.reduce import ReduceOp, _get_reduce_op
+from paddle.distributed.communication.reduce import (
+    ReduceOp,
+    _get_reduce_op,
+    _to_inplace_op,
+)
+from paddle.framework import in_pir_mode
 
 
 def _all_reduce_in_dygraph(tensor, op, group, sync_op, use_calc_stream):
@@ -58,6 +63,11 @@ def _all_reduce_in_static_mode(tensor, op, group, sync_op, use_calc_stream):
     if not isinstance(ring_id, int):
         raise ValueError("The type of 'ring_id' for all_reduce should be int.")
 
+    if in_pir_mode():
+        op_type: str = _to_inplace_op(op_type)
+        getattr(_C_ops, op_type)(tensor, ring_id, sync_op, False)
+        return
+
     # TODO: Support task and use task.wait in static graph mode
     #       Use use_calc_stream rather than sync_op
     helper = framework.LayerHelper(op_type, **locals())
diff --git a/test/collective/CMakeLists.txt b/test/collective/CMakeLists.txt
@@ -208,6 +208,21 @@ if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   )
   set_tests_properties(test_collective_process_group PROPERTIES TIMEOUT "350")
 endif()
+if((WITH_GPU OR WITH_ROCM) AND (LINUX))
+  bash_test_modules(
+    test_collective_process_group_pir
+    START_BASH
+    ../legacy_test/dist_test.sh
+    TIMEOUT
+    "350"
+    LABELS
+    "RUN_TYPE=DIST"
+    ENVS
+    "PADDLE_DIST_UT_PORT=21294;http_proxy=;https_proxy=;PYTHONPATH=..:${PADDLE_BINARY_DIR}/python"
+  )
+  set_tests_properties(test_collective_process_group_pir PROPERTIES TIMEOUT
+                                                                    "350")
+endif()
 if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   py_test_modules(
     test_collective_reduce MODULES test_collective_reduce ENVS
diff --git a/test/collective/process_group_nccl_pir.py b/test/collective/process_group_nccl_pir.py
diff --git a/test/collective/test_collective_process_group_pir.py b/test/collective/test_collective_process_group_pir.py

Original file line number	Diff line number	Diff line change
`@@ -315,7 +315,9 @@ std::set<std::string> GetRegisterDataType(const std::string& op_name) {`
`315`	`315`	`data_type.insert(phi::DataTypeToString(info_pair.first.dtype()));`
`316`	`316`	`}`
`317`	`317`	`}`
`318`		`-`
	`318`	`+ if (data_type.empty()) {`
	`319`	`+ VLOG(6) << "No data type is registered for " << op_name;`
	`320`	`+ }`
`319`	`321`	`return data_type;`
`320`	`322`	`}`
`321`	`323`