add reduce_avg op (PaddlePaddle#64120)

co63oc · May 10, 2024 · 661d035 · 661d035
1 parent ae876f3
commit 661d035
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -70,6 +70,8 @@
 
 NEED_GEN_STATIC_ONLY_APIS = [
     'c_allreduce_avg_',
+    'c_reduce_avg',
+    'c_reduce_avg_',
     'c_allreduce_min_',
     'c_allreduce_prod_',
     'distributed_fused_lamb_init',
@@ -187,8 +189,6 @@
     'soft_relu',
     'uniform_random_batch_size_like',
     'match_matrix_tensor',
-    'c_reduce_avg',
-    'c_reduce_avg_',
     'c_reduce_max',
     'c_reduce_max_',
     'c_reduce_min',

diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -679,6 +679,12 @@ void BindOperation(py::module *m) {
              }
              return attrs_dict;
            })
+      .def("set_execution_stream",
+           [](Operation &self, const std::string &exe_stream) {
+             self.set_attribute("execution_stream",
+                                pir::StrAttribute::get(
+                                    pir::IrContext::Instance(), exe_stream));
+           })
       .def("set_scheduling_priority",
            [](Operation &self, int64_t priority) {
              self.set_attribute("scheduling_priority",