[PIR] Add check before fetch attr in map (#57284)

chen2016013 · web-flow · commit b452fb81cd45 · 2023-09-14T10:21:41.000+08:00
* Add check before fetch attr in map

* Add check before fetch attr in map

* change error text

* remove wrong changes which should be in another PR

* change text format

* fix code style

* fix code style
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -669,21 +669,41 @@ def gen_build_func_str(
     )
 
     GET_ATTRIBUTES_FROM_MAP_TEMPLATE = """
+  PADDLE_ENFORCE(
+      attributes.find("{attribute_name}") != attributes.end(),
+      phi::errors::NotFound(
+          "'{attribute_name}' Attribute is expected for {op_name}. "));
   {attr_type} {attribute_name} = attributes.at("{attribute_name}").dyn_cast<{attr_ir_type}>().data();
 """
     GET_STR_ATTRIBUTES_FROM_MAP_TEMPLATE = """
+  PADDLE_ENFORCE(
+      attributes.find("{attribute_name}") != attributes.end(),
+      phi::errors::NotFound(
+          "'{attribute_name}' Attribute is expected for {op_name}. "));
   {attr_type} {attribute_name} = attributes.at("{attribute_name}").dyn_cast<pir::StrAttribute>().AsString();
 """
     GET_ARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE = """
+  PADDLE_ENFORCE(
+      attributes.find("{attribute_name}") != attributes.end(),
+      phi::errors::NotFound(
+          "'{attribute_name}' Attribute is expected for {op_name}. "));
   {attr_type} {attribute_name};
   for (size_t i = 0; i < attributes.at("{attribute_name}").dyn_cast<pir::ArrayAttribute>().size(); i++) {{
     {attribute_name}.push_back(attributes.at("{attribute_name}").dyn_cast<pir::ArrayAttribute>().at(i).dyn_cast<{inner_type}>().{data_name}());
   }}
 """
     GET_INTARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE = """
+  PADDLE_ENFORCE(
+      attributes.find("{attribute_name}") != attributes.end(),
+      phi::errors::NotFound(
+          "'{attribute_name}' Attribute is expected for {op_name}. "));
   {attr_type} {attribute_name} = attributes.at("{attribute_name}").dyn_cast<paddle::dialect::IntArrayAttribute>().data().GetData();
 """
     GET_SCALAR_ATTRIBUTE_FROM_MAP_TEMPLATE = """
+  PADDLE_ENFORCE(
+      attributes.find("{attribute_name}") != attributes.end(),
+      phi::errors::NotFound(
+          "'{attribute_name}' Attribute is expected for {op_name}. "));
   {attr_type} {attribute_name} = attributes.at("{attribute_name}").dyn_cast<paddle::dialect::ScalarAttribute>().data().to<{attr_type}>();
 """
 
@@ -706,6 +726,7 @@ def gen_build_func_str(
                     data_name = "AsString"
                 get_attributes_str += (
                     GET_ARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
+                        op_name=op_class_name,
                         attr_type=attr_type,
                         attribute_name=op_attribute_name_list[idx],
                         inner_type=inner_type,
@@ -718,6 +739,7 @@ def gen_build_func_str(
             ):
                 get_attributes_str += (
                     GET_INTARRAY_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
+                        op_name=op_class_name,
                         attr_type=attr_type,
                         attribute_name=op_attribute_name_list[idx],
                     )
@@ -728,20 +750,23 @@ def gen_build_func_str(
             ):
                 get_attributes_str += (
                     GET_SCALAR_ATTRIBUTE_FROM_MAP_TEMPLATE.format(
+                        op_name=op_class_name,
                         attr_type=attr_type,
                         attribute_name=op_attribute_name_list[idx],
                     )
                 )
             elif "pir::StrAttribute" in op_attribute_type_list[idx]:
                 get_attributes_str += (
                     GET_STR_ATTRIBUTES_FROM_MAP_TEMPLATE.format(
+                        op_name=op_class_name,
                         attr_type=attr_type,
                         attribute_name=op_attribute_name_list[idx],
                         attr_ir_type=op_attribute_type_list[idx],
                     )
                 )
             else:
                 get_attributes_str += GET_ATTRIBUTES_FROM_MAP_TEMPLATE.format(
+                    op_name=op_class_name,
                     attr_type=attr_type,
                     attribute_name=op_attribute_name_list[idx],
                     attr_ir_type=op_attribute_type_list[idx],
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -455,10 +455,22 @@ void FusedGemmEpilogueOp::Build(pir::Builder &builder,
                                 pir::OpResult y_,
                                 pir::OpResult bias_,
                                 pir::AttributeMap attributes) {
+  PADDLE_ENFORCE(
+      attributes.find("trans_x") != attributes.end(),
+      phi::errors::NotFound(
+          "'trans_x' Attribute is expected for FusedGemmEpilogueOp"));
   bool trans_x = attributes.at("trans_x").dyn_cast<pir::BoolAttribute>().data();
 
+  PADDLE_ENFORCE(
+      attributes.find("trans_y") != attributes.end(),
+      phi::errors::NotFound(
+          "'trans_y' Attribute is expected for FusedGemmEpilogueOp"));
   bool trans_y = attributes.at("trans_y").dyn_cast<pir::BoolAttribute>().data();
 
+  PADDLE_ENFORCE(
+      attributes.find("activation") != attributes.end(),
+      phi::errors::NotFound(
+          "'activation' Attribute is expected for FusedGemmEpilogueOp"));
   std::string activation =
       attributes.at("activation").dyn_cast<pir::StrAttribute>().AsString();
 
@@ -697,10 +709,22 @@ void FusedGemmEpilogueGradOp::Build(pir::Builder &builder,
                                     pir::OpResult reserve_space_,
                                     pir::OpResult out_grad_,
                                     pir::AttributeMap attributes) {
+  PADDLE_ENFORCE(
+      attributes.find("trans_x") != attributes.end(),
+      phi::errors::NotFound(
+          "'trans_x' Attribute is expected for FusedGemmEpilogueGradOp"));
   bool trans_x = attributes.at("trans_x").dyn_cast<pir::BoolAttribute>().data();
 
+  PADDLE_ENFORCE(
+      attributes.find("trans_y") != attributes.end(),
+      phi::errors::NotFound(
+          "'trans_y' Attribute is expected for FusedGemmEpilogueGradOp"));
   bool trans_y = attributes.at("trans_y").dyn_cast<pir::BoolAttribute>().data();
 
+  PADDLE_ENFORCE(
+      attributes.find("activation_grad") != attributes.end(),
+      phi::errors::NotFound("'activation_grad' Attribute is expected for"
+                            "FusedGemmEpilogueGradOp"));
   std::string activation_grad =
       attributes.at("activation_grad").dyn_cast<pir::StrAttribute>().AsString();
 
