[CINN] Complete the op_role and chunk_id attributes of the operator in decompose process (#69612)

* refine code

* add ut

* refine

Author: zhangbo9674

paddle/fluid/primitive/base/decomp_trans.cc

@@ -491,8 +491,18 @@ void DecompProgram::decomp_block(
     if (enable_prim) {
       VLOG(4) << "[Prim] decomp op name " << op->name();
       check_decomp_dynamic_shape(op);
-      auto& builder = *(paddle::dialect::ApiBuilder::Instance().GetBuilder());
-      builder.set_insertion_point(op);
+      std::shared_ptr<pir::Builder> builder =
+          paddle::dialect::ApiBuilder::Instance().GetBuilder();
+      builder->set_insertion_point(op);
+
+      int op_role = (op->attribute<pir::Int32Attribute>("op_role"))
+                        ? op->attribute<pir::Int32Attribute>("op_role").data()
+                        : -1;
+      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::vector<std::vector<pir::Value>> decomp_res = call_decomp_rule(op);
       if (decomp_res.size() == 0) {
         // if we don't decomp this op, then leave it intact.
@@ -563,8 +573,9 @@ void DecompProgram::decomp_block(
       tar_vars[i] = src_vars_[i];
     }
   }
-  auto& builder = *(paddle::dialect::ApiBuilder::Instance().GetBuilder());
-  builder.SetInsertionPointToBlockEnd(block);
+  std::shared_ptr<pir::Builder> builder =
+      paddle::dialect::ApiBuilder::Instance().GetBuilder();
+  builder->SetInsertionPointToBlockEnd(block);
 }
 
 }  // namespace paddle

paddle/pir/include/core/builder.h

@@ -193,4 +193,20 @@ OpTy Builder::Build(Args &&...args) {
   return OpTy(op);
 }
 
+class BuilderAttrGuard {
+ public:
+  BuilderAttrGuard(std::shared_ptr<Builder> builder, int op_role, int chunk_id);
+
+  ~BuilderAttrGuard();
+
+  // forbid copy and operator=
+  BuilderAttrGuard(const BuilderAttrGuard &guard) = delete;
+  BuilderAttrGuard &operator=(const BuilderAttrGuard &guard) = delete;
+
+ private:
+  std::shared_ptr<Builder> builder_;
+  int pre_op_role_ = -1;
+  int pre_chunk_id_ = -1;
+};
+
 }  // namespace pir

paddle/pir/src/core/builder.cc

@@ -106,4 +106,23 @@ TensorNameAttribute Builder::tensor_name_attr(const std::string &value) {
   return TensorNameAttribute::get(context_, value);
 }
 
+BuilderAttrGuard::BuilderAttrGuard(std::shared_ptr<Builder> builder,
+                                   int op_role,
+                                   int chunk_id)
+    : builder_(builder),
+      pre_op_role_(builder_->op_role()),
+      pre_chunk_id_(builder_->chunk_id()) {
+  if (pre_op_role_ != op_role) {
+    builder_->set_op_role(op_role);
+  }
+  if (pre_chunk_id_ != chunk_id) {
+    builder_->set_chunk_id(chunk_id);
+  }
+}
+
+BuilderAttrGuard::~BuilderAttrGuard() {  // NOLINT
+  builder_->set_op_role(pre_op_role_);
+  builder_->set_chunk_id(pre_chunk_id_);
+}
+
 }  // namespace pir

test/prim/pir_prim/test_decomp_op.py

@@ -38,6 +38,11 @@ def get_ir_program():
             y_s = paddle.mean(y_s)
             y_s = paddle.tanh(y_s)
         pir_program = pir.translate_to_pir(main_program.desc)
+
+        all_ops = pir_program.global_block().ops
+        for op in all_ops:
+            op.op_role = 1
+
         return pir_program
 
 
@@ -68,6 +73,10 @@ def test_build_op(self):
                     'pd_op.tanh',
                 ],
             )
+            op_role_list = [op.op_role for op in pir_program.global_block().ops]
+            self.assertEqual(
+                all(op_role == 1 for op_role in op_role_list), True
+            )
 
 
 if __name__ == "__main__":