Add indented raw_ostream class

[jpienaar]

Summary

Class simplifies keeping track of the indentation while emitting. For every new line the current indentation is simply prefixed (if not at start of line, then it just emits as normal). Add a simple Region helper that makes it easy to have the C++ scope match the emitted scope.

Use this in op doc generator and rewrite generator.

[rriddle]

Any reason why you can't just use/improve llvm::ScopedPrinter instead? https://github.com/llvm/llvm-project/blob/2e0ee477127c44b86c018da7f24a622ec8f8fccd/llvm/include/llvm/Support/ScopedPrinter.h

Seems like this serves the exact same use case. Also, this seems like something that would go into llvm/Support and not mlir/Support.

[jpienaar]

In D84107#2160458, @rriddle wrote:

Any reason why you can't just use/improve llvm::ScopedPrinter instead? https://github.com/llvm/llvm-project/blob/2e0ee477127c44b86c018da7f24a622ec8f8fccd/llvm/include/llvm/Support/ScopedPrinter.h

Don't they serve different purposes? E.g., ScopedPrinter is meant to be used locally, it adds labels everywhere, can't be used where raw_ostream can and uses different calling conventions rather than << methods.

[antiagainst]

This looks awesome! Sorry about the delay... LGTM; just a few nits. I see River has comments so would be good to get his consent too.

mlir/include/mlir/Support/IndentedOstream.h

40

Re-indents?

44

Increases?

50

Decreases?

65

Probably can delete this comment; not very useful.

70

currentPos?

78

s/uses/used/

mlir/tools/mlir-tblgen/RewriterGen.cpp

510

Are these tabs that will need clang-format?

[jpienaar]

Addressed typos and used more similar naming to ScopedPrinter

[mehdi_amini]

Isn't it something that belongs to the LLVM Support library?

[mehdi_amini]

In D84107#2160458, @rriddle wrote:

Also, this seems like something that would go into llvm/Support and not mlir/Support.

Actually River commented on this but I don't think you addressed it?

[jpienaar]

In D84107#2313417, @mehdi_amini wrote:

In D84107#2160458, @rriddle wrote:

Also, this seems like something that would go into llvm/Support and not mlir/Support.

That depends, almost everything in mlir/Support could also be in llvm/Support, as these are mostly general utils (e.g., nothing in StorageUniquer.h is MLIR unique it would seem). Wouldn't the question be is it used beyond MLIR too? E.g., just moving parts which have no current use in larger project to more core location doesn't make as much to me, moving it when it becomes used there seems more of a fit.

Actually River commented on this but I don't think you addressed it?

River asked if it should be an extension of ScopedPrinter which I did address.

mlir/include/mlir/Support/IndentedOstream.h

70

Overriding base class member, but yes inconsistent.

mlir/tools/mlir-tblgen/RewriterGen.cpp

510

Artifact of new phab interface (no tabs here)

[mehdi_amini]

In D84107#2314551, @jpienaar wrote:

In D84107#2313417, @mehdi_amini wrote:

In D84107#2160458, @rriddle wrote:

Also, this seems like something that would go into llvm/Support and not mlir/Support.

That depends, almost everything in mlir/Support could also be in llvm/Support,

We have been moving everything we could, and if things are remaining from when MLIR was not in-tree, they are just "pending".

as these are mostly general utils (e.g., nothing in StorageUniquer.h is MLIR unique it would seem). Wouldn't the question be is it used beyond MLIR too? E.g., just moving parts which have no current use in larger project to more core location doesn't make as much to me, moving it when it becomes used there seems more of a fit.

I think this is backward: things won't get used unless they are in LLVM Support, and also this not in line with the existing practice either: we have been moving things from MLIR Support to LLVM support as much as possible.

Actually River commented on this but I don't think you addressed it?

River asked if it should be an extension of ScopedPrinter which I did address.

I quoted River above. His sentence was independent from the comment on ScopedPrinter: "Also, this seems like something that would go into llvm/Support and not mlir/Support."
I believe he was asking the same thing as I do.

mlir/include/mlir/Support/IndentedOstream.h

+//===- IndentedOstream.h - raw ostream wrapper to indent --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// raw_ostream subclass that keeps track of indentation for textual output
+// where indentation helps readability.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_INDENTEDOSTREAM_H_
+#define MLIR_SUPPORT_INDENTEDOSTREAM_H_
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+
+/// raw_ostream subclass that simplifies indention a sequence of code.
+class raw_indented_ostream : public raw_ostream {
+public:
+  explicit raw_indented_ostream(llvm::raw_ostream &os) : os(os) {
+    SetUnbuffered();
+  }
+
+  /// Simple RAII struct to use to indentation around entering/exiting region.
+  struct DelimitedScope {
+    explicit DelimitedScope(raw_indented_ostream &os, StringRef open = "",
+                            StringRef close = "")
+        : os(os), open(open), close(close) {
+      os << open;
+      os.indent();
+    }
+    ~DelimitedScope() {
+      os.unindent();
+      os << close;
+    }
+
+    raw_indented_ostream &os;
+
+  private:
+    llvm::StringRef open, close;
+  };
+
+  /// Returns DelimitedScope.
+  DelimitedScope scope(StringRef open = "", StringRef close = "") {
+    return DelimitedScope(*this, open, close);
+  }
+
+  /// Re-indents by removing the leading whitespace from the first non-empty
+  /// line from every line of the the string, skipping over empty lines at the
+  /// start.
+  raw_indented_ostream &reindent(StringRef str);
+
+  /// Increases the indent and returning this raw_indented_ostream.
+  raw_indented_ostream &indent() {
+    currentIndent += indentSize;
+    return *this;
+  }
+
+  /// Decreases the indent and returning this raw_indented_ostream.
+  raw_indented_ostream &unindent() {
+    currentIndent = std::max(0, currentIndent - indentSize);
+    return *this;
+  }
+
+  /// Emits whitespace and sets the indendation for the stream.
+  raw_indented_ostream &indent(int with) {
+    os.indent(with);
+    atStartOfLine = false;
+    currentIndent = with;
+    return *this;
+  }
+
+private:
+  void write_impl(const char *ptr, size_t size) override;
+
+  /// Return the current position within the stream, not counting the bytes
+  /// currently in the buffer.
+  uint64_t current_pos() const override { return os.tell(); }
+
+  /// Constant indent added/removed.
+  static constexpr int indentSize = 2;
+
+  // Tracker for current indentation.
+  int currentIndent = 0;
+
+  // The leading whitespace of the string being printed, if reindent is used.
+  int leadingWs = 0;
+
+  // Tracks whether at start of line and so indent is required or not.
+  bool atStartOfLine = true;
+
+  // The underlying raw_ostream.
+  raw_ostream &os;
+};
+
+} // namespace mlir
+#endif // MLIR_SUPPORT_INDENTEDOSTREAM_H_

mlir/lib/Support/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   FileUtilities.cpp
+  IndentedOstream.cpp
   MlirOptMain.cpp
   StorageUniquer.cpp
   ToolUtilities.cpp
 )
 
 add_mlir_library(MLIRSupport
   FileUtilities.cpp
   StorageUniquer.cpp
add_mlir_library(MLIROptLib
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRParser
   MLIRSupport
   )
+
+# This doesn't use add_mlir_library as it is used in mlir-tblgen and else
+# mlir-tblgen ends up depending on mlir-generic-headers.
+add_llvm_library(MLIRSupportIdentedOstream
+  IndentedOstream.cpp
+
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support)

mlir/lib/Support/IndentedOstream.cpp

+//===- IndentedOstream.cpp - raw ostream wrapper to indent ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// raw_ostream subclass that keeps track of indentation for textual output
+// where indentation helps readability.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/IndentedOstream.h"
+
+using namespace mlir;
+
+raw_indented_ostream &mlir::raw_indented_ostream::reindent(StringRef str) {
+  StringRef remaining = str;
+  // Find leading whitespace indent.
+  while (!remaining.empty()) {
+    auto split = remaining.split('\n');
+    size_t indent = split.first.find_first_not_of(" \t");
+    if (indent != StringRef::npos) {
+      leadingWs = indent;
+      break;
+    }
+    remaining = split.second;
+  }
+  // Print, skipping the empty lines.
+  *this << remaining;
+  leadingWs = 0;
+  return *this;
+}
+
+void mlir::raw_indented_ostream::write_impl(const char *ptr, size_t size) {
+  StringRef str(ptr, size);
+  // Print out indented.
+  auto print = [this](StringRef str) {
+    if (atStartOfLine)
+      os.indent(currentIndent) << str.substr(leadingWs);
+    else
+      os << str.substr(leadingWs);
+  };
+
+  while (!str.empty()) {
+    size_t idx = str.find('\n');
+    if (idx == StringRef::npos) {
+      if (!str.substr(leadingWs).empty()) {
+        print(str);
+        atStartOfLine = false;
+      }
+      break;
+    }
+
+    auto split =
+        std::make_pair(str.slice(0, idx), str.slice(idx + 1, StringRef::npos));
+    // Print empty new line without spaces if line only has spaces.
+    if (!split.first.ltrim().empty())
+      print(split.first);
+    os << '\n';
+    atStartOfLine = true;
+    str = split.second;
+  }
+}

mlir/tools/mlir-tblgen/CMakeLists.txt

add_tablegen(mlir-tblgen MLIR
   RewriterGen.cpp
   SPIRVUtilsGen.cpp
   StructsGen.cpp
   )
 
 set_target_properties(mlir-tblgen PROPERTIES FOLDER "Tablegenning")
 target_link_libraries(mlir-tblgen
   PRIVATE
+  MLIRSupportIdentedOstream
   MLIRTableGen)
 
 mlir_check_all_link_libraries(mlir-tblgen)

mlir/tools/mlir-tblgen/OpDocGen.cpp

 //===- OpDocGen.cpp - MLIR operation documentation generator --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // OpDocGen uses the description of operations to generate documentation for the
 // operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "DocGenUtilities.h"
+#include "mlir/Support/IndentedOstream.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Operator.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 #include "llvm/TableGen/TableGenBackend.h"
 
 using namespace llvm;
 using namespace mlir;
 using namespace mlir::tblgen;
 
 using mlir::tblgen::Operator;
 
 // Emit the description by aligning the text to the left per line (e.g.,
 // removing the minimum indentation across the block).
 //
 // This expects that the description in the tablegen file is already formatted
 // in a way the user wanted but has some additional indenting due to being
 // nested in the op definition.
 void mlir::tblgen::emitDescription(StringRef description, raw_ostream &os) {
-  // Determine the minimum number of spaces in a line.
-  size_t min_indent = -1;
-  StringRef remaining = description;
-  while (!remaining.empty()) {
-    auto split = remaining.split('\n');
-    size_t indent = split.first.find_first_not_of(" \t");
-    if (indent != StringRef::npos)
-      min_indent = std::min(indent, min_indent);
-    remaining = split.second;
-  }
-
-  // Print out the description indented.
-  os << "\n";
-  remaining = description;
-  bool printed = false;
-  while (!remaining.empty()) {
-    auto split = remaining.split('\n');
-    if (split.second.empty()) {
-      // Skip last line with just spaces.
-      if (split.first.ltrim().empty())
-        break;
-    }
-    // Print empty new line without spaces if line only has spaces, unless no
-    // text has been emitted before.
-    if (split.first.ltrim().empty()) {
-      if (printed)
-        os << "\n";
-    } else {
-      os << split.first.substr(min_indent) << "\n";
-      printed = true;
-    }
-    remaining = split.second;
-  }
+  raw_indented_ostream ros(os);
+  ros.reindent(description.rtrim(" \t"));
 }
 
 // Emits `str` with trailing newline if not empty.
 static void emitIfNotEmpty(StringRef str, raw_ostream &os) {
   if (!str.empty()) {
     emitDescription(str, os);
     os << "\n";
   }

 }
 
 static void emitOpDoc(Operator op, raw_ostream &os) {
   os << llvm::formatv("### `{0}` ({1})\n", op.getOperationName(),
                       op.getQualCppClassName());
 
   // Emit the summary, syntax, and description if present.
   if (op.hasSummary())
-    os << "\n" << op.getSummary() << "\n";
+    os << "\n" << op.getSummary() << "\n\n";
   if (op.hasAssemblyFormat())
     emitAssemblyFormat(op.getOperationName(), op.getAssemblyFormat().trim(),
                        os);
   if (op.hasDescription())
     mlir::tblgen::emitDescription(op.getDescription(), os);
 
   // Emit attributes.
   if (op.getNumAttributes() != 0) {
static void emitDialectDoc(const RecordKeeper &recordKeeper, raw_ostream &os) {
   }
   for (auto *typeDef : typeDefs) {
     Type type(typeDef);
     if (auto dialect = type.getDialect())
       dialectTypes[dialect].push_back(type);
   }
 
   os << "<!-- Autogenerated by mlir-tblgen; don't manually edit -->\n";
-  for (auto dialectWithOps : dialectOps)
+  for (const auto &dialectWithOps : dialectOps)
     emitDialectDoc(dialectWithOps.first, dialectWithOps.second,
                    dialectTypes[dialectWithOps.first], os);
 }
 
 //===----------------------------------------------------------------------===//
 // Gen Registration
 //===----------------------------------------------------------------------===//
 

mlir/tools/mlir-tblgen/RewriterGen.cpp

 //===- RewriterGen.cpp - MLIR pattern rewriter generator ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // RewriterGen uses pattern rewrite definitions to generate rewriter matchers.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Support/IndentedOstream.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/Pattern.h"
 #include "mlir/TableGen/Predicate.h"
 #include "mlir/TableGen/Type.h"
 #include "llvm/ADT/StringExtras.h"
private:
   //===--------------------------------------------------------------------===//
 
   // Emits C++ statements for matching the op constrained by the given DAG
   // `tree`.
   void emitOpMatch(DagNode tree, int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an operand.
-  void emitOperandMatch(DagNode tree, int argIndex, int depth, int indent);
+  void emitOperandMatch(DagNode tree, int argIndex, int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an attribute.
-  void emitAttributeMatch(DagNode tree, int argIndex, int depth, int indent);
+  void emitAttributeMatch(DagNode tree, int argIndex, int depth);
 
   // Emits C++ for checking a match with a corresponding match failure
   // diagnostic.
   void emitMatchCheck(int depth, const FmtObjectBase &matchFmt,
                       const llvm::formatv_object_base &failureFmt);
 
   //===--------------------------------------------------------------------===//
   // Rewrite utilities
private:
   Pattern pattern;
 
   // Map for all bound symbols' info.
   SymbolInfoMap symbolInfoMap;
 
   // The next unused ID for newly created values.
   unsigned nextValueId;
 
-  raw_ostream &os;
+  raw_indented_ostream os;
 
   // Format contexts containing placeholder substitutions.
   FmtContext fmtCtx;
 
   // Number of op processed.
   int opCounter = 0;
 };
 } // end anonymous namespace
void PatternEmitter::emitOpMatch(DagNode tree, int depth) {
 
   int indent = 4 + 2 * depth;
   os.indent(indent) << formatv(
       "auto castedOp{0} = dyn_cast_or_null<{1}>(op{0}); (void)castedOp{0};\n",
       depth, op.getQualCppClassName());
   // Skip the operand matching at depth 0 as the pattern rewriter already does.
   if (depth != 0) {
     // Skip if there is no defining operation (e.g., arguments to function).
-    os.indent(indent) << formatv("if (!castedOp{0}) return failure();\n",
-                                 depth);
+    os << formatv("if (!castedOp{0})\n  return failure();\n", depth);
   }
   if (tree.getNumArgs() != op.getNumArgs()) {
     PrintFatalError(loc, formatv("op '{0}' argument number mismatch: {1} in "
                                  "pattern vs. {2} in definition",
                                  op.getOperationName(), tree.getNumArgs(),
                                  op.getNumArgs()));
   }
 
   // If the operand's name is set, set to that variable.
   auto name = tree.getSymbol();
   if (!name.empty())
-    os.indent(indent) << formatv("{0} = castedOp{1};\n", name, depth);
+    os << formatv("{0} = castedOp{1};\n", name, depth);
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     auto opArg = op.getArg(i);
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
         if (operand->isVariableLength()) {
           auto error = formatv("use nested DAG construct to match op {0}'s "
                                "variadic operand #{1} unsupported now",
                                op.getOperationName(), i);
           PrintFatalError(loc, error);
         }
       }
-      os.indent(indent) << "{\n";
+      os << "{\n";
 
-      os.indent(indent + 2) << formatv(
+      os.indent() << formatv(
           "auto *op{0} = "
           "(*castedOp{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
           depth + 1, depth, i);
       emitOpMatch(argTree, depth + 1);
-      os.indent(indent + 2)
-          << formatv("tblgen_ops[{0}] = op{1};\n", ++opCounter, depth + 1);
-      os.indent(indent) << "}\n";
+      os << formatv("tblgen_ops[{0}] = op{1};\n", ++opCounter, depth + 1);
+      os.unindent() << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
-      emitOperandMatch(tree, i, depth, indent);
+      emitOperandMatch(tree, i, depth);
     } else if (opArg.is<NamedAttribute *>()) {
-      emitAttributeMatch(tree, i, depth, indent);
+      emitAttributeMatch(tree, i, depth);
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 }
 
-void PatternEmitter::emitOperandMatch(DagNode tree, int argIndex, int depth,
-                                      int indent) {
+void PatternEmitter::emitOperandMatch(DagNode tree, int argIndex, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *operand = op.getArg(argIndex).get<NamedTypeConstraint *>();
   auto matcher = tree.getArgAsLeaf(argIndex);
 
   // If a constraint is specified, we need to generate C++ statements to
   // check the constraint.
   if (!matcher.isUnspecified()) {
     if (!matcher.isOperandMatcher()) {
void PatternEmitter::emitOperandMatch(DagNode tree, int argIndex, int depth) {
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     // We need to subtract the number of attributes before this operand to get
     // the index in the operand list.
     auto numPrevAttrs = std::count_if(
         op.arg_begin(), op.arg_begin() + argIndex,
         [](const Argument &arg) { return arg.is<NamedAttribute *>(); });
 
-    os.indent(indent) << formatv("{0} = castedOp{1}.getODSOperands({2});\n",
-                                 name, depth, argIndex - numPrevAttrs);
+    os << formatv("{0} = castedOp{1}.getODSOperands({2});\n", name, depth,
+                  argIndex - numPrevAttrs);
   }
 }
 
-void PatternEmitter::emitAttributeMatch(DagNode tree, int argIndex, int depth,
-                                        int indent) {
+void PatternEmitter::emitAttributeMatch(DagNode tree, int argIndex, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *namedAttr = op.getArg(argIndex).get<NamedAttribute *>();
   const auto &attr = namedAttr->attr;
 
-  os.indent(indent) << "{\n";
-  indent += 2;
-  os.indent(indent) << formatv(
-      "auto tblgen_attr = op{0}->getAttrOfType<{1}>(\"{2}\");"
+  os << "{\n";
+  os.indent() << formatv(
+      "auto tblgen_attr = op{0}->getAttrOfType<{1}>(\"{2}\"); "
       "(void)tblgen_attr;\n",
       depth, attr.getStorageType(), namedAttr->name);
 
   // TODO: This should use getter method to avoid duplication.
   if (attr.hasDefaultValue()) {
-    os.indent(indent) << "if (!tblgen_attr) tblgen_attr = "
-                      << std::string(tgfmt(attr.getConstBuilderTemplate(),
-                                           &fmtCtx, attr.getDefaultValue()))
-                      << ";\n";
+    os << "if (!tblgen_attr) tblgen_attr = "
+       << std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx,
+                            attr.getDefaultValue()))
   } else if (attr.isOptional()) {
     // For a missing attribute that is optional according to definition, we
     // should just capture a mlir::Attribute() to signal the missing state.
     // That is precisely what getAttr() returns on missing attributes.
   } else {
     emitMatchCheck(depth, tgfmt("tblgen_attr", &fmtCtx),
                    formatv("\"expected op '{0}' to have attribute '{1}' "
                            "of type '{2}'\"",
emitMatchCheck(
                 op.getOperationName(), namedAttr->name,
                 matcher.getAsConstraint().getDescription()));
   }
 
   // Capture the value
   auto name = tree.getArgName(argIndex);
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
-    os.indent(indent) << formatv("{0} = tblgen_attr;\n", name);
+    os << formatv("{0} = tblgen_attr;\n", name);
   }
 
-  indent -= 2;
-  os.indent(indent) << "}\n";
+  os.unindent() << "}\n";
 }
 
 void PatternEmitter::emitMatchCheck(
     int depth, const FmtObjectBase &matchFmt,
     const llvm::formatv_object_base &failureFmt) {
-  // {0} The match depth (used to get the operation that failed to match).
-  // {1} The format for the match string.
-  // {2} The format for the failure string.
-  const char *matchStr = R"(
-    if (!({1})) {
-      return rewriter.notifyMatchFailure(op{0}, [&](::mlir::Diagnostic &diag) {
-        diag << {2};
-      });
-    })";
-  os << llvm::formatv(matchStr, depth, matchFmt.str(), failureFmt.str())
-     << "\n";
+  os << "if (!(" << matchFmt.str() << "))";
+  os.scope("{\n", "\n}\n").os
+      << "return rewriter.notifyMatchFailure(op" << depth
+      << ", [&](::mlir::Diagnostic &diag) {\n  diag << " << failureFmt.str()
+      << ";\n});";
 }
 
 void PatternEmitter::emitMatchLogic(DagNode tree) {
   LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n");
   int depth = 0;
   emitOpMatch(tree, depth);
 
   for (auto &appliedConstraint : pattern.getConstraints()) {
void PatternEmitter::emit(StringRef rewriteName) {
   LLVM_DEBUG(llvm::dbgs() << "start collecting ops used in result patterns\n");
   for (unsigned i = 0, e = pattern.getNumResultPatterns(); i != e; ++i) {
     collectOps(pattern.getResultPattern(i), resultOps);
   }
   LLVM_DEBUG(llvm::dbgs() << "done collecting ops used in result patterns\n");
 
   // Emit RewritePattern for Pattern.
   auto locs = pattern.getLocation();
-  os << formatv("/* Generated from:\n\t{0:$[ instantiating\n\t]}\n*/\n",
+  os << formatv("/* Generated from:\n    {0:$[ instantiating\n    ]}\n*/\n",
                 make_range(locs.rbegin(), locs.rend()));
   os << formatv(R"(struct {0} : public ::mlir::RewritePattern {
   {0}(::mlir::MLIRContext *context)
       : ::mlir::RewritePattern("{1}", {{)",
                 rewriteName, rootName);
   // Sort result operators by name.
   llvm::SmallVector<const Operator *, 4> sortedResultOps(resultOps.begin(),
                                                          resultOps.end());
   llvm::sort(sortedResultOps, [&](const Operator *lhs, const Operator *rhs) {
     return lhs->getOperationName() < rhs->getOperationName();
   });
   llvm::interleaveComma(sortedResultOps, os, [&](const Operator *op) {
     os << '"' << op->getOperationName() << '"';
   });
   os << formatv(R"(}, {0}, context) {{})", pattern.getBenefit()) << "\n";
 
   // Emit matchAndRewrite() function.
-  os << R"(
-  ::mlir::LogicalResult
-  matchAndRewrite(::mlir::Operation *op0,
-                  ::mlir::PatternRewriter &rewriter) const override {
-)";
+  {
+    auto classScope = os.scope();
+    os.reindent(R"(
+    ::mlir::LogicalResult matchAndRewrite(::mlir::Operation *op0,
+        ::mlir::PatternRewriter &rewriter) const override {)")
+        << '\n';
+    {
+      auto functionScope = os.scope();
 
       // Register all symbols bound in the source pattern.
       pattern.collectSourcePatternBoundSymbols(symbolInfoMap);
 
-  LLVM_DEBUG(
-      llvm::dbgs() << "start creating local variables for capturing matches\n");
-  os.indent(4) << "// Variables for capturing values and attributes used for "
-                  "creating ops\n";
+      LLVM_DEBUG(llvm::dbgs()
+                 << "start creating local variables for capturing matches\n");
+      os << "// Variables for capturing values and attributes used while "
       // Create local variables for storing the arguments and results bound
       // to symbols.
       for (const auto &symbolInfoPair : symbolInfoMap) {
         StringRef symbol = symbolInfoPair.getKey();
         auto &info = symbolInfoPair.getValue();
-    os.indent(4) << info.getVarDecl(symbol);
+        os << info.getVarDecl(symbol);
       }
       // TODO: capture ops with consistent numbering so that it can be
       // reused for fused loc.
-  os.indent(4) << formatv("::mlir::Operation *tblgen_ops[{0}];\n\n",
-                          pattern.getSourcePattern().getNumOps());
-  LLVM_DEBUG(
-      llvm::dbgs() << "done creating local variables for capturing matches\n");
+      os << formatv("::mlir::Operation *tblgen_ops[{0}];\n\n",
+      LLVM_DEBUG(llvm::dbgs()
+                 << "done creating local variables for capturing matches\n");
 
-  os.indent(4) << "// Match\n";
-  os.indent(4) << "tblgen_ops[0] = op0;\n";
+      os << "// Match\n";
+      os << "tblgen_ops[0] = op0;\n";
       emitMatchLogic(sourceTree);
-  os << "\n";
 
-  os.indent(4) << "// Rewrite\n";
+      os << "\n// Rewrite\n";
       emitRewriteLogic();
 
-  os.indent(4) << "return success();\n";
-  os << "  };\n";
+      os << "return success();\n";
+    }
     os << "};\n";
   }
+  os << "};\n\n";
+}
 
 void PatternEmitter::emitRewriteLogic() {
   LLVM_DEBUG(llvm::dbgs() << "--- start emitting rewrite logic ---\n");
   const Operator &rootOp = pattern.getSourceRootOp();
   int numExpectedResults = rootOp.getNumResults();
   int numResultPatterns = pattern.getNumResultPatterns();
 
   // First register all symbols bound to ops generated in result patterns.
for (int i = numResultPatterns - 1; i >= 0; --i) {
     }
   }
 
   if (offsets.front() > 0) {
     const char error[] = "no enough values generated to replace the matched op";
     PrintFatalError(loc, error);
   }
 
-  os.indent(4) << "auto odsLoc = rewriter.getFusedLoc({";
+  os << "auto odsLoc = rewriter.getFusedLoc({";
   for (int i = 0, e = pattern.getSourcePattern().getNumOps(); i != e; ++i) {
     os << (i ? ", " : "") << "tblgen_ops[" << i << "]->getLoc()";
   }
   os << "}); (void)odsLoc;\n";
 
   // Process auxiliary result patterns.
   for (int i = 0; i < replStartIndex; ++i) {
     DagNode resultTree = pattern.getResultPattern(i);
     auto val = handleResultPattern(resultTree, offsets[i], 0);
     // Normal op creation will be streamed to `os` by the above call; but
     // NativeCodeCall will only be materialized to `os` if it is used. Here
     // we are handling auxiliary patterns so we want the side effect even if
     // NativeCodeCall is not replacing matched root op's results.
     if (resultTree.isNativeCodeCall())
-      os.indent(4) << val << ";\n";
+      os << val << ";\n";
   }
 
   if (numExpectedResults == 0) {
     assert(replStartIndex >= numResultPatterns &&
            "invalid auxiliary vs. replacement pattern division!");
     // No result to replace. Just erase the op.
-    os.indent(4) << "rewriter.eraseOp(op0);\n";
+    os << "rewriter.eraseOp(op0);\n";
   } else {
     // Process replacement result patterns.
-    os.indent(4)
-        << "::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;\n";
+    os << "::llvm::SmallVector<::mlir::Value, 4> tblgen_repl_values;\n";
     for (int i = replStartIndex; i < numResultPatterns; ++i) {
       DagNode resultTree = pattern.getResultPattern(i);
       auto val = handleResultPattern(resultTree, offsets[i], 0);
-      os.indent(4) << "\n";
+      os << "\n";
       // Resolve each symbol for all range use so that we can loop over them.
       // We need an explicit cast to `SmallVector` to capture the cases where
       // `{0}` resolves to an `Operation::result_range` as well as cases that
       // are not iterable (e.g. vector that gets wrapped in additional braces by
       // RewriterGen).
       // TODO: Revisit the need for materializing a vector.
       os << symbolInfoMap.getAllRangeUse(
           val,
-          "    for (auto v : ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{ "
-          "tblgen_repl_values.push_back(v); }",
+          "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
+          "  tblgen_repl_values.push_back(v);\n}\n",
           "\n");
     }
-    os.indent(4) << "\n";
-    os.indent(4) << "rewriter.replaceOp(op0, tblgen_repl_values);\n";
+    os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }
 
   LLVM_DEBUG(llvm::dbgs() << "--- done emitting rewrite logic ---\n");
 }
 
 std::string PatternEmitter::getUniqueSymbol(const Operator *op) {
   return std::string(
       formatv("tblgen_{0}_{1}", op->getCppClassName(), nextValueId++));
std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
   } else {
     resultValue = std::string(tree.getSymbol());
     // Strip the index to get the name for the value pack and use it to name the
     // local variable for the op.
     valuePackName = std::string(SymbolInfoMap::getValuePackName(resultValue));
   }
 
   // Create the local variable for this op.
-  os.indent(4) << formatv("{0} {1};\n", resultOp.getQualCppClassName(),
-                          valuePackName);
-  os.indent(4) << "{\n";
+  os << formatv("{0} {1};\n{{\n", resultOp.getQualCppClassName(),
 
   // Right now ODS don't have general type inference support. Except a few
   // special cases listed below, DRR needs to supply types for all results
   // when building an op.
   bool isSameOperandsAndResultType =
       resultOp.getTrait("::mlir::OpTrait::SameOperandsAndResultType");
   bool useFirstAttr =
       resultOp.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType");
 
   if (isSameOperandsAndResultType || useFirstAttr) {
     // We know how to deduce the result type for ops with these traits and we've
     // generated builders taking aggregate parameters. Use those builders to
     // create the ops.
 
     // First prepare local variables for op arguments used in builder call.
     createAggregateLocalVarsForOpArgs(tree, childNodeNames);
 
     // Then create the op.
-    os.indent(6) << formatv(
-        "{0} = rewriter.create<{1}>({2}, tblgen_values, tblgen_attrs);\n",
+    os.scope("", "\n}\n").os << formatv(
+        "{0} = rewriter.create<{1}>({2}, tblgen_values, tblgen_attrs);",
         valuePackName, resultOp.getQualCppClassName(), locToUse);
-    os.indent(4) << "}\n";
     return resultValue;
   }
 
   bool usePartialResults = valuePackName != resultValue;
 
   if (usePartialResults || depth > 0 || resultIndex < 0) {
     // For these cases (broadcastable ops, op results used both as auxiliary
     // values and replacement values, ops in nested patterns, auxiliary ops), we
     // still need to supply the result types when building the op. But because
     // we don't generate a builder automatically with ODS for them, it's the
     // developer's responsibility to make sure such a builder (with result type
     // deduction ability) exists. We go through the separate-parameter builder
     // here given that it's easier for developers to write compared to
     // aggregate-parameter builders.
     createSeparateLocalVarsForOpArgs(tree, childNodeNames);
 
-    os.indent(6) << formatv("{0} = rewriter.create<{1}>({2}", valuePackName,
+    os.scope().os << formatv("{0} = rewriter.create<{1}>({2}", valuePackName,
                              resultOp.getQualCppClassName(), locToUse);
     supplyValuesForOpArgs(tree, childNodeNames);
-    os << "\n      );\n";
-    os.indent(4) << "}\n";
+    os << "\n  );\n}\n";
     return resultValue;
   }
 
   // If depth == 0 and resultIndex >= 0, it means we are replacing the values
   // generated from the source pattern root op. Then we can use the source
   // pattern's value types to determine the value type of the generated op
   // here.
 
   // First prepare local variables for op arguments used in builder call.
   createAggregateLocalVarsForOpArgs(tree, childNodeNames);
 
   // Then prepare the result types. We need to specify the types for all
   // results.
-  os.indent(6) << formatv("::mlir::SmallVector<::mlir::Type, 4> tblgen_types; "
-                          "(void)tblgen_types;\n");
+  os.indent() << formatv("::mlir::SmallVector<::mlir::Type, 4> tblgen_types; "
   int numResults = resultOp.getNumResults();
   if (numResults != 0) {
     for (int i = 0; i < numResults; ++i)
-      os.indent(6) << formatv("for (auto v : castedOp0.getODSResults({0})) {{"
-                              "tblgen_types.push_back(v.getType()); }\n",
-                              resultIndex + i);
+      os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n"
+                    "  tblgen_types.push_back(v.getType());\n}\n",
   }
-  os.indent(6) << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, "
-                          "tblgen_values, tblgen_attrs);\n",
-                          valuePackName, resultOp.getQualCppClassName(),
-                          locToUse);
-  os.indent(4) << "}\n";
+  os << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, "
+                valuePackName, resultOp.getQualCppClassName(), locToUse);
+  os.unindent() << "}\n";
   return resultValue;
 }
 
 void PatternEmitter::createSeparateLocalVarsForOpArgs(
     DagNode node, ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);
 
   // Now prepare operands used for building this op:
   // * If the operand is non-variadic, we create a `Value` local variable.
   // * If the operand is variadic, we create a `SmallVector<Value>` local
   //   variable.
 
   int valueIndex = 0; // An index for uniquing local variable names.
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     const auto *operand =
         resultOp.getArg(argIndex).dyn_cast<NamedTypeConstraint *>();
-    if (!operand) {
-      // We do not need special handling for attributes.
+    if (!operand)
       continue;
-    }
 
+    raw_indented_ostream::DelimitedScope scope(os);
     std::string varName;
     if (operand->isVariadic()) {
       varName = std::string(formatv("tblgen_values_{0}", valueIndex++));
-      os.indent(6) << formatv("::mlir::SmallVector<::mlir::Value, 4> {0};\n",
-                              varName);
+      os << formatv("::mlir::SmallVector<::mlir::Value, 4> {0};\n", varName);
       std::string range;
       if (node.isNestedDagArg(argIndex)) {
         range = childNodeNames[argIndex];
       } else {
         range = std::string(node.getArgName(argIndex));
       }
       // Resolve the symbol for all range use so that we have a uniform way of
       // capturing the values.
       range = symbolInfoMap.getValueAndRangeUse(range);
-      os.indent(6) << formatv("for (auto v : {0}) {1}.push_back(v);\n", range,
-                              varName);
+      os << formatv("for (auto v: {0}) {{\n  {1}.push_back(v);\n}\n", range,
     } else {
       varName = std::string(formatv("tblgen_value_{0}", valueIndex++));
-      os.indent(6) << formatv("::mlir::Value {0} = ", varName);
+      os << formatv("::mlir::Value {0} = ", varName);
       if (node.isNestedDagArg(argIndex)) {
         os << symbolInfoMap.getValueAndRangeUse(childNodeNames[argIndex]);
       } else {
         DagLeaf leaf = node.getArgAsLeaf(argIndex);
         auto symbol =
             symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex));
         if (leaf.isNativeCodeCall()) {
           os << std::string(

 }
 
 void PatternEmitter::supplyValuesForOpArgs(
     DagNode node, const ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);
   for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();
        argIndex != numOpArgs; ++argIndex) {
     // Start each argument on its own line.
-    (os << ",\n").indent(8);
+    os << ",\n    ";
 
     Argument opArg = resultOp.getArg(argIndex);
     // Handle the case of operand first.
     if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
       if (!operand->name.empty())
         os << "/*" << operand->name << "=*/";
       os << childNodeNames.lookup(argIndex);
       continue;
for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();
     }
   }
 }
 
 void PatternEmitter::createAggregateLocalVarsForOpArgs(
     DagNode node, const ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);
 
-  os.indent(6) << formatv("::mlir::SmallVector<::mlir::Value, 4> "
-                          "tblgen_values; (void)tblgen_values;\n");
-  os.indent(6) << formatv("::mlir::SmallVector<::mlir::NamedAttribute, 4> "
-                          "tblgen_attrs; (void)tblgen_attrs;\n");
+  auto scope = os.scope();
+  os << formatv("::mlir::SmallVector<::mlir::Value, 4> "
+  os << formatv("::mlir::SmallVector<::mlir::NamedAttribute, 4> "
 
   const char *addAttrCmd =
-      "if (auto tmpAttr = {1}) "
-      "tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), tmpAttr);\n";
+      "if (auto tmpAttr = {1}) {\n"
+      "  tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), "
+      "tmpAttr);\n}\n";
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
       // The argument in the op definition.
       auto opArgName = resultOp.getArgName(argIndex);
       if (auto subTree = node.getArgAsNestedDag(argIndex)) {
         if (!subTree.isNativeCodeCall())
           PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                                "for creating attribute");
-        os.indent(6) << formatv(addAttrCmd, opArgName,
-                                handleReplaceWithNativeCodeCall(subTree));
+        os << formatv(addAttrCmd, opArgName,
       } else {
         auto leaf = node.getArgAsLeaf(argIndex);
         // The argument in the result DAG pattern.
         auto patArgName = node.getArgName(argIndex);
-        os.indent(6) << formatv(addAttrCmd, opArgName,
-                                handleOpArgument(leaf, patArgName));
+        os << formatv(addAttrCmd, opArgName,
       }
       continue;
     }
 
     const auto *operand =
         resultOp.getArg(argIndex).get<NamedTypeConstraint *>();
     std::string varName;
     if (operand->isVariadic()) {
       std::string range;
       if (node.isNestedDagArg(argIndex)) {
         range = childNodeNames.lookup(argIndex);
       } else {
         range = std::string(node.getArgName(argIndex));
       }
       // Resolve the symbol for all range use so that we have a uniform way of
       // capturing the values.
       range = symbolInfoMap.getValueAndRangeUse(range);
-      os.indent(6) << formatv(
-          "for (auto v : {0}) tblgen_values.push_back(v);\n", range);
+      os << formatv("for (auto v: {0}) {{\n  tblgen_values.push_back(v);\n}\n",
+                    range);
     } else {
-      os.indent(6) << formatv("tblgen_values.push_back(", varName);
+      os << formatv("tblgen_values.push_back(", varName);
       if (node.isNestedDagArg(argIndex)) {
         os << symbolInfoMap.getValueAndRangeUse(
             childNodeNames.lookup(argIndex));
       } else {
         DagLeaf leaf = node.getArgAsLeaf(argIndex);
         auto symbol =
             symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex));
         if (leaf.isNativeCodeCall()) {

mlir/unittests/Support/CMakeLists.txt

+add_mlir_unittest(MLIRSupportTests
+  IndentedOstreamTest.cpp
+)
+
+target_link_libraries(MLIRSupportTests
+  PRIVATE MLIRSupportIdentedOstream MLIRSupport)

mlir/unittests/Support/IndentedOstreamTest.cpp

+//===- IndentedOstreamTest.cpp - Indented raw ostream Tests ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/IndentedOstream.h"
+#include "gmock/gmock.h"
+
+using namespace mlir;
+using ::testing::StrEq;
+
+TEST(FormatTest, SingleLine) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  raw_indented_ostream ros(os);
+  ros << 10;
+  ros.flush();
+  EXPECT_THAT(os.str(), StrEq("10"));
+}
+
+TEST(FormatTest, SimpleMultiLine) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  raw_indented_ostream ros(os);
+  ros << "a";
+  ros << "b";
+  ros << "\n";
+  ros << "c";
+  ros << "\n";
+  ros.flush();
+  EXPECT_THAT(os.str(), StrEq("ab\nc\n"));
+}
+
+TEST(FormatTest, SimpleMultiLineIndent) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  raw_indented_ostream ros(os);
+  ros.indent(2) << "a";
+  ros.indent(4) << "b";
+  ros << "\n";
+  ros << "c";
+  ros << "\n";
+  ros.flush();
+  EXPECT_THAT(os.str(), StrEq("  a    b\n    c\n"));
+}
+
+TEST(FormatTest, SingleRegion) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  raw_indented_ostream ros(os);
+  ros << "before\n";
+  {
+    raw_indented_ostream::DelimitedScope scope(ros);
+    ros << "inside " << 10;
+    ros << "\n   two\n";
+    {
+      raw_indented_ostream::DelimitedScope scope(ros, "{\n", "\n}\n");
+      ros << "inner inner";
+    }
+  }
+  ros << "after";
+  ros.flush();
+  const auto *expected =
+      R"(before
+  inside 10
+     two
+  {
+    inner inner
+  }
+after)";
+  EXPECT_THAT(os.str(), StrEq(expected));
+
+  // Repeat the above with inline form.
+  str.clear();
+  ros << "before\n";
+  ros.scope().os << "inside " << 10 << "\n   two\n";
+  ros.scope().os.scope("{\n", "\n}\n").os << "inner inner";
+  ros << "after";
+  ros.flush();
+  EXPECT_THAT(os.str(), StrEq(expected));
+}
+
+TEST(FormatTest, Reindent) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  raw_indented_ostream ros(os);
+
+  // String to print with some additional empty lines at the start and lines
+  // with just spaces.
+  const auto *desc = R"(
+       
+       
+         First line
+                 second line
+                 
+                 
+  )";
+  ros.reindent(desc);
+  ros.flush();
+  const auto *expected =
+      R"(First line
+        second line
+
+
+)";
+  EXPECT_THAT(os.str(), StrEq(expected));
+}