[mlir] Print a comment when printing an op that failed to verify

[Mogball]

Summary

This patch changes the printer to print a comment when an op fails to
verify.

func @foo() {
  // Verification failed, printing generic form
  "test.invalid_op"() : () -> ()
  return
}

The verification of the op is done lazily instead.

Fixes #57435

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

2783–2810

This is going to run on every single operation (which in turn will recursively verify), we should not be doing that. AsmState is still the right place to do verification.

[rriddle]

I would also rather have the comment at the top of the IR dump anyways, otherwise it's going to be difficult to spot if the dump is non-trivial.

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

2783–2810

Right. But I'm struggling to find a place to print the comment otherwise. It can't be printed until an OperationPrinter is instantiated. I can pass a flag from AsmState to OperationPrinter and print the comment whenever the latter is instantiated but that isn't very pretty :( Any suggestions?

[lattner]

Awesome, thank you for improving this!

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

2783–2810

Passing a flag seems like the best thing to me IMO (which doesn't need to be publicly exposed, you could expose whatever is necessary just in AsmStateImpl). The reasoning is that we need the flags to be correct in AsmState, because that directly impacts how values are numbered, aliases are initialized, etc. AsmState is also potentially used to separately print operations, and if we don't have that be the source of truth for flags, we'll end up with inconsistent results when printing different operations that use the same AsmState.

mlir/include/mlir-c/IR.h

 
 /// 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.
 MLIR_CAPI_EXPORTED void
 mlirOpPrintingFlagsUseLocalScope(MlirOpPrintingFlags flags);
 
+/// Do not verify the operation when using custom operation printers.
+MLIR_CAPI_EXPORTED void
+mlirOpPrintingFlagsAssumeVerified(MlirOpPrintingFlags flags);
+
 //===----------------------------------------------------------------------===//
 // Operation API.
 //===----------------------------------------------------------------------===//
 
 /// Creates an operation and transfers ownership to the caller.
 /// Note that caller owned child objects are transferred in this call and must
 /// not be further used. Particularly, this applies to any regions added to
 /// the state (the implementation may invalidate any such pointers).

mlir/lib/Bindings/Python/IRCore.cpp

void PyOperationBase::print(py::object fileObject, bool binary,
                             bool enableDebugInfo, bool prettyDebugInfo,
                             bool printGenericOpForm, bool useLocalScope,
                             bool assumeVerified) {
   PyOperation &operation = getOperation();
   operation.checkValid();
   if (fileObject.is_none())
     fileObject = py::module::import("sys").attr("stdout");
 
-  if (!assumeVerified && !printGenericOpForm &&
-      !mlirOperationVerify(operation)) {
-    std::string message("// Verification failed, printing generic form\n");
-    if (binary) {
-      fileObject.attr("write")(py::bytes(message));
-    } else {
-      fileObject.attr("write")(py::str(message));
-    }
-    printGenericOpForm = true;
-  }
-
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   if (largeElementsLimit)
     mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
   if (enableDebugInfo)
     mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/prettyDebugInfo);
   if (printGenericOpForm)
     mlirOpPrintingFlagsPrintGenericOpForm(flags);
   if (useLocalScope)
     mlirOpPrintingFlagsUseLocalScope(flags);
+  if (assumeVerified)
+    mlirOpPrintingFlagsAssumeVerified(flags);
 
   PyFileAccumulator accum(fileObject, binary);
   mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
                               accum.getUserData());
   mlirOpPrintingFlagsDestroy(flags);
 }
 
 py::object PyOperationBase::getAsm(bool binary,

mlir/lib/CAPI/IR/IR.cpp

 void mlirOpPrintingFlagsPrintGenericOpForm(MlirOpPrintingFlags flags) {
   unwrap(flags)->printGenericOpForm();
 }
 
 void mlirOpPrintingFlagsUseLocalScope(MlirOpPrintingFlags flags) {
   unwrap(flags)->useLocalScope();
 }
 
+void mlirOpPrintingFlagsAssumeVerified(MlirOpPrintingFlags flags) {
+  unwrap(flags)->assumeVerified();
+}
+
 //===----------------------------------------------------------------------===//
 // Location API.
 //===----------------------------------------------------------------------===//
 
 MlirLocation mlirLocationFileLineColGet(MlirContext context,
                                         MlirStringRef filename, unsigned line,
                                         unsigned col) {
   return wrap(Location(

mlir/lib/IR/AsmPrinter.cpp

private:
   AsmState::LocationMap *locationMap;
 
   // Allow direct access to the impl fields.
   friend AsmState;
 };
 } // namespace detail
 } // namespace mlir
 
-/// Verifies the operation and switches to generic op printing if verification
-/// fails. We need to do this because custom print functions may fail for
-/// invalid ops.
-static OpPrintingFlags verifyOpAndAdjustFlags(Operation *op,
-                                              OpPrintingFlags printerFlags) {
-  if (printerFlags.shouldPrintGenericOpForm() ||
-      printerFlags.shouldAssumeVerified())
-    return printerFlags;
-
-  LLVM_DEBUG(llvm::dbgs() << DEBUG_TYPE << ": Verifying operation: "
-                          << op->getName() << "\n");
-
-  // Ignore errors emitted by the verifier. We check the thread id to avoid
-  // consuming other threads' errors.
-  auto parentThreadId = llvm::get_threadid();
-  ScopedDiagnosticHandler diagHandler(op->getContext(), [&](Diagnostic &diag) {
-    if (parentThreadId == llvm::get_threadid()) {
-      LLVM_DEBUG({
-        diag.print(llvm::dbgs());
-        llvm::dbgs() << "\n";
-      });
-      return success();
-    }
-    return failure();
-  });
-  if (failed(verify(op))) {
-    LLVM_DEBUG(llvm::dbgs()
-               << DEBUG_TYPE << ": '" << op->getName()
-               << "' failed to verify and will be printed in generic form\n");
-    printerFlags.printGenericOpForm();
-  }
-
-  return printerFlags;
-}
-
 AsmState::AsmState(Operation *op, const OpPrintingFlags &printerFlags,
                    LocationMap *locationMap)
-    : impl(std::make_unique<AsmStateImpl>(
-          op, verifyOpAndAdjustFlags(op, printerFlags), locationMap)) {}
+    : impl(std::make_unique<AsmStateImpl>(op, printerFlags, locationMap)) {}
 AsmState::~AsmState() = default;
 
 const OpPrintingFlags &AsmState::getPrinterFlags() const {
   return impl->getPrinterFlags();
 }
 
 void AsmState::attachResourcePrinter(
     std::unique_ptr<AsmResourcePrinter> printer) {
void OperationPrinter::print(Operation *op) {
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
   if (printerFlags.shouldPrintValueUsers())
     printUsersComment(op);
 }
 
 void OperationPrinter::printOperation(Operation *op) {
+  // Switch to the generic printer if verification fails. We need to do this
+  // because custom print functions may fail for invalid ops.
+  if (!printerFlags.shouldPrintGenericOpForm() &&
+      !printerFlags.shouldAssumeVerified()) {
+    LLVM_DEBUG(llvm::dbgs() << DEBUG_TYPE << ": Verifying operation: "
+                            << op->getName() << "\n");
+    // Ignore errors emitted by the verifier. We check the thread id to avoid
+    // consuming other threads' errors.
+    auto parentThreadId = llvm::get_threadid();
+    auto handler = [&](Diagnostic &diag) {
+      if (parentThreadId != llvm::get_threadid())
+        return failure();
+      LLVM_DEBUG({
+        diag.print(llvm::dbgs());
+        llvm::dbgs() << "\n";
+      });
+      return success();
+    };
+    ScopedDiagnosticHandler diagHandler(op->getContext(), handler);
+    if (failed(verify(op))) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << DEBUG_TYPE << ": '" << op->getName()
+                 << "' failed to verify and will be printed in generic form\n");
+      os << "// Verification failed, printing generic form";
+      printNewline();
+      printerFlags.printGenericOpForm();
+    }
+  }
+
   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.

mlir/test/IR/print-ir-invalid.mlir

 // # RUN: mlir-opt -test-print-invalid %s | FileCheck %s
 // # RUN: mlir-opt -test-print-invalid %s --mlir-print-assume-verified  | FileCheck %s --check-prefix=ASSUME-VERIFIED
 
 // The pass creates some ops and prints them to stdout, the input is just an
 // empty module.
 module {}
 
 // The operation is invalid because the body does not have a terminator, print
 // the generic form.
 // CHECK:      Invalid operation:
+// CHECK-NEXT: // Verification failed, printing generic form
 // CHECK-NEXT: "func.func"() ({
 // CHECK-NEXT: ^bb0:
 // CHECK-NEXT: })
 // CHECK-SAME: sym_name = "test"
 
 // The operation is valid because the body has a terminator, print the custom
 // form.
 // CHECK:      Valid operation: