From 661d0358b0247669aad23e4b4ba354811d858a38 Mon Sep 17 00:00:00 2001
From: lzydev <lizhiyu02@baidu.com>
Date: Fri, 10 May 2024 16:35:53 +0800
Subject: [PATCH] add reduce_avg op (#64120)

---
 paddle/fluid/pir/dialect/op_generator/ops_api_gen.py | 4 ++--
 paddle/fluid/pybind/pir.cc                           | 6 ++++++
 2 files changed, 8 insertions(+), 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
index 5f7cfce1f5e18d..c53eaa708229d1 100644
--- 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
index 18febc9ff6754f..133289501c5a62 100644
--- 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",