Add a `skipRegion()` feature to the OpPrintingFlags for MLIR ASM printer

[mehdi_amini]

Summary

This is a convenient flag for context where we intend to summarize a top-level
operation without the full-blown regions it may hold.

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

681

Why the comment, instead of something like ...?

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

681

Felt more explicit, but I didn't consider ... before either...

[mehdi_amini]

Use ... instead of "/*skip region*/" comment

mlir/include/mlir/IR/OperationSupport.h

public:
   /// 'prettyForm' is set to true, debug information is printed in a more
   /// readable 'pretty' form. Note: The IR generated with 'prettyForm' is not
   /// parsable.
   OpPrintingFlags &enableDebugInfo(bool enable = true, bool prettyForm = false);
 
   /// Always print operations in the generic form.
   OpPrintingFlags &printGenericOpForm();
 
+  /// Skip printing regions.
+  OpPrintingFlags &skipRegions();
+
   /// Do not verify the operation when using custom operation printers.
   OpPrintingFlags &assumeVerified();
 
   /// Use local scope when printing the operation. This allows for using the
   /// printer in a more localized and thread-safe setting, but may not
   /// necessarily be identical to what the IR will look like when dumping
   /// the full module.
   OpPrintingFlags &useLocalScope();
public:
   bool shouldPrintDebugInfo() const;
 
   /// Return if debug information should be printed in the pretty form.
   bool shouldPrintDebugInfoPrettyForm() const;
 
   /// Return if operations should be printed in the generic form.
   bool shouldPrintGenericOpForm() const;
 
+  /// Return if regions should be skipped.
+  bool shouldSkipRegions() const;
+
   /// Return if operation verification should be skipped.
   bool shouldAssumeVerified() const;
 
   /// Return if the printer should use local scope when dumping the IR.
   bool shouldUseLocalScope() const;
 
   /// Return if the printer should print users of values.
   bool shouldPrintValueUsers() const;
 
 private:
   /// Elide large elements attributes if the number of elements is larger than
   /// the upper limit.
   std::optional<int64_t> elementsAttrElementLimit;
 
   /// Print debug information.
   bool printDebugInfoFlag : 1;
   bool printDebugInfoPrettyFormFlag : 1;
 
   /// Print operations in the generic form.
   bool printGenericOpFormFlag : 1;
 
+  /// Always skip Regions.
+  bool skipRegionsFlag : 1;
+
   /// Skip operation verification.
   bool assumeVerifiedFlag : 1;
 
   /// Print operations with numberings local to the current operation.
   bool printLocalScope : 1;
 
   /// Print users of values.
   bool printValueUsersFlag : 1;

mlir/lib/IR/AsmPrinter.cpp

 void mlir::registerAsmPrinterCLOptions() {
   // Make sure that the options struct has been initialized.
   *clOptions;
 }
 
 /// Initialize the printing flags with default supplied by the cl::opts above.
 OpPrintingFlags::OpPrintingFlags()
     : printDebugInfoFlag(false), printDebugInfoPrettyFormFlag(false),
-      printGenericOpFormFlag(false), assumeVerifiedFlag(false),
-      printLocalScope(false), printValueUsersFlag(false) {
+      printGenericOpFormFlag(false), skipRegionsFlag(false),
+      assumeVerifiedFlag(false), printLocalScope(false),
+      printValueUsersFlag(false) {
   // Initialize based upon command line options, if they are available.
   if (!clOptions.isConstructed())
     return;
   if (clOptions->elideElementsAttrIfLarger.getNumOccurrences())
     elementsAttrElementLimit = clOptions->elideElementsAttrIfLarger;
   printDebugInfoFlag = clOptions->printDebugInfoOpt;
   printDebugInfoPrettyFormFlag = clOptions->printPrettyDebugInfoOpt;
   printGenericOpFormFlag = clOptions->printGenericOpFormOpt;

 }
 
 /// Always print operations in the generic form.
 OpPrintingFlags &OpPrintingFlags::printGenericOpForm() {
   printGenericOpFormFlag = true;
   return *this;
 }
 
+/// Always skip Regions.
+OpPrintingFlags &OpPrintingFlags::skipRegions() {
+  skipRegionsFlag = true;
+  return *this;
+}
+
 /// Do not verify the operation when using custom operation printers.
 OpPrintingFlags &OpPrintingFlags::assumeVerified() {
   assumeVerifiedFlag = true;
   return *this;
 }
 
 /// Use local scope when printing the operation. This allows for using the
 /// printer in a more localized and thread-safe setting, but may not necessarily
bool OpPrintingFlags::shouldPrintDebugInfoPrettyForm() const {
   return printDebugInfoPrettyFormFlag;
 }
 
 /// Return if operations should be printed in the generic form.
 bool OpPrintingFlags::shouldPrintGenericOpForm() const {
   return printGenericOpFormFlag;
 }
 
+/// Return if Region should be skipped.
+bool OpPrintingFlags::shouldSkipRegions() const { return skipRegionsFlag; }
+
 /// Return if operation verification should be skipped.
 bool OpPrintingFlags::shouldAssumeVerified() const {
   return assumeVerifiedFlag;
 }
 
 /// Return if the printer should use local scope when dumping the IR.
 bool OpPrintingFlags::shouldUseLocalScope() const { return printLocalScope; }
 
void printCustomOrGenericOp(Operation *op) override {
     // Otherwise print with the generic assembly form.
     printGenericOp(op);
   }
 
 private:
   /// Print the given operation in the generic form.
   void printGenericOp(Operation *op, bool printOpName = true) override {
     // Consider nested operations for aliases.
+    if (!printerFlags.shouldSkipRegions()) {
       for (Region &region : op->getRegions())
         printRegion(region, /*printEntryBlockArgs=*/true,
                     /*printBlockTerminators=*/true);
+    }
 
     // Visit all the types used in the operation.
     for (Type type : op->getOperandTypes())
       printType(type);
     for (Type type : op->getResultTypes())
       printType(type);
 
     // Consider the attributes of the operation for aliases.
private:
   }
 
   /// Print the given region.
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators,
                    bool printEmptyBlock = false) override {
     if (region.empty())
       return;
+    if (printerFlags.shouldSkipRegions()) {
+      os << "{/*skip region*/}";
+      return;
+    }
 
     auto *entryBlock = &region.front();
     print(entryBlock, printEntryBlockArgs, printBlockTerminators);
     for (Block &b : llvm::drop_begin(region, 1))
       print(&b);
   }
 
   void printRegionArgument(BlockArgument arg, ArrayRef<NamedAttribute> argAttrs,
if (op->getNumSuccessors() != 0) {
     interleaveComma(op->getSuccessors(),
                     [&](Block *successor) { printBlockName(successor); });
     os << ']';
   }
 
   // Print regions.
   if (op->getNumRegions() != 0) {
     os << " (";
+    if (!printerFlags.shouldSkipRegions()) {
       interleaveComma(op->getRegions(), [&](Region &region) {
         printRegion(region, /*printEntryBlockArgs=*/true,
                     /*printBlockTerminators=*/true, /*printEmptyBlock=*/true);
       });
+    } else {
+      os << "/*skip " << op->getNumRegions() << " regions*/";
+    }
     os << ')';
   }
 
   auto attrs = op->getAttrs();
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
   os << " : ";
void OperationPrinter::printSuccessorAndUseList(Block *successor,
   interleaveComma(succOperands,
                   [this](Value operand) { printType(operand.getType()); });
   os << ')';
 }
 
 void OperationPrinter::printRegion(Region &region, bool printEntryBlockArgs,
                                    bool printBlockTerminators,
                                    bool printEmptyBlock) {
+  if (printerFlags.shouldSkipRegions()) {
+    os << "{/*skip region*/}";
+    return;
+  }
   os << "{" << newLine;
   if (!region.empty()) {
     auto restoreDefaultDialect =
         llvm::make_scope_exit([&]() { defaultDialectStack.pop_back(); });
     if (auto iface = dyn_cast<OpAsmOpInterface>(region.getParentOp()))
       defaultDialectStack.push_back(iface.getDefaultDialect());
     else
       defaultDialectStack.push_back("");

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

std::optional<lsp::Hover> MLIRDocument::buildHoverForOperation(
   // Add the operation name to the hover.
   os << "\"" << op.op->getName() << "\"";
   if (SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op.op))
     os << " : " << symbol.getVisibility() << " @" << symbol.getName() << "";
   os << "\n\n";
 
   os << "Generic Form:\n\n```mlir\n";
 
-  // Temporary drop the regions of this operation so that they don't get
-  // printed in the output. This helps keeps the size of the output hover
-  // small.
-  SmallVector<std::unique_ptr<Region>> regions;
-  for (Region &region : op.op->getRegions()) {
-    regions.emplace_back(std::make_unique<Region>());
-    regions.back()->takeBody(region);
-  }
-
-  op.op->print(
-      os, OpPrintingFlags().printGenericOpForm().elideLargeElementsAttrs());
+  op.op->print(os, OpPrintingFlags()
+                       .printGenericOpForm()
+                       .elideLargeElementsAttrs()
+                       .skipRegions());
   os << "\n```\n";
 
-  // Move the regions back to the current operation.
-  for (Region &region : op.op->getRegions())
-    region.takeBody(*regions.back());
-
   return hover;
 }
 
 lsp::Hover MLIRDocument::buildHoverForOperationResult(SMRange hoverRange,
                                                       Operation *op,
                                                       unsigned resultStart,
                                                       unsigned resultEnd,
                                                       SMLoc posLoc) {

mlir/test/mlir-lsp-server/hover.test

// Hover on a region operation.
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":0,"character":1}
 }}
 //      CHECK:  "id": 5,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
-// CHECK-NEXT:      "value": "\"func.func\" : public @foo\n\nGeneric Form:\n\n```mlir\n\"func.func\"() ({\n}) {function_type = (i1) -> (), sym_name = \"foo\"} : () -> ()\n```\n"
+// CHECK-NEXT:      "value": "\"func.func\" : public @foo\n\nGeneric Form:\n\n```mlir\n\"func.func\"() (/*skip 1 regions*/) {function_type = (i1) -> (), sym_name = \"foo\"} : () -> ()\n```\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 9,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 0,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a symbol reference.
 {"jsonrpc":"2.0","id":6,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":0,"character":13}
 }}
 //      CHECK:  "id": 6,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
-// CHECK-NEXT:      "value": "\"func.func\" : public @foo\n\nGeneric Form:\n\n```mlir\n\"func.func\"() ({\n}) {function_type = (i1) -> (), sym_name = \"foo\"} : () -> ()\n```\n"
+// CHECK-NEXT:      "value": "\"func.func\" : public @foo\n\nGeneric Form:\n\n```mlir\n\"func.func\"() (/*skip 1 regions*/) {function_type = (i1) -> (), sym_name = \"foo\"} : () -> ()\n```\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 14,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 10,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":7,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}