Enable printing newlines and indents in attribute and type printers

This commit moves the code responsible for adding newlines and tracking indent, so that it can be used not only for operation printers, but also for attribute and type printers.

It could be useful for nested attributes, where proper formatting with newlines and indents would benefit the readability of the IR. Currently, everything is printed on one line, which makes it difficult to read if the attribute is more verbose and there are multiple levels of nesting.

Author: Andruszkiewicz, Jacenty

mlir/include/mlir/IR/OpImplementation.h

@@ -130,6 +130,16 @@ class AsmPrinter {
   /// Return the raw output stream used by this printer.
   virtual raw_ostream &getStream() const;
 
+  /// Print a newline and indent the printer to the start of the current
+  /// operation.
+  virtual void printNewline();
+
+  /// Increase indentation.
+  virtual void increaseIndent();
+
+  /// Decrease indentation.
+  virtual void decreaseIndent();
+
   /// Print the given floating point value in a stabilized form that can be
   /// roundtripped through the IR. This is the companion to the 'parseFloat'
   /// hook on the AsmParser.
@@ -448,16 +458,6 @@ class OpAsmPrinter : public AsmPrinter {
   /// Print a loc(...) specifier if printing debug info is enabled.
   virtual void printOptionalLocationSpecifier(Location loc) = 0;
 
-  /// Print a newline and indent the printer to the start of the current
-  /// operation.
-  virtual void printNewline() = 0;
-
-  /// Increase indentation.
-  virtual void increaseIndent() = 0;
-
-  /// Decrease indentation.
-  virtual void decreaseIndent() = 0;
-
   /// Print a block argument in the usual format of:
   ///   %ssaName : type {attr1=42} loc("here")
   /// where location printing is controlled by the standard internal option.

mlir/lib/IR/AsmPrinter.cpp

@@ -418,6 +418,19 @@ class AsmPrinter::Impl {
   /// Returns the output stream of the printer.
   raw_ostream &getStream() { return os; }
 
+  /// Print a newline and indent the printer to the start of the current
+  /// operation.
+  void printNewline() {
+    os << newLine;
+    os.indent(currentIndent);
+  }
+
+  /// Increase indentation.
+  void increaseIndent() { currentIndent += indentWidth; }
+
+  /// Decrease indentation.
+  void decreaseIndent() { currentIndent -= indentWidth; }
+
   template <typename Container, typename UnaryFunctor>
   inline void interleaveComma(const Container &c, UnaryFunctor eachFn) const {
     llvm::interleaveComma(c, os, eachFn);
@@ -532,6 +545,12 @@ class AsmPrinter::Impl {
 
   /// A tracker for the number of new lines emitted during printing.
   NewLineCounter newLine;
+
+  /// The number of spaces used for indenting nested operations.
+  const static unsigned indentWidth = 2;
+
+  /// This is the current indentation level for nested structures.
+  unsigned currentIndent = 0;
 };
 } // namespace mlir
 
@@ -1004,6 +1023,9 @@ class DummyAliasDialectAsmPrinter : public DialectAsmPrinter {
 
   /// The following are hooks of `DialectAsmPrinter` that are not necessary for
   /// determining potential aliases.
+  void printNewline() override {}
+  void increaseIndent() override {}
+  void decreaseIndent() override {}
   void printFloat(const APFloat &) override {}
   void printKeywordOrString(StringRef) override {}
   void printString(StringRef) override {}
@@ -2892,6 +2914,13 @@ void AsmPrinter::Impl::printDialectAttribute(Attribute attr) {
   {
     llvm::raw_string_ostream attrNameStr(attrName);
     Impl subPrinter(attrNameStr, state);
+
+    // The values of currentIndent and newLine are assigned to the created
+    // subprinter, so that the indent level and number of printed lines can be
+    // tracked.
+    subPrinter.currentIndent = currentIndent;
+    subPrinter.newLine = newLine;
+
     DialectAsmPrinter printer(subPrinter);
     dialect.printAttribute(attr, printer);
   }
@@ -2906,6 +2935,13 @@ void AsmPrinter::Impl::printDialectType(Type type) {
   {
     llvm::raw_string_ostream typeNameStr(typeName);
     Impl subPrinter(typeNameStr, state);
+
+    // The values of currentIndent and newLine are assigned to the created
+    // subprinter, so that the indent level and number of printed lines can be
+    // tracked.
+    subPrinter.currentIndent = currentIndent;
+    subPrinter.newLine = newLine;
+
     DialectAsmPrinter printer(subPrinter);
     dialect.printType(type, printer);
   }
@@ -2946,6 +2982,21 @@ raw_ostream &AsmPrinter::getStream() const {
   return impl->getStream();
 }
 
+void AsmPrinter::printNewline() {
+  assert(impl && "expected AsmPrinter::printNewLine to be overriden");
+  impl->printNewline();
+}
+
+void AsmPrinter::increaseIndent() {
+  assert(impl && "expected AsmPrinter::increaseIndent to be overriden");
+  impl->increaseIndent();
+}
+
+void AsmPrinter::decreaseIndent() {
+  assert(impl && "expected AsmPrinter::decreaseIndent to be overriden");
+  impl->decreaseIndent();
+}
+
 /// Print the given floating point value in a stablized form.
 void AsmPrinter::printFloat(const APFloat &value) {
   assert(impl && "expected AsmPrinter::printFloat to be overriden");
@@ -3276,19 +3327,6 @@ class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
     printTrailingLocation(loc);
   }
 
-  /// Print a newline and indent the printer to the start of the current
-  /// operation.
-  void printNewline() override {
-    os << newLine;
-    os.indent(currentIndent);
-  }
-
-  /// Increase indentation.
-  void increaseIndent() override { currentIndent += indentWidth; }
-
-  /// Decrease indentation.
-  void decreaseIndent() override { currentIndent -= indentWidth; }
-
   /// Print a block argument in the usual format of:
   ///   %ssaName : type {attr1=42} loc("here")
   /// where location printing is controlled by the standard internal option.
@@ -3414,12 +3452,6 @@ class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
   // top-level we start with "builtin" as the default, so that the top-level
   // `module` operation prints as-is.
   SmallVector<StringRef> defaultDialectStack{"builtin"};
-
-  /// The number of spaces used for indenting nested operations.
-  const static unsigned indentWidth = 2;
-
-  // This is the current indentation level for nested structures.
-  unsigned currentIndent = 0;
 };
 } // namespace
 

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

@@ -439,4 +439,10 @@ def TestCustomStorageCtorAttr : Test_Attr<"TestCustomStorageCtorAttr"> {
     let hasStorageCustomConstructor = 1;
 }
 
+def TestAttrNewlineAndIndent : Test_Attr<"TestAttrNewlineAndIndent"> {
+  let mnemonic = "newline_and_indent";
+  let parameters = (ins "::mlir::Type":$indentType);
+  let hasCustomAssemblyFormat = 1;
+}
+
 #endif // TEST_ATTRDEFS

mlir/test/lib/Dialect/Test/TestAttributes.cpp

@@ -430,6 +430,35 @@ bool TestConstMemorySpaceAttr::isValidPtrIntCast(
   return false;
 }
 
+//===----------------------------------------------------------------------===//
+// TestAttrNewlineAndIndent
+//===----------------------------------------------------------------------===//
+
+Attribute TestAttrNewlineAndIndentAttr::parse(::mlir::AsmParser &parser,
+                                              ::mlir::Type type) {
+  Type indentType;
+  if (parser.parseLess()) {
+    return {};
+  }
+  if (parser.parseType(indentType)) {
+    return {};
+  }
+  if (parser.parseGreater()) {
+    return {};
+  }
+  return get(parser.getContext(), indentType);
+}
+
+void TestAttrNewlineAndIndentAttr::print(::mlir::AsmPrinter &printer) const {
+  printer << "<";
+  printer.increaseIndent();
+  printer.printNewline();
+  printer << getIndentType();
+  printer.decreaseIndent();
+  printer.printNewline();
+  printer << ">";
+}
+
 //===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//

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

@@ -470,4 +470,9 @@ def TestMemrefType : Test_Type<"TestMemref",
   }];
 }
 
+def TestTypeNewlineAndIndent : Test_Type<"TestTypeNewlineAndIndent"> {
+  let mnemonic = "newline_and_indent";
+  let hasCustomAssemblyFormat = 1;
+}
+
 #endif // TEST_TYPEDEFS

mlir/test/lib/Dialect/Test/TestTypes.cpp

@@ -573,3 +573,30 @@ ::mlir::LogicalResult TestTensorType::verifyCompatibleBufferType(
                      getElementType() == testMemref.getElementType();
   return mlir::success(valid);
 }
+
+//===----------------------------------------------------------------------===//
+// TestTypeNewlineAndIndent
+//===----------------------------------------------------------------------===//
+
+Type TestTypeNewlineAndIndentType::parse(::mlir::AsmParser &parser) {
+  if (parser.parseLess()) {
+    return {};
+  }
+  if (parser.parseKeyword("indented_content")) {
+    return {};
+  }
+  if (parser.parseGreater()) {
+    return {};
+  }
+  return get(parser.getContext());
+}
+
+void TestTypeNewlineAndIndentType::print(::mlir::AsmPrinter &printer) const {
+  printer << "<";
+  printer.increaseIndent();
+  printer.printNewline();
+  printer << "indented_content";
+  printer.decreaseIndent();
+  printer.printNewline();
+  printer << ">";
+}

mlir/test/mlir-tblgen/testdialect-attrdefs.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s --strict-whitespace
 
 // CHECK-LABEL: func private @compoundA()
 // CHECK-SAME: #test.cmpnd_a<1, !test.smpla, [5, 6]>
@@ -44,3 +44,19 @@ func.func private @hexdecimalInteger() attributes {
 // expected-error @below {{expected an integer}}
   sdg = #test.decimal_shape<1x0xb>
 }
+
+// -----
+
+// CHECK-LABEL: @newlineAndIndent
+// CHECK-SAME:  indent = #test.newline_and_indent<
+// CHECK-NEXT:  {{^    }}!test.newline_and_indent<    
+// CHECK-NEXT:  {{^      }}indented_content
+// CHECK-NEXT:  {{^    }}>
+// CHECK-NEXT:  {{^  }}>
+func.func private @newlineAndIndent() attributes {
+  indent = #test.newline_and_indent<
+    !test.newline_and_indent<
+      indented_content
+    >
+  >
+}

mlir/test/mlir-tblgen/testdialect-typedefs.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s | mlir-opt -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s | mlir-opt -verify-diagnostics | FileCheck %s --strict-whitespace
 
 //////////////
 // Tests the types in the 'Test' dialect, not the ones in 'typedefs.mlir'
@@ -42,3 +42,13 @@ func.func @testInt(%A : !test.int<s, 8>, %B : !test.int<unsigned, 2>, %C : !test
 func.func @structTest (%A : !test.struct< {field1, !test.smpla}, {field2, !test.int<none, 3>} > ) {
   return
 }
+
+// CHECK-LABEL: @newlineAndIndent
+// CHECK-SAME:  !test.newline_and_indent<
+// CHECK-NEXT:  {{^    }}indented_content
+// CHECK-NEXT:  {{^  }}>
+func.func @newlineAndIndent(%A : !test.newline_and_indent<
+  indented_content
+>) {
+  return
+}

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

@@ -958,9 +958,7 @@ void DefFormat::genOptionalGroupPrinter(OptionalElement *el, FmtContext &ctx,
 void DefFormat::genWhitespacePrinter(WhitespaceElement *el, FmtContext &ctx,
                                      MethodBody &os) {
   if (el->getValue() == "\\n") {
-    // FIXME: The newline should be `printer.printNewLine()`, i.e., handled by
-    // the printer.
-    os << tgfmt("$_printer << '\\n';\n", &ctx);
+    os << tgfmt("$_printer.printNewline();\n", &ctx);
   } else if (!el->getValue().empty()) {
     os << tgfmt("$_printer << \"$0\";\n", &ctx, el->getValue());
   } else {