add comp_op_name (#71098)

changeyoung98 · web-flow · commit 179d15784e3d · 2025-02-13T14:29:11.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/api_builder.h b/paddle/fluid/pir/dialect/operator/ir/api_builder.h
@@ -81,6 +81,11 @@ class ApiBuilder {
   void SetChunkId(int chunk_id) { builder_->set_chunk_id(chunk_id); }
   int GetChunkId() const { return builder_->chunk_id(); }
 
+  void SetCompOpName(std::string op_name) {
+    builder_->set_comp_op_name(op_name);
+  }
+  std::string GetCompOpName() const { return builder_->comp_op_name(); }
+
  private:
   ApiBuilder();
 
diff --git a/paddle/fluid/primitive/base/decomp_trans.cc b/paddle/fluid/primitive/base/decomp_trans.cc
@@ -504,7 +504,8 @@ void DecompProgram::decomp_block(
       int chunk_id = (op->attribute<pir::Int32Attribute>("chunk_id"))
                          ? op->attribute<pir::Int32Attribute>("chunk_id").data()
                          : -1;
-      pir::BuilderAttrGuard guard(builder, op_role, chunk_id);
+      std::string comp_op_name = op->name();
+      pir::BuilderAttrGuard guard(builder, op_role, chunk_id, comp_op_name);
 
       std::vector<std::vector<pir::Value>> decomp_res = call_decomp_rule(op);
       if (decomp_res.size() == 0) {
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -2370,6 +2370,11 @@ void BindUtils(pybind11::module *m) {
   m->def("set_op_role",
          [](int op_role) { ApiBuilder::Instance().SetOpRole(op_role); });
   m->def("get_op_role", []() { return ApiBuilder::Instance().GetOpRole(); });
+  m->def("set_comp_op_name", [](std::string comp_op_name) {
+    ApiBuilder::Instance().SetCompOpName(comp_op_name);
+  });
+  m->def("get_comp_op_name",
+         []() { return ApiBuilder::Instance().GetCompOpName(); });
   m->def("register_paddle_dialect", []() {
     pir::IrContext::Instance()
         ->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
diff --git a/paddle/pir/include/core/builder.h b/paddle/pir/include/core/builder.h
@@ -122,6 +122,13 @@ class Builder {
   /// Set/Get the chunk_id
   void set_chunk_id(int chunk_id) { chunk_id_ = chunk_id; }
   int chunk_id() const { return chunk_id_; }
+
+  /// Set/Get the comp_op_name
+  void set_comp_op_name(std::string comp_op_name) {
+    comp_op_name_ = comp_op_name;
+  }
+  std::string comp_op_name() const { return comp_op_name_; }
+
   IrContext *ir_context() const { return context_; }
 
   Block *block() const { return insertion_point_.first; }
@@ -183,6 +190,9 @@ class Builder {
   // bw, opt region.
   int op_role_ = -1;
   int chunk_id_ = -1;
+
+  // comp_op_name is used to specify the composite op name after decomposition
+  std::string comp_op_name_ = "";
 };
 
 template <typename OpTy, typename... Args>
@@ -195,7 +205,10 @@ OpTy Builder::Build(Args &&...args) {
 
 class BuilderAttrGuard {
  public:
-  BuilderAttrGuard(std::shared_ptr<Builder> builder, int op_role, int chunk_id);
+  BuilderAttrGuard(std::shared_ptr<Builder> builder,
+                   int op_role,
+                   int chunk_id,
+                   std::string comp_op_name);
 
   ~BuilderAttrGuard();
 
@@ -207,6 +220,7 @@ class BuilderAttrGuard {
   std::shared_ptr<Builder> builder_;
   int pre_op_role_ = -1;
   int pre_chunk_id_ = -1;
+  std::string pre_comp_op_name_ = "";
 };
 
 }  // namespace pir
diff --git a/paddle/pir/src/core/builder.cc b/paddle/pir/src/core/builder.cc
@@ -32,6 +32,12 @@ Operation *Builder::Build(OperationArgument &&argument) {
   if (chunk_id_ != -1) {
     op->set_attribute("chunk_id", Int32Attribute::get(context_, chunk_id_));
   }
+  if (comp_op_name_ != "") {
+    op->set_attribute("comp_op_name",
+                      StrAttribute::get(context_, comp_op_name_));
+    VLOG(6) << "Add comp_op_name: " << comp_op_name_
+            << " to op: " << op->name();
+  }
   return op;
 }
 
@@ -108,21 +114,27 @@ TensorNameAttribute Builder::tensor_name_attr(const std::string &value) {
 
 BuilderAttrGuard::BuilderAttrGuard(std::shared_ptr<Builder> builder,
                                    int op_role,
-                                   int chunk_id)
+                                   int chunk_id,
+                                   std::string comp_op_name)
     : builder_(builder),
       pre_op_role_(builder_->op_role()),
-      pre_chunk_id_(builder_->chunk_id()) {
+      pre_chunk_id_(builder_->chunk_id()),
+      pre_comp_op_name_(builder_->comp_op_name()) {
   if (pre_op_role_ != op_role) {
     builder_->set_op_role(op_role);
   }
   if (pre_chunk_id_ != chunk_id) {
     builder_->set_chunk_id(chunk_id);
   }
+  if (pre_comp_op_name_ != comp_op_name) {
+    builder_->set_comp_op_name(comp_op_name);
+  }
 }
 
 BuilderAttrGuard::~BuilderAttrGuard() {  // NOLINT
   builder_->set_op_role(pre_op_role_);
   builder_->set_chunk_id(pre_chunk_id_);
+  builder_->set_comp_op_name(pre_comp_op_name_);
 }
 
 }  // namespace pir
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -51,6 +51,7 @@
     warning_once,
     while_prune_check,
 )
+from paddle.base.framework import pir_op_name_guard
 from paddle.base.libpaddle.pir import (
     build_pipe_for_block,
     get_used_external_value,
@@ -908,7 +909,9 @@ def append_yield(
                         # create grad_op
 
                         before_ops_num = len(bwd_block.ops)
-                        with dynamic_shape_prim_vjp_guard(op, inputs):
+                        with dynamic_shape_prim_vjp_guard(
+                            op, inputs
+                        ), pir_op_name_guard(op.name() + '_grad'):
                             input_grads = paddle.framework.core.call_vjp(
                                 op,
                                 inputs,
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8506,3 +8506,16 @@ def pir_chunk_id_guard(chunk_id: int - 1) -> Generator[None, None, None]:
     finally:
         if paddle.framework.in_pir_mode():
             pir.set_chunk_id(original_chunk_id)
+
+
+@signature_safe_contextmanager
+def pir_op_name_guard(op_name: str) -> Generator[None, None, None]:
+
+    if paddle.framework.in_pir_mode() and core._is_bwd_prim_enabled():
+        original_comp_op_name = pir.get_comp_op_name()
+        pir.set_comp_op_name(op_name)
+    try:
+        yield
+    finally:
+        if paddle.framework.in_pir_mode() and core._is_bwd_prim_enabled():
+            pir.set_comp_op_name(original_comp_op_name)
diff --git a/python/paddle/decomposition/decomp.py b/python/paddle/decomposition/decomp.py
@@ -26,7 +26,11 @@
     has_decomp_rule,
     has_decomp_vjp,
 )
-from paddle.base.framework import pir_chunk_id_guard, pir_op_role_guard
+from paddle.base.framework import (
+    pir_chunk_id_guard,
+    pir_op_name_guard,
+    pir_op_role_guard,
+)
 from paddle.base.libpaddle.pir import Block, Operation
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.decomposition.recompute import DebugPrint, auto_recompute
@@ -869,9 +873,9 @@ def decompose_dist_program(pir_program):
                 ) and _check_prim_dynamic(op):
                     skip_decomp = True
                 if not skip_decomp:
-                    with pir_op_role_guard(op.op_role), pir_chunk_id_guard(
-                        op.chunk_id
-                    ):
+                    with pir_op_name_guard(op.name()), pir_op_role_guard(
+                        op.op_role
+                    ), pir_chunk_id_guard(op.chunk_id):
                         pir.set_insertion_point(op)
                         orig_outs = op.results()
 
diff --git a/python/paddle/pir/__init__.py b/python/paddle/pir/__init__.py
@@ -29,6 +29,7 @@
     create_shaped_type,
     fake_value,
     get_chunk_id,
+    get_comp_op_name,
     get_current_insertion_point,
     get_op_role,
     is_fake_value,
@@ -37,6 +38,7 @@
     reset_insertion_point_to_end,
     reset_insertion_point_to_start,
     set_chunk_id,
+    set_comp_op_name,
     set_insertion_point,
     set_insertion_point_after,
     set_insertion_point_to_block_end,
diff --git a/test/prim/pir_prim/test_prim_program.py b/test/prim/pir_prim/test_prim_program.py
@@ -99,5 +99,60 @@ def test_prim_all(self):
             np.testing.assert_allclose(ref, actual, rtol=1e-6)
 
 
+class TestCompOpName(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.shape_y = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "all":
+            core._set_prim_all_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            y = paddle.static.data('y', self.shape_y, dtype='float32')
+            x.stop_gradient = False
+            y.stop_gradient = False
+            divide_out = paddle.divide(x, y)
+            sum_out = paddle.mean(divide_out, axis=0)
+            [new_out] = decomp.decompose(main_program, [sum_out])
+            gradients = grad(new_out, (x, y))
+
+            exe = paddle.static.Executor()
+            [fwd, dx, dy] = exe.run(
+                feed={'x': self.x, 'y': self.y}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        if flag == "all":
+            core._set_prim_all_enabled(False)
+            assert (
+                'pd_op.mean' not in whole_ops
+                and 'pd_op.divide_grad' not in whole_ops
+            )
+        else:
+            assert (
+                'pd_op.mean' in whole_ops and 'pd_op.divide_grad' in whole_ops
+            )
+        return main_program
+
+    def test_set_comp_op_name(self):
+        decomp_program = self.base_net("all")
+        for op in decomp_program.global_block().ops:
+            if op.name() == 'pd_op.sum':
+                assert op.attrs()['comp_op_name'] == 'pd_op.mean'
+            if op.name() == 'pd_op.expand':
+                assert op.attrs()['comp_op_name'] == 'pd_op.sum_grad'
+        origin_program = self.base_net()
+        for op in decomp_program.global_block().ops:
+            if op.name() == 'pd_op.mean':
+                assert not op.has_attr('comp_op_name')
+            if op.name() == 'pd_op.sum_grad':
+                assert not op.has_attr('comp_op_name')
+
+
 if __name__ == "__main__":
     unittest.main()
