support struct_name for pir (#67174)

AndSonder · web-flow · commit c1a5199be229 · 2024-08-09T11:04:00.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -853,6 +853,11 @@ void BindOperation(py::module *m) {
                                                 phi::IntArray(val));
              self.set_attribute(attr_name, attr);
            })
+      .def("set_str_attr",
+           [](Operation &self, std::string &attr_name, std::string &val) {
+             self.set_attribute(
+                 attr_name, StrAttribute::get(pir::IrContext::Instance(), val));
+           })
       .def("attrs",
            [](Operation &self) -> py::dict {
              py::dict attrs_dict;
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -1257,9 +1257,25 @@ def name_struct(prefix=None):
         assert prefix, "namescope prefix can not be empty."
         global _name_struct
         _name_struct = _name_struct.child(prefix)
+        if in_pir_mode():
+            op_num_before = len(
+                paddle.static.default_main_program().global_block().ops
+            )
         try:
             yield
         finally:
+            if in_pir_mode():
+                all_ops = (
+                    paddle.static.default_main_program().global_block().ops
+                )
+                op_num = len(all_ops)
+
+                for idx in reversed(range(op_num_before, op_num)):
+                    op = all_ops[idx]
+                    if op.has_attr("struct_name"):
+                        break
+                    op.set_str_attr("struct_name", _full_name_struct())
+
             _name_struct = _name_struct.parent()
 
 
