Refine hardcoded attr setting/getting in ir (#8420)

* use names in trait static func * more changes on op name attr * use wrapped func
Oneflow-Inc · Jun 14, 2022 · b58da14 · b58da14
1 parent bed02f4
commit b58da14
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 24 deletions.
diff --git a/oneflow/ir/lib/OneFlow/OneFlowOpFolders.cpp b/oneflow/ir/lib/OneFlow/OneFlowOpFolders.cpp
@@ -51,10 +51,11 @@ OpFoldResult UnaryFold(MLIRContext* ctx, ArrayRef<Attribute> operands,
   const auto attr_dict = operands.front().cast<mlir::DictionaryAttr>();
   auto attrs = NamedAttrList(attr_dict);
   const auto tensor = support::DenseElementsAttrToTensor(
-      attr_dict.get("value"), attr_dict.get("device_tag"), attr_dict.get("device_name"));
+      attr_dict.get("value"), attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceTagAttr()),
+      attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceNameAttr()));
   const auto result = f(tensor).GetPtrOrThrow();
   attrs.set("value", support::TensorToDenseElementsAttr(result, ctx));
-  attrs.set("op_name", GenNewVariableOpName(ctx));
+  attrs.set(OpTrait::IsOpConfCompatible<void>::getOpNameAttr(), GenNewVariableOpName(ctx));
 
   return attrs.getDictionary(ctx);
 }
@@ -67,17 +68,19 @@ OpFoldResult BinaryFold(MLIRContext* ctx, ArrayRef<Attribute> operands,
   auto rhs_attr_dict = operands.back().cast<mlir::DictionaryAttr>();
 
   auto attrs = NamedAttrList(lhs_attr_dict);
-  const auto lhs_tensor = support::DenseElementsAttrToTensor(lhs_attr_dict.get("value"),
-                                                             lhs_attr_dict.get("device_tag"),
-                                                             lhs_attr_dict.get("device_name"));
-  const auto rhs_tensor = support::DenseElementsAttrToTensor(rhs_attr_dict.get("value"),
-                                                             rhs_attr_dict.get("device_tag"),
-                                                             rhs_attr_dict.get("device_name"));
+  const auto lhs_tensor = support::DenseElementsAttrToTensor(
+      lhs_attr_dict.get("value"),
+      lhs_attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceTagAttr()),
+      lhs_attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceNameAttr()));
+  const auto rhs_tensor = support::DenseElementsAttrToTensor(
+      rhs_attr_dict.get("value"),
+      rhs_attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceTagAttr()),
+      rhs_attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceNameAttr()));
 
   const auto result = f(lhs_tensor, rhs_tensor).GetPtrOrThrow();
 
   attrs.set("value", support::TensorToDenseElementsAttr(result, ctx));
-  attrs.set("op_name", GenNewVariableOpName(ctx));
+  attrs.set(OpTrait::IsOpConfCompatible<void>::getOpNameAttr(), GenNewVariableOpName(ctx));
 
   return attrs.getDictionary(ctx);
 }

diff --git a/oneflow/ir/lib/OneFlow/Passes.cpp b/oneflow/ir/lib/OneFlow/Passes.cpp
@@ -349,9 +349,9 @@ ::llvm::SmallVector<::mlir::Value, 4> CreateConv2dAndErasePad(::mlir::PatternRew
 
 NamedAttrList GetUserOpCommonAttrs(MLIRContext* ctx, const std::string& op_name) {
   NamedAttrList attrs;
-  attrs.set("op_name", StringAttr::get(ctx, op_name));
-  attrs.set("device_tag", StringAttr::get(ctx, "cpu"));
-  attrs.set("device_name",
+  attrs.set(OpTrait::IsOpConfCompatible<void>::getOpNameAttr(), StringAttr::get(ctx, op_name));
+  attrs.set(OpTrait::IsOpConfCompatible<void>::getDeviceTagAttr(), StringAttr::get(ctx, "cpu"));
+  attrs.set(OpTrait::IsOpConfCompatible<void>::getDeviceNameAttr(),
             ArrayAttr::get(ctx, llvm::to_vector<8>(llvm::map_range(ArrayRef<StringRef>({"@0:0"}),
                                                                    [&](StringRef v) -> Attribute {
                                                                      return StringAttr::get(ctx, v);
@@ -569,7 +569,8 @@ llvm::SmallVector<mlir::Value, 4> getInputOperandTransposeOp(NCHWCompatible op,
                                                              PatternRewriter& rewriter) {
   std::string transpose_name = OpTrait::IsOpConfCompatible<void>::getOpName(op).str()
                                + "_transpose_input_" + std::to_string(num_transposed_operand);
-  transpose_attributes.set(llvm::StringRef("op_name"), rewriter.getStringAttr(transpose_name));
+  transpose_attributes.set(llvm::StringRef(OpTrait::IsOpConfCompatible<void>::getOpNameAttr()),
+                           rewriter.getStringAttr(transpose_name));
   SmallVector<Value, 4> input_operands;
   input_operands.push_back(val);
   auto res = rewriter
@@ -583,7 +584,8 @@ TransposeOp getResultTransposeOp(NCHWCompatible op, Value val, NamedAttrList tra
                                  int num_transposed_result, PatternRewriter& rewriter) {
   std::string transpose_name = OpTrait::IsOpConfCompatible<void>::getOpName(op).str()
                                + "_transpose_output_" + std::to_string(num_transposed_result);
-  transpose_attributes.set(llvm::StringRef("op_name"), rewriter.getStringAttr(transpose_name));
+  transpose_attributes.set(llvm::StringRef(OpTrait::IsOpConfCompatible<void>::getOpNameAttr()),
+                           rewriter.getStringAttr(transpose_name));
   SmallVector<Value, 4> operands;
   operands.push_back(val);
   TransposeOp transpose_op = rewriter.create<oneflow::TransposeOp>(op.getLoc(), val.getType(),

diff --git a/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp b/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp
@@ -492,10 +492,7 @@ LogicalResult ConvertCtrlInputs(Operation* op, ::oneflow::OperatorConf& op_conf)
   if (auto ctrl_ins = GetCtrlIntputOperands(op)) {
     for (auto ctrl_in : ctrl_ins.getValue()) {
       op_conf.add_ctrl_in_op_name(
-          ctrl_in.getDefiningOp()
-              ->getAttrOfType<StringAttr>(OpTrait::IsOpConfCompatible<void>::getOpNameAttr())
-              .getValue()
-              .str());
+          OpTrait::IsOpConfCompatible<void>::getOpName(ctrl_in.getDefiningOp()).str());
     }
   }
   return success();
@@ -675,9 +672,8 @@ llvm::Optional<std::string> GetOutputLbn(OpResult result) {
       auto size = std::get<1>(name_size_tuple);
       if ((size_sum + size) > result_number) {
         const uint32_t bn_i = result_number - size_sum;
-        return def_op->getAttrOfType<StringAttr>(OpTrait::IsOpConfCompatible<void>::getOpNameAttr())
-                   .str()
-               + "/" + name + "_" + std::to_string(bn_i);
+        return OpTrait::IsOpConfCompatible<void>::getOpName(def_op).str() + "/" + name + "_"
+               + std::to_string(bn_i);
       }
       size_sum += size;
     }
@@ -946,7 +942,7 @@ LogicalResult ConvertVariableOpConf(VariableOp op, ::oneflow::OperatorConf* op_c
   // all operands are ctrl_inputs
   for (const auto& operand : op->getOperands()) {
     op_conf->add_ctrl_in_op_name(
-        operand.getDefiningOp()->getAttrOfType<StringAttr>("op_name").getValue().str());
+        OpTrait::IsOpConfCompatible<void>::getOpName(operand.getDefiningOp()).str());
   }
   if (auto floatInit = op.float_initializer()) {
     var_op_conf->mutable_initializer()->mutable_constant_conf()->set_value(
@@ -1002,7 +998,7 @@ LogicalResult ConvertInputOpConf(InputOp op, ::oneflow::OperatorConf* op_conf) {
   // operand 0 is block argument, others are ctrl_inputs
   for (size_t i = 1; i < op->getNumOperands(); ++i) {
     op_conf->add_ctrl_in_op_name(
-        op->getOperand(i).getDefiningOp()->getAttrOfType<StringAttr>("op_name").getValue().str());
+        OpTrait::IsOpConfCompatible<void>::getOpName(op->getOperand(i).getDefiningOp()).str());
   }
 
   return success();
@@ -1054,7 +1050,7 @@ LogicalResult ConvertOutputOpConf(OutputOp op, ::oneflow::OperatorConf* op_conf)
   output_op_conf->set_in(output_lbn);
   for (size_t i = 1; i < op->getNumOperands(); ++i) {
     op_conf->add_ctrl_in_op_name(
-        op->getOperand(i).getDefiningOp()->getAttrOfType<StringAttr>("op_name").getValue().str());
+        OpTrait::IsOpConfCompatible<void>::getOpName(op->getOperand(i).getDefiningOp()).str());
   }
   return success();
 }