[mlir][OpFormatGen] Add support for optional enum attributes

[rriddle]

Summary

The check for formatting enum attributes was missing a call to get the base attribute, which is necessary to strip off the top-level OptionalAttr<> wrapper.

[jpienaar]

Ha ran into this last week, thanks!

mlir/test/Target/llvmir.mlir

llvm.func @useInlineAsm(%arg0: !llvm.i32) {
 
   // CHECK-NEXT:  call i8 asm sideeffect "foo", "=r,r,r"(i32 {{.*}}, i32 {{.*}}), !dbg !11
   %2 = llvm.inline_asm has_side_effects "foo", "=r,r,r" %arg0, %arg0 : (!llvm.i32, !llvm.i32) -> !llvm.i8
 
   // CHECK-NEXT:  call i8 asm alignstack "foo", "=r,r,r"(i32 {{.*}}, i32 {{.*}}), !dbg !12
   %3 = llvm.inline_asm is_align_stack "foo", "=r,r,r" %arg0, %arg0 : (!llvm.i32, !llvm.i32) -> !llvm.i8
 
   // CHECK-NEXT:  call i8 asm inteldialect "foo", "=r,r,r"(i32 {{.*}}, i32 {{.*}}), !dbg !13
-  %4 = llvm.inline_asm asm_dialect = 1 "foo", "=r,r,r" %arg0, %arg0 : (!llvm.i32, !llvm.i32) -> !llvm.i8
+  %4 = llvm.inline_asm asm_dialect = "intel" "foo", "=r,r,r" %arg0, %arg0 : (!llvm.i32, !llvm.i32) -> !llvm.i8
 
   // CHECK-NEXT:  call { i8, i8 } asm "foo", "=r,=r,r"(i32 {{.*}}), !dbg !14
   %5 = llvm.inline_asm "foo", "=r,=r,r" %arg0 : (!llvm.i32) -> !llvm.struct<(i8, i8)>
 
   llvm.return
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def FormatOptionalUnitAttrNoElide
     : TEST_Op<"format_optional_unit_attribute_no_elide"> {
   let arguments = (ins UnitAttr:$is_optional);
   let assemblyFormat = "($is_optional^)? attr-dict";
 }
 
+def FormatOptionalEnumAttr : TEST_Op<"format_optional_enum_attr"> {
+  let arguments = (ins OptionalAttr<SomeI64Enum>:$attr);
+  let assemblyFormat = "($attr^)? attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // Custom Directives
 
 def FormatCustomDirectiveOperands
     : TEST_Op<"format_custom_directive_operands", [AttrSizedOperandSegments]> {
   let arguments = (ins I64:$operand, Optional<I64>:$optOperand,
                        Variadic<I64>:$varOperands);
   let assemblyFormat = [{

mlir/test/mlir-tblgen/op-format.mlir

 
 // CHECK: test.format_optional_unit_attribute
 // CHECK-NOT: is_optional
 test.format_optional_unit_attribute
 
 // CHECK: test.format_optional_unit_attribute_no_elide unit
 test.format_optional_unit_attribute_no_elide unit
 
+// CHECK: test.format_optional_enum_attr "case5"
+test.format_optional_enum_attr "case5"
+
+// CHECK: test.format_optional_enum_attr
+// CHECK-NOT: "case5"
+test.format_optional_enum_attr
+
 //===----------------------------------------------------------------------===//
 // Format optional operands and results
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_optional_operand_result_a_op(%[[I64]] : i64) : i64
 test.format_optional_operand_result_a_op(%i64 : i64) : i64
 
 // CHECK: test.format_optional_operand_result_a_op( : ) : i64

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // Parser Gen
 
 /// Returns true if we can format the given attribute as an EnumAttr in the
 /// parser format.
 static bool canFormatEnumAttr(const NamedAttribute *attr) {
-  const EnumAttr *enumAttr = dyn_cast<EnumAttr>(&attr->attr);
+  Attribute baseAttr = attr->attr.getBaseAttr();
+  const EnumAttr *enumAttr = dyn_cast<EnumAttr>(&baseAttr);
   if (!enumAttr)
     return false;
 
   // The attribute must have a valid underlying type and a constant builder.
   return !enumAttr->getUnderlyingType().empty() &&
          !enumAttr->getConstBuilderTemplate().empty();
 }
 
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     // Nothing to parse.
 
     /// Arguments.
   } else if (auto *attr = dyn_cast<AttributeVariable>(element)) {
     const NamedAttribute *var = attr->getVar();
 
     // Check to see if we can parse this as an enum attribute.
     if (canFormatEnumAttr(var)) {
-      const EnumAttr &enumAttr = cast<EnumAttr>(var->attr);
+      Attribute baseAttr = var->attr.getBaseAttr();
+      const EnumAttr &enumAttr = cast<EnumAttr>(baseAttr);
 
       // Generate the code for building an attribute for this enum.
       std::string attrBuilderStr;
       {
         llvm::raw_string_ostream os(attrBuilderStr);
         os << tgfmt(enumAttr.getConstBuilderTemplate(), &attrTypeCtx,
                     "attrOptional.getValue()");
       }
void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
   lastWasPunctuation = false;
   shouldEmitSpace = true;
 
   if (auto *attr = dyn_cast<AttributeVariable>(element)) {
     const NamedAttribute *var = attr->getVar();
 
     // If we are formatting as an enum, symbolize the attribute as a string.
     if (canFormatEnumAttr(var)) {
-      const EnumAttr &enumAttr = cast<EnumAttr>(var->attr);
+      Attribute baseAttr = var->attr.getBaseAttr();
+      const EnumAttr &enumAttr = cast<EnumAttr>(baseAttr);
       body << "  p << '\"' << " << enumAttr.getSymbolToStringFnName() << "("
-           << var->name << "()) << '\"';\n";
+           << (var->attr.isOptional() ? "*" : "") << var->name
+           << "()) << '\"';\n";
       return;
     }
 
     // If we are formatting as a symbol name, handle it as a symbol name.
     if (shouldFormatSymbolNameAttr(var)) {
       body << "  p.printSymbolName(" << var->name << "Attr().getValue());\n";
       return;
     }