[mlir][OpFormat] Add support for emitting newlines from the custom format of an operation

[rriddle]

Summary

This revision adds a new printNewline hook to OpAsmPrinter that allows for printing a newline within the custom format of an operation, that is then indented to the start of the operation. Support for the declarative assembly format is also added, in the form of a \n literal.

[mehdi_amini]

That will be really neat for a few ops out thread :)

[mehdi_amini]

In D93151#2450041, @mehdi_amini wrote:

That will be really neat for a few ops out thread :)

I meant to LGTM

mlir/docs/OpDefinitions.md

 #### Literals
 
 A literal is either a keyword or punctuation surrounded by \`\`.
 
 The following are the set of valid punctuation:
 
 `:`, `,`, `=`, `<`, `>`, `(`, `)`, `{`, `}`, `[`, `]`, `->`, `?`, `+`, `*`
 
+The following are valid whitespace punctuation:
+
+`\n`, ` `
+
+The `\n` literal emits a newline an indents to the start of the operation. An
+example is shown below:
+
+```tablegen
+let assemblyFormat = [{
+  `{` `\n` ` ` ` ` `this_is_on_a_newline` `\n` `}` attr-dict
+}];
+```
+
+```mlir
+%results = my.operation {
+  this_is_on_a_newline
+}
+```
+
+An empty literal \`\` may be used to remove a space that is inserted implicitly
+after certain literal elements, such as `)`/`]`/etc. For example, "`]`" may
+result in an output of `]` it is not the last element in the format. "`]` \`\`"
+would trim the trailing space in this situation.
+
 #### Variables
 
 A variable is an entity that has been registered on the operation itself, i.e.
 an argument(attribute or operand), region, result, successor, etc. In the
 `CallOp` example above, the variables would be `$callee` and `$args`.
 
 Attribute variables are printed with their respective value type, unless that
 value type is buildable. In those cases, the type of the attribute is elided.

mlir/include/mlir/IR/OpImplementation.h

 /// This is a pure-virtual base class that exposes the asmprinter hooks
 /// necessary to implement a custom print() method.
 class OpAsmPrinter {
 public:
   OpAsmPrinter() {}
   virtual ~OpAsmPrinter();
   virtual raw_ostream &getStream() const = 0;
 
+  /// Print a newline and indent the printer to the start of the current
+  /// operation.
+  virtual void printNewline() = 0;
+
   /// Print implementations for various things an operation contains.
   virtual void printOperand(Value value) = 0;
   virtual void printOperand(Value value, raw_ostream &os) = 0;
 
   /// Print a comma separated list of operands.
   template <typename ContainerType>
   void printOperands(const ContainerType &container) {
     printOperands(container.begin(), container.end());

mlir/lib/IR/AsmPrinter.cpp

private:
   }
 
   /// Return 'nulls' as the output stream, this will ignore any data fed to it.
   raw_ostream &getStream() const override { return llvm::nulls(); }
 
   /// The following are hooks of `OpAsmPrinter` that are not necessary for
   /// determining potential aliases.
   void printAffineMapOfSSAIds(AffineMapAttr, ValueRange) override {}
+  void printNewline() override {}
   void printOperand(Value) override {}
   void printOperand(Value, raw_ostream &os) override {
     // Users expect the output string to have at least the prefixed % to signal
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
   void printSymbolName(StringRef) override {}
public:
 
   //===--------------------------------------------------------------------===//
   // OpAsmPrinter methods
   //===--------------------------------------------------------------------===//
 
   /// Return the current stream of the printer.
   raw_ostream &getStream() const override { return os; }
 
+  /// Print a newline and indent the printer to the start of the current
+  /// operation.
+  void printNewline() override {
+    os << newLine;
+    os.indent(currentIndent);
+  }
+
   /// Print the given type.
   void printType(Type type) override { ModulePrinter::printType(type); }
 
   /// Print the given attribute.
   void printAttribute(Attribute attr) override {
     ModulePrinter::printAttribute(attr);
   }
 

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

 }
 
 //===----------------------------------------------------------------------===//
 // Test Op Asm Format
 //===----------------------------------------------------------------------===//
 
 def FormatLiteralOp : TEST_Op<"format_literal_op"> {
   let assemblyFormat = [{
-    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` `` `(` ` ` `)` `?` `+` `*` attr-dict
+    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` `` `(` ` ` `)`
+    `?` `+` `*` `{` `\n` `}` attr-dict
   }];
 }
 
 // Test that we elide attributes that are within the syntax.
 def FormatAttrOp : TEST_Op<"format_attr_op"> {
   let arguments = (ins I64Attr:$attr);
   let assemblyFormat = "$attr attr-dict";
 }

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

 }]>;
 // CHECK: error: unexpected end of file in literal
 // CHECK: error: expected directive, literal, variable, or optional group
 def LiteralInvalidB : TestFormat_Op<"literal_invalid_b", [{
   `
 }]>;
 // CHECK-NOT: error
 def LiteralValid : TestFormat_Op<"literal_valid", [{
-  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` `?` `+` `*` ` ` `` `->` `abc$._`
+  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` `?` `+` `*` ` ` `` `->` `\n` `abc$._`
   attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: optional groups can only be used as top-level elements

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

 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -verify-diagnostics | FileCheck %s
 
 // CHECK: %[[I64:.*]] =
 %i64 = "foo.op"() : () -> (i64)
 // CHECK: %[[I32:.*]] =
 %i32 = "foo.op"() : () -> (i32)
 // CHECK: %[[MEMREF:.*]] =
 %memref = "foo.op"() : () -> (memref<1xf64>)
 
-// CHECK: test.format_literal_op keyword_$. -> :, = <> () []( ) ? + * {foo.some_attr}
-test.format_literal_op keyword_$. -> :, = <> () []( ) ? + * {foo.some_attr}
+// CHECK: test.format_literal_op keyword_$. -> :, = <> () []( ) ? + * {
+// CHECK-NEXT: } {foo.some_attr}
+test.format_literal_op keyword_$. -> :, = <> () []( ) ? + * {
+} {foo.some_attr}
 
 // CHECK: test.format_attr_op 10
 // CHECK-NOT: {attr
 test.format_attr_op 10
 
 // CHECK: test.format_opt_attr_op_a(10)
 // CHECK-NOT: {opt_attr
 test.format_opt_attr_op_a(10)

mlir/tools/mlir-tblgen/OpFormatGen.cpp

enum class Kind {
     ResultsDirective,
     SuccessorsDirective,
     TypeDirective,
     TypeRefDirective,
 
     /// This element is a literal.
     Literal,
 
-    /// This element prints or omits a space. It is ignored by the parser.
+    /// This element is a whitespace.
+    Newline,
     Space,
 
     /// This element is an variable value.
     AttributeVariable,
     OperandVariable,
     RegionVariable,
     ResultVariable,
     SuccessorVariable,
bool LiteralElement::isValidLiteral(StringRef value) {
   if (!isalpha(front) && front != '_')
     return false;
   return llvm::all_of(value.drop_front(), [](char c) {
     return isalnum(c) || c == '_' || c == '$' || c == '.';
   });
 }
 
 //===----------------------------------------------------------------------===//
-// SpaceElement
+// WhitespaceElement
 
 namespace {
+/// This class represents a whitespace element, e.g. newline or space. It's a
+/// literal that is printed but never parsed.
+class WhitespaceElement : public Element {
+public:
+  WhitespaceElement(Kind kind) : Element{kind} {}
+  static bool classof(const Element *element) {
+    Kind kind = element->getKind();
+    return kind == Kind::Newline || kind == Kind::Space;
+  }
+};
+
+/// This class represents an instance of a newline element. It's a literal that
+/// prints a newline. It is ignored by the parser.
+class NewlineElement : public WhitespaceElement {
+public:
+  NewlineElement() : WhitespaceElement(Kind::Newline) {}
+  static bool classof(const Element *element) {
+    return element->getKind() == Kind::Newline;
+  }
+};
+
 /// This class represents an instance of a space element. It's a literal that
 /// prints or omits printing a space. It is ignored by the parser.
-class SpaceElement : public Element {
+class SpaceElement : public WhitespaceElement {
 public:
-  SpaceElement(bool value) : Element{Kind::Space}, value(value) {}
+  SpaceElement(bool value) : WhitespaceElement(Kind::Space), value(value) {}
   static bool classof(const Element *element) {
     return element->getKind() == Kind::Space;
   }
 
   /// Returns true if this element should print as a space. Otherwise, the
   /// element should omit printing a space between the surrounding elements.
   bool getValue() const { return value; }
 
public:
   /// Return the index of the first element that needs to be parsed.
   unsigned getParseStart() const { return parseStart; }
 
 private:
   /// The child elements of this optional.
   std::vector<std::unique_ptr<Element>> elements;
   /// The index of the element that acts as the anchor for the optional group.
   unsigned anchor;
-  /// The index of the first element that is parsed (is not a SpaceElement).
+  /// The index of the first element that is parsed (is not a
+  /// WhitespaceElement).
   unsigned parseStart;
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // OperationFormat
 //===----------------------------------------------------------------------===//
 
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     body << "  }\n";
 
     /// Literals.
   } else if (LiteralElement *literal = dyn_cast<LiteralElement>(element)) {
     body << "  if (parser.parse";
     genLiteralParser(literal->getLiteral(), body);
     body << ")\n    return ::mlir::failure();\n";
 
-    /// Spaces.
-  } else if (isa<SpaceElement>(element)) {
+    /// Whitespaces.
+  } else if (isa<WhitespaceElement>(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)) {

 
 void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
                                         Operator &op, bool &shouldEmitSpace,
                                         bool &lastWasPunctuation) {
   if (LiteralElement *literal = dyn_cast<LiteralElement>(element))
     return genLiteralPrinter(literal->getLiteral(), body, shouldEmitSpace,
                              lastWasPunctuation);
 
+  // Emit a whitespace element.
+  if (NewlineElement *newline = dyn_cast<NewlineElement>(element)) {
+    body << "  p.printNewline();\n";
+    return;
+  }
   if (SpaceElement *space = dyn_cast<SpaceElement>(element))
     return genSpacePrinter(space->getValue(), body, shouldEmitSpace,
                            lastWasPunctuation);
 
   // Emit an optional group.
   if (OptionalElement *optional = dyn_cast<OptionalElement>(element)) {
     // Emit the check for the presence of the anchor element.
     Element *anchor = optional->getAnchor();
while (it != e) {
     auto *prevAttr = dyn_cast<AttributeVariable>(element);
     if (!prevAttr || prevAttr->getTypeBuilder())
       continue;
 
     // Check the next iterator within the stack for literal elements.
     for (auto &nextItPair : iteratorStack) {
       ElementsIterT nextIt = nextItPair.first, nextE = nextItPair.second;
       for (; nextIt != nextE; ++nextIt) {
-        // Skip any trailing spaces, attribute dictionaries, or optional groups.
-        if (isa<SpaceElement>(*nextIt) || isa<AttrDictDirective>(*nextIt) ||
-            isa<OptionalElement>(*nextIt))
+        // Skip any trailing whitespace, attribute dictionaries, or optional
+        // groups.
+        if (isa<WhitespaceElement>(*nextIt) ||
+            isa<AttrDictDirective>(*nextIt) || isa<OptionalElement>(*nextIt))
           continue;
 
         // We are only interested in `:` literals.
         auto *literal = dyn_cast<LiteralElement>(&*nextIt);
         if (!literal || literal->getLiteral() != ":")
           break;
 
         // TODO: Use the location of the literal element itself.
LogicalResult FormatParser::parseLiteral(std::unique_ptr<Element> &element) {
 
   StringRef value = literalTok.getSpelling().drop_front().drop_back();
 
   // The parsed literal is a space element (`` or ` `).
   if (value.empty() || (value.size() == 1 && value.front() == ' ')) {
     element = std::make_unique<SpaceElement>(!value.empty());
     return ::mlir::success();
   }
+  // The parsed literal is a newline element.
+  if (value == "\\n") {
+    element = std::make_unique<NewlineElement>();
+    return ::mlir::success();
+  }
 
   // Check that the parsed literal is valid.
   if (!LiteralElement::isValidLiteral(value))
     return emitError(literalTok.getLoc(), "expected valid literal");
 
   element = std::make_unique<LiteralElement>(value);
   return ::mlir::success();
 }
if (failed(parseToken(Token::question, "expected '?' after optional group")))
     return ::mlir::failure();
 
   // The optional group is required to have an anchor.
   if (!anchorIdx)
     return emitError(curLoc, "optional group specified no anchor element");
 
   // The first parsable element of the group must be able to be parsed in an
   // optional fashion.
-  auto parseBegin = llvm::find_if_not(
-      elements, [](auto &element) { return isa<SpaceElement>(element.get()); });
+  auto parseBegin = llvm::find_if_not(elements, [](auto &element) {
+    return isa<WhitespaceElement>(element.get());
+  });
   Element *firstElement = parseBegin->get();
   if (!isa<AttributeVariable>(firstElement) &&
       !isa<LiteralElement>(firstElement) &&
       !isa<OperandVariable>(firstElement) && !isa<RegionVariable>(firstElement))
     return emitError(curLoc,
                      "first parsable element of an operand group must be "
                      "an attribute, literal, operand, or region");
 
return TypeSwitch<Element *, LogicalResult>(childElements.back().get())
         seenVariables.insert(ele->getVar());
         return ::mlir::success();
       })
       .Case<RegionVariable>([&](RegionVariable *) {
         // TODO: When ODS has proper support for marking "optional" regions, add
         // a check here.
         return ::mlir::success();
       })
-      // Literals, spaces, custom directives, and type directives may be used,
-      // but they can't anchor the group.
-      .Case<LiteralElement, SpaceElement, CustomDirective,
+      // Literals, whitespace, custom directives, and type directives may be
+      // used, but they can't anchor the group.
+      .Case<LiteralElement, WhitespaceElement, CustomDirective,
             FunctionalTypeDirective, OptionalElement, TypeRefDirective,
             TypeDirective>([&](Element *) {
         if (isAnchor)
           return emitError(childLoc, "only variables can be used to anchor "
                                      "an optional group");
         return ::mlir::success();
       })
       .Default([&](Element *) {