[mlir][NFC] Make 'printOp' public in AsmPrinter

[dcaballe]

Summary

This patch moves the 'printOp' functionality to the public API of
AsmPrinter. No 'parseOp' seems to be needed at this time as existing
APIs are able to parse operations producing results where results are
omitted in the textual form (the LHS of an operation is redundant
when it comes to building the operation itself as it only contains
the result names).

[dcaballe]

mlir/lib/IR/AsmPrinter.cpp

2674

We have too many similar print methods for Operation *. We should rename them but I prefer to get your feedback first.

[rriddle]

Can you add a test in the TestDialect?

mlir/lib/IR/AsmPrinter.cpp

2674

Yeah, it'd be nice to rename them to be explicit about how they are printing. The current ones are really hard to reason about.

[dcaballe]

Renaming + testing

[bondhugula]

mlir/lib/IR/AsmPrinter.cpp

2665

What exactly is "etc." here?

2666

The names print and printFullOp look confusing now: the former is printing more than the latter now.

[dcaballe]

mlir/lib/IR/AsmPrinter.cpp

2665

I guess it refers to the location context but we only have indentation and location for now. I'll remove the etc. part.

2666

I kept print because I saw another print method for blocks but I now see it's the only one? I can rename it to printFullOpWithIndent or printFullOpWithIndentAndLoc but I'm not sure if this will scale well if we add more information to be printed. Perhaps something like printFullOpInContext or something like that?

[bondhugula]

mlir/lib/IR/AsmPrinter.cpp

2666

I don't have any more comments. Please feel free to go with what you think is better.

mlir/include/mlir/IR/OpImplementation.h

virtual void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                                      ArrayRef<StringRef> elidedAttrs = {}) = 0;
 
   /// If the specified operation has attributes, print out an attribute
   /// dictionary prefixed with 'attributes'.
   virtual void
   printOptionalAttrDictWithKeyword(ArrayRef<NamedAttribute> attrs,
                                    ArrayRef<StringRef> elidedAttrs = {}) = 0;
 
+  /// Prints the entire operation with the custom assembly form, if available,
+  /// or the generic assembly form, otherwise.
+  virtual void printCustomOrGenericOp(Operation *op) = 0;
+
   /// Print the entire operation with the default generic assembly form.
   /// If `printOpName` is true, then the operation name is printed (the default)
   /// otherwise it is omitted and the print will start with the operand list.
   virtual void printGenericOp(Operation *op, bool printOpName = true) = 0;
 
   /// Prints a region.
   /// If 'printEntryBlockArgs' is false, the arguments of the
   /// block are not printed. If 'printBlockTerminator' is false, the terminator

mlir/lib/IR/AsmPrinter.cpp

 /// for us to generate aliases only for the attributes and types that would be
 /// in the output, and trims down unnecessary output.
 class DummyAliasOperationPrinter : private OpAsmPrinter {
 public:
   explicit DummyAliasOperationPrinter(const OpPrintingFlags &printerFlags,
                                       AliasInitializer &initializer)
       : printerFlags(printerFlags), initializer(initializer) {}
 
-  /// Print the given operation.
-  void print(Operation *op) {
+  /// Prints the entire operation with the custom assembly form, if available,
+  /// or the generic assembly form, otherwise.
+  void printCustomOrGenericOp(Operation *op) override {
     // Visit the operation location.
     if (printerFlags.shouldPrintDebugInfo())
       initializer.visit(op->getLoc(), /*canBeDeferred=*/true);
 
     // If requested, always print the generic form.
     if (!printerFlags.shouldPrintGenericOpForm()) {
       // Check to see if this is a known operation.  If so, use the registered
       // custom printer hook.
void print(Block *block, bool printBlockArgs = true,
     // requested.
     bool hasTerminator =
         !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
     auto range = llvm::make_range(
         block->begin(),
         std::prev(block->end(),
                   (!hasTerminator || printBlockTerminator) ? 0 : 1));
     for (Operation &op : range)
-      print(&op);
+      printCustomOrGenericOp(&op);
   }
 
   /// Print the given region.
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators,
                    bool printEmptyBlock = false) override {
     if (region.empty())
       return;

 void AliasInitializer::initialize(
     Operation *op, const OpPrintingFlags &printerFlags,
     llvm::MapVector<Attribute, SymbolAlias> &attrToAlias,
     llvm::MapVector<Type, SymbolAlias> &typeToAlias) {
   // Use a dummy printer when walking the IR so that we can collect the
   // attributes/types that will actually be used during printing when
   // considering aliases.
   DummyAliasOperationPrinter aliasPrinter(printerFlags, *this);
-  aliasPrinter.print(op);
+  aliasPrinter.printCustomOrGenericOp(op);
 
   // Initialize the aliases sorted by name.
   initializeAliases(aliasToAttr, attrToAlias, &deferrableAttributes);
   initializeAliases(aliasToType, typeToAlias);
 }
 
 void AliasInitializer::visit(Attribute attr, bool canBeDeferred) {
   if (!visitedAttributes.insert(attr).second) {
public:
   using Impl::printType;
 
   explicit OperationPrinter(raw_ostream &os, AsmStateImpl &state)
       : Impl(os, state), OpAsmPrinter(static_cast<Impl &>(*this)) {}
 
   /// Print the given top-level operation.
   void printTopLevelOperation(Operation *op);
 
-  /// Print the given operation with its indent and location.
-  void print(Operation *op);
-  /// Print the bare location, not including indentation/location/etc.
-  void printOperation(Operation *op);
-  /// Print the given operation in the generic form.
+  /// Print the given operation, including its left-hand side and its right-hand
+  /// side, with its indent and location.
+  void printFullOpWithIndentAndLoc(Operation *op);
+  /// Print the given operation, including its left-hand side and its right-hand
+  /// side, but not including indentation and location.
+  void printFullOp(Operation *op);
+  /// Print the right-hand size of the given operation in the custom or generic
+  /// form.
+  void printCustomOrGenericOp(Operation *op) override;
+  /// Print the right-hand side of the given operation in the generic form.
   void printGenericOp(Operation *op, bool printOpName) override;
 
   /// Print the name of the given block.
   void printBlockName(Block *block);
 
   /// Print the given block. If 'printBlockArgs' is false, the arguments of the
   /// block are not printed. If 'printBlockTerminator' is false, the terminator
   /// operation of the block is not printed.
   void print(Block *block, bool printBlockArgs = true,

 };
 } // namespace
 
 void OperationPrinter::printTopLevelOperation(Operation *op) {
   // Output the aliases at the top level that can't be deferred.
   state.getAliasState().printNonDeferredAliases(os, newLine);
 
   // Print the module.
-  print(op);
+  printFullOpWithIndentAndLoc(op);
   os << newLine;
 
   // Output the aliases at the top level that can be deferred.
   state.getAliasState().printDeferredAliases(os, newLine);
 
   // Output any file level metadata.
   printFileMetadataDictionary(op);
 }
if (!omitType) {
     os << ": ";
     printType(arg.getType());
   }
   printOptionalAttrDict(argAttrs);
   // TODO: We should allow location aliases on block arguments.
   printTrailingLocation(arg.getLoc(), /*allowAlias*/ false);
 }
 
-void OperationPrinter::print(Operation *op) {
+void OperationPrinter::printFullOpWithIndentAndLoc(Operation *op) {
   // Track the location of this operation.
   state.registerOperationLocation(op, newLine.curLine, currentIndent);
 
   os.indent(currentIndent);
-  printOperation(op);
+  printFullOp(op);
   printTrailingLocation(op->getLoc());
   if (printerFlags.shouldPrintValueUsers())
     printUsersComment(op);
 }
 
-void OperationPrinter::printOperation(Operation *op) {
+void OperationPrinter::printFullOp(Operation *op) {
   if (size_t numResults = op->getNumResults()) {
     auto printResultGroup = [&](size_t resultNo, size_t resultCount) {
       printValueID(op->getResult(resultNo), /*printResultNo=*/false);
       if (resultCount > 1)
         os << ':' << resultCount;
     };
 
     // Check to see if this operation has multiple result groups.
if (size_t numResults = op->getNumResults()) {
 
     } else {
       printResultGroup(/*resultNo=*/0, /*resultCount=*/numResults);
     }
 
     os << " = ";
   }
 
-  // If requested, always print the generic form.
-  if (!printerFlags.shouldPrintGenericOpForm()) {
-    // Check to see if this is a known operation. If so, use the registered
-    // custom printer hook.
-    if (auto opInfo = op->getRegisteredInfo()) {
-      opInfo->printAssembly(op, *this, defaultDialectStack.back());
-      return;
-    }
-    // Otherwise try to dispatch to the dialect, if available.
-    if (Dialect *dialect = op->getDialect()) {
-      if (auto opPrinter = dialect->getOperationPrinter(op)) {
-        // Print the op name first.
-        StringRef name = op->getName().getStringRef();
-        // Only drop the default dialect prefix when it cannot lead to
-        // ambiguities.
-        if (name.count('.') == 1)
-          name.consume_front((defaultDialectStack.back() + ".").str());
-        os << name;
-
-        // Print the rest of the op now.
-        opPrinter(op, *this);
-        return;
-      }
-    }
-  }
-
-  // Otherwise print with the generic assembly form.
-  printGenericOp(op, /*printOpName=*/true);
+  printCustomOrGenericOp(op);
 }
 
 void OperationPrinter::printUsersComment(Operation *op) {
   unsigned numResults = op->getNumResults();
   if (!numResults && op->getNumOperands()) {
     os << " // id: ";
     printOperationID(op);
   } else if (numResults && op->use_empty()) {
void OperationPrinter::printUserIDs(Operation *user, bool prefixComma) {
   if (!user->getNumResults()) {
     printOperationID(user);
   } else {
     interleaveComma(user->getResults(),
                     [this](Value result) { printValueID(result); });
   }
 }
 
+void OperationPrinter::printCustomOrGenericOp(Operation *op) {
+  // If requested, always print the generic form.
+  if (!printerFlags.shouldPrintGenericOpForm()) {
+    // Check to see if this is a known operation. If so, use the registered
+    // custom printer hook.
+    if (auto opInfo = op->getRegisteredInfo()) {
+      opInfo->printAssembly(op, *this, defaultDialectStack.back());
+      return;
+    }
+    // Otherwise try to dispatch to the dialect, if available.
+    if (Dialect *dialect = op->getDialect()) {
+      if (auto opPrinter = dialect->getOperationPrinter(op)) {
+        // Print the op name first.
+        StringRef name = op->getName().getStringRef();
+        // Only drop the default dialect prefix when it cannot lead to
+        // ambiguities.
+        if (name.count('.') == 1)
+          name.consume_front((defaultDialectStack.back() + ".").str());
+        os << name;
+
+        // Print the rest of the op now.
+        opPrinter(op, *this);
+        return;
+      }
+    }
+  }
+
+  // Otherwise print with the generic assembly form.
+  printGenericOp(op, /*printOpName=*/true);
+}
+
 void OperationPrinter::printGenericOp(Operation *op, bool printOpName) {
   if (printOpName)
     printEscapedString(op->getName().getStringRef());
   os << '(';
   interleaveComma(op->getOperands(), [&](Value value) { printValueID(value); });
   os << ')';
 
   // For terminators, print the list of successors and their operands.
void OperationPrinter::print(Block *block, bool printBlockArgs,
 
   bool hasTerminator =
       !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   auto range = llvm::make_range(
       block->begin(),
       std::prev(block->end(),
                 (!hasTerminator || printBlockTerminator) ? 0 : 1));
   for (auto &op : range) {
-    print(&op);
+    printFullOpWithIndentAndLoc(&op);
     os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
 void OperationPrinter::printValueID(Value value, bool printResultNo,
                                     raw_ostream *streamOverride) const {
   state.getSSANameState().printValueID(value, printResultNo,
void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
   print(os, state);
 }
 void Operation::print(raw_ostream &os, AsmState &state) {
   OperationPrinter printer(os, state.getImpl());
   if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope()) {
     state.getImpl().initializeAliases(this);
     printer.printTopLevelOperation(this);
   } else {
-    printer.print(this);
+    printer.printFullOpWithIndentAndLoc(this);
   }
 }
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());
   llvm::errs() << "\n";
 }
 

mlir/test/IR/print-op-custom-or-generic.mlir

+// # RUN: mlir-opt %s -split-input-file | FileCheck %s
+// # RUN: mlir-opt %s -mlir-print-op-generic -split-input-file  | FileCheck %s --check-prefix=GENERIC
+
+// Check that `printCustomOrGenericOp` and `printGenericOp` print the right
+// assembly format. For operations without custom format, both should print the
+// generic format.
+
+// CHECK-LABEL: func @op_with_custom_printer
+// CHECK-GENERIC-LABEL: "func"()
+func.func @op_with_custom_printer() {
+  %x = test.string_attr_pretty_name
+  // CHECK: %x = test.string_attr_pretty_name
+  // GENERIC: %0 = "test.string_attr_pretty_name"()
+  return
+  // CHECK: return
+  // GENERIC: "func.return"()
+}
+
+// -----
+
+// CHECK-LABEL: func @op_without_custom_printer
+// CHECK-GENERIC: "func"()
+func.func @op_without_custom_printer() {
+  // CHECK: "test.result_type_with_trait"() : () -> !test.test_type_with_trait
+  // GENERIC: "test.result_type_with_trait"() : () -> !test.test_type_with_trait
+  "test.result_type_with_trait"() : () -> !test.test_type_with_trait
+  return
+}

mlir/test/IR/print-op-generic.mlir

-// # RUN: mlir-opt %s | FileCheck %s
-// # RUN: mlir-opt %s --mlir-print-op-generic  | FileCheck %s --check-prefix=GENERIC
-
-// CHECK-LABEL: func @pretty_names
-// CHECK-GENERIC: "func"()
-func.func @pretty_names() {
-  %x = test.string_attr_pretty_name
-  // CHECK: %x = test.string_attr_pretty_name
-  // GENERIC: %0 = "test.string_attr_pretty_name"()
-  return
-  // CHECK: return
-  // GENERIC: "func.return"()
-}