From d14cfe10341681d18edf05ac98da2c5241b0864e Mon Sep 17 00:00:00 2001
From: Mike Urbach <mikeurbach@gmail.com>
Date: Wed, 23 Sep 2020 18:01:39 +0000
Subject: [PATCH] [mlir][OpFormatGen] Update "custom" directives for
 attributes.

This tweaks the generated code for parsing attributes with a custom
directive to call `addAttribute` on the `OperationState` directly,
and adds a newline after this call. Previously, the generated code
would call `addAttribute` on the `OperationState` field `attributes`,
which has no such method and fails to compile. Furthermore, the lack
of newline would generate code with incorrectly formatted single line
`if` statements. Added tests for parsing and printing attributes with
a custom directive.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D87860
---
 mlir/test/lib/Dialect/Test/TestDialect.cpp | 18 ++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td      | 11 +++++++++++
 mlir/test/mlir-tblgen/op-format.mlir       |  6 ++++++
 mlir/tools/mlir-tblgen/OpFormatGen.cpp     |  4 ++--
 4 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 7bdd0fd3815b..c84a7717abe7 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -364,6 +364,17 @@ parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,
   varSuccessors.append(2, varSuccessor);
   return success();
 }
+static ParseResult parseCustomDirectiveAttributes(OpAsmParser &parser,
+                                                  IntegerAttr &attr,
+                                                  IntegerAttr &optAttr) {
+  if (parser.parseAttribute(attr))
+    return failure();
+  if (succeeded(parser.parseOptionalComma())) {
+    if (parser.parseAttribute(optAttr))
+      return failure();
+  }
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // Printing
@@ -417,6 +428,13 @@ static void printCustomDirectiveSuccessors(OpAsmPrinter &printer,
   if (!varSuccessors.empty())
     printer << ", " << varSuccessors.front();
 }
+static void printCustomDirectiveAttributes(OpAsmPrinter &printer,
+                                           Attribute attribute,
+                                           Attribute optAttribute) {
+  printer << attribute;
+  if (optAttribute)
+    printer << ", " << optAttribute;
+}
 
 //===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index a2cf5b0d28e1..6d4e58c70c02 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1547,6 +1547,17 @@ def FormatCustomDirectiveSuccessors
   }];
 }
 
+def FormatCustomDirectiveAttributes
+    : TEST_Op<"format_custom_directive_attributes"> {
+  let arguments = (ins I64Attr:$attr, OptionalAttr<I64Attr>:$optAttr);
+  let assemblyFormat = [{
+    custom<CustomDirectiveAttributes>(
+      $attr, $optAttr
+    )
+    attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // AllTypesMatch type inference
 
diff --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index 5066fe5a24e6..57360e31ebfb 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -213,6 +213,12 @@ test.format_custom_directive_operands_and_types %i64, %i64 -> (%i64) : i64, i64
 // CHECK: test.format_custom_directive_operands_and_types %[[I64]] -> (%[[I64]]) : i64 -> (i64)
 test.format_custom_directive_operands_and_types %i64 -> (%i64) : i64 -> (i64)
 
+// CHECK: test.format_custom_directive_attributes 54 : i64
+test.format_custom_directive_attributes 54 : i64
+
+// CHECK: test.format_custom_directive_attributes 54 : i64, 46 : i64
+test.format_custom_directive_attributes 54 : i64, 46 : i64
+
 // CHECK: test.format_custom_directive_regions {
 // CHECK-NEXT: test.return
 // CHECK-NEXT: }
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 7658963d6f9b..a910d95d52c9 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -943,8 +943,8 @@ static void genCustomDirectiveParser(CustomDirective *dir, OpMethodBody &body) {
       if (var->attr.isOptional())
         body << llvm::formatv("    if ({0}Attr)\n  ", var->name);
 
-      body << llvm::formatv(
-          "    result.attributes.addAttribute(\"{0}\", {0}Attr);", var->name);
+      body << llvm::formatv("    result.addAttribute(\"{0}\", {0}Attr);\n",
+                            var->name);
     } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
       const NamedTypeConstraint *var = operand->getVar();
       if (!var->isOptional())