diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h index f2bae58c80b3a2..7e43418f6f156e 100644 --- a/mlir/include/mlir/IR/OperationSupport.h +++ b/mlir/include/mlir/IR/OperationSupport.h @@ -1096,6 +1096,10 @@ class OpPrintingFlags { /// elements. OpPrintingFlags &elideLargeElementsAttrs(int64_t largeElementLimit = 16); + /// Enables breaking attributes on individual lines when there are more than + /// the given number of attributes on an operation. + OpPrintingFlags& newlineAfterAttribute(int64_t attributeLimit = 2); + /// Enable or disable printing of debug information (based on `enable`). If /// 'prettyForm' is set to true, debug information is printed in a more /// readable 'pretty' form. Note: The IR generated with 'prettyForm' is not @@ -1126,6 +1130,9 @@ class OpPrintingFlags { /// Return the size limit for printing large ElementsAttr. std::optional getLargeElementsAttrLimit() const; + /// Return the size limit for printing newlines after attributes. + std::optional getNewlineAfterAttrLimit() const; + /// Return if debug information should be printed. bool shouldPrintDebugInfo() const; @@ -1152,6 +1159,10 @@ class OpPrintingFlags { /// the upper limit. std::optional elementsAttrElementLimit; + /// Print newlines after each attribute when an operation has more than + /// the given number of attributes. + std::optional newlineAfterAttr; + /// Print debug information. bool printDebugInfoFlag : 1; bool printDebugInfoPrettyFormFlag : 1; diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp index ef0bf75a9cd674..906dda75403d8c 100644 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -140,6 +140,11 @@ struct AsmPrinterOptions { llvm::cl::desc("Elide ElementsAttrs with \"...\" that have " "more elements than the given upper limit")}; + llvm::cl::opt newlineAfterAttr{ + "mlir-newline-after-attr", + llvm::cl::desc("Break attributes on ops into multiple lines with more " + "than the given upper limit")}; + llvm::cl::opt printDebugInfoOpt{ "mlir-print-debuginfo", llvm::cl::init(false), llvm::cl::desc("Print debug info in MLIR output")}; @@ -191,6 +196,8 @@ OpPrintingFlags::OpPrintingFlags() return; if (clOptions->elideElementsAttrIfLarger.getNumOccurrences()) elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger; + if (clOptions->newlineAfterAttr.getNumOccurrences()) + newlineAfterAttr = clOptions->newlineAfterAttr; printDebugInfoFlag = clOptions->printDebugInfoOpt; printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt; printGenericOpFormFlag = clOptions->printGenericOpFormOpt; @@ -209,6 +216,14 @@ OpPrintingFlags::elideLargeElementsAttrs(int64_t largeElementLimit) { return *this; } +/// Enables breaking attributes on individual lines when there are more than +/// the given number of attributes on an operation. +OpPrintingFlags & +OpPrintingFlags::newlineAfterAttribute(int64_t attributeLimit) { + newlineAfterAttr = attributeLimit; + return *this; +} + /// Enable printing of debug information. If 'prettyForm' is set to true, /// debug information is printed in a more readable 'pretty' form. OpPrintingFlags &OpPrintingFlags::enableDebugInfo(bool enable, @@ -262,6 +277,11 @@ std::optional OpPrintingFlags::getLargeElementsAttrLimit() const { return elementsAttrElementLimit; } +/// Return the size limit for printing newlines after attributes. +std::optional OpPrintingFlags::getNewlineAfterAttrLimit() const { + return newlineAfterAttr; +} + /// Return if debug information should be printed. bool OpPrintingFlags::shouldPrintDebugInfo() const { return printDebugInfoFlag; @@ -346,6 +366,12 @@ class AsmPrinter::Impl { llvm::interleaveComma(c, os, eachFn); } + template + inline void interleave(const Container &c, UnaryFunctor eachFn, + StringRef separator) const { + llvm::interleave(c, os, eachFn, separator); + } + /// This enum describes the different kinds of elision for the type of an /// attribute when printing it. enum class AttrTypeElision { @@ -396,7 +422,7 @@ class AsmPrinter::Impl { protected: void printOptionalAttrDict(ArrayRef attrs, ArrayRef elidedAttrs = {}, - bool withKeyword = false); + unsigned currentIndent = 0, bool withKeyword = false); void printNamedAttribute(NamedAttribute attr); void printTrailingLocation(Location loc, bool allowAlias = true); void printLocationInternal(LocationAttr loc, bool pretty = false, @@ -1916,7 +1942,7 @@ void AsmPrinter::Impl::printLocationInternal(LocationAttr loc, bool pretty, os << '>'; } os << '['; - interleave( + llvm::interleave( loc.getLocations(), [&](Location loc) { printLocationInternal(loc, pretty); }, [&]() { os << ", "; }); @@ -2546,6 +2572,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) { void AsmPrinter::Impl::printOptionalAttrDict(ArrayRef attrs, ArrayRef elidedAttrs, + unsigned currentIndent, bool withKeyword) { // If there are no attributes, then there is nothing to be done. if (attrs.empty()) @@ -2556,11 +2583,30 @@ void AsmPrinter::Impl::printOptionalAttrDict(ArrayRef attrs, // Print the 'attributes' keyword if necessary. if (withKeyword) os << " attributes"; + os << " {"; + + SmallString<16> separator = StringRef(", "); + if (printerFlags.getNewlineAfterAttrLimit() && + attrs.size() > *printerFlags.getNewlineAfterAttrLimit()) { + + // Increase indent to match the visually match the "{ " below. + //currentIndent += 2; + + separator.clear(); + separator.reserve(currentIndent + 2); + separator.append(",\n"); + for (size_t i = 0; i < currentIndent; ++i) + separator.push_back(' '); + + // Already put the first attribute on its own line. + os << "\n"; + os.indent(currentIndent); + } // Otherwise, print them all out in braces. - os << " {"; - interleaveComma(filteredAttrs, - [&](NamedAttribute attr) { printNamedAttribute(attr); }); + interleave( + filteredAttrs, [&](NamedAttribute attr) { printNamedAttribute(attr); }, + separator); os << '}'; }; @@ -2980,12 +3026,12 @@ class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter { /// Print an optional attribute dictionary with a given set of elided values. void printOptionalAttrDict(ArrayRef attrs, ArrayRef elidedAttrs = {}) override { - Impl::printOptionalAttrDict(attrs, elidedAttrs); + Impl::printOptionalAttrDict(attrs, elidedAttrs, currentIndent + indentWidth); } void printOptionalAttrDictWithKeyword( ArrayRef attrs, ArrayRef elidedAttrs = {}) override { - Impl::printOptionalAttrDict(attrs, elidedAttrs, + Impl::printOptionalAttrDict(attrs, elidedAttrs, currentIndent + indentWidth, /*withKeyword=*/true); } diff --git a/mlir/test/IR/mlir-newline-after-attr.mlir b/mlir/test/IR/mlir-newline-after-attr.mlir new file mode 100644 index 00000000000000..5657f6e5ad706e --- /dev/null +++ b/mlir/test/IR/mlir-newline-after-attr.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-opt %s -mlir-newline-after-attr=2 | FileCheck %s +// Ensure that the printed version is still parseable. +// RUN: mlir-opt %s -mlir-newline-after-attr=2 | mlir-opt + +// CHECK: foo.dense_attr = +"test.op"() {foo.dense_attr = dense<1> : tensor<3xi32>} : () -> () + +// CHECK: foo.dense_attr = +// CHECK: foo.second_attr = +"test.op"() {foo.dense_attr = dense<1> : tensor<3xi32>, foo.second_attr = dense<2> : tensor<3xi32>} : () -> ()