[mlir] Fix dumping invalid ops

[sgrechanik]

Summary

This patch fixes the crash when printing some ops (like affine.for and
scf.for) when they are dumped in invalid state, e.g. during pattern
application. Now the AsmState constructor verifies the operation
first and switches to generic operation printing when the verification
fails. Also operations are now printed in generic form when emitting
diagnostics and the severity level is Error.

[rriddle]

We shouldn't be baking in verification checking into the custom printers. If anything we should have a dump/printer API that falls back to generic in the case of verification failure. It isn't feasible/maintainable to re-encode all of the verifier constraints, so we should keep custom printers clean from things they should be able to assume.

[sgrechanik]

We can call opInfo->verifyInvariants(op) in OperationPrinter::printOperation and switch to generic printing on failure, something along the lines:

if (auto opInfo = op->getRegisteredInfo()) {
  // First verify the operation. If the operation fails to verify then the
  // custom print function may fail and we have to use the generic one.
  {
    // Ignore errors emitted by the verifier.
    ScopedDiagnosticHandler diagHandler(
        op->getContext(), [](Diagnostic &) { return success(); });
    if (failed(opInfo->verifyInvariants(op))) {
      printGenericOp(op, /*printOpName=*/true);
      return;
    }
  }
  opInfo->printAssembly(op, *this, defaultDialectStack.back());
  return;
}

However, doing this exposes some bugs in other ops. I'll try to investigate.

[mehdi_amini]

I talked about this with River in the past: I'd call the verifier on the entry point of the printer (the top-level operation to print), which will verify recursively the IR down. And switch to generically printing for the whole IR and not on an op-by-op basis.

[sgrechanik]

Not sure at which level to put the verification, Operation::dump seems to be a good candidate, it will fix the original crash and won't affect other kinds of printing.

[rriddle]

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

1–3

LG, except for these. I'm not sure we should be adding these to the tests. Realistically tests shouldn't be crashing with -debug, but we shouldn't need to add debug individually to all of the tests to verify that.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

2976

I think this should belong to the printer and not just "dump()"

The behavior does not seem desirable to be coupled to the llvm::errs() in itself. We don't want to crash when printing in general.
We had to hardcode this in the python bindings, but it fits better the framework itself.

[sgrechanik]

Fixing this not only for dump uncovered a bug in verification of spirv::ConstantOp, which should be fixed first: https://reviews.llvm.org/D118939

[jpienaar]

mlir/lib/IR/AsmPrinter.cpp

2971

There is (I believe) still a change that an error can be emitted here in multi threaded environment, it would require a new handler to be set after this but before verify below. Not a high probability event though.

[sgrechanik]

• Moved the check to Operation::print and Block::print
• Removed error suppression: errors may be helpful.

mlir/lib/IR/AsmPrinter.cpp

2971

Hmm, yes, you are right, and I don't immediately see how to solve it.

But maybe actually we shouldn't try to suppress error messages here, because without error messages it's hard to tell why the printer switched to generic op printing.

2976

Moved it to Operation::print and Block::print

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

1–3

Do you have a better idea how to tests this?

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

2960

Please run git clang-format on your commit (arc should do it for you if clang-format is in the path I think)

2971

I wouldn't expect diagnostic about verification failure in a place where I call "print" though.

River: is there a safe way to silence these in the context of the printer?

2987

Is this check needed/useful?

Also you checked in the Operation::print is the generic printer flag isn't already set before verifying, should you do the same here?

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

2971

Generally if you only want diagnostics for this thread, you should check the result of llvm::get_threadid() within the handler against the one for the main thread. Something like:

auto parentThreadId = llvm::get_threadid();
auto handlerFn = [&](Diagnostic &diag) {
  return success(parentThreadId == llvm::get_threadid());
};

[sgrechanik]

• Error suppression is back
• fixed formatting errors

mlir/lib/IR/AsmPrinter.cpp

2971

Thanks, I'll use that.

2987

I inserted this check to avoid messing with diagnostic handling if the block is empty.

Here the flags are always the default ones. AsmState contains its own set of flags, but as far as I understand, they are not used for the actual printing. However, I wonder whether AsmState and OperationPrinter flags mismatch would cause any trouble. Maybe we shouldn't do the verification in functions taking AsmState, and verify in more higher level Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) and Block::print(raw_ostream &os) instead?

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

2987

It's generally too late to do verification if the AsmState already exists, the construction of that may call back into operations (and thus crash if invalid). I'd just verify in the non-AsmState functions.

Tangentially though, we should just expose the OpPrintingFlags held by the AsmState and use those (same applies for the other function).

[sgrechanik]

mlir/lib/IR/AsmPrinter.cpp

2987

Tangentially though, we should just expose the OpPrintingFlags held by the AsmState and use those (same applies for the other function).

Then it probably makes sense to do that and perform verification in the AsmState constructor. It doesn't seem to be a lot of code change.
However, I ran into an issue when a verifier tries to emit an error and print the op, and printing the op requires calling the verifier first, resulting in infinite recursion. I need to think how to solve it, something like temporary changing printOpOnDiagnostic, but there doesn't seem to be a thread-safe way to do it.

[sgrechanik]

Ok, let's do it gradually and expose printing flags in AsmState first: https://reviews.llvm.org/D119870

[sgrechanik]

Rebased and moved the verification to the AsmState constructor.

[stellaraccident]

At the python level, we have a flag to assume verified and not reverify. When using systematically, there are often cases where you know verification is successful and can skip doing it again.

We can add that in after this patch if needed, but I feel we should do so.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1256

Shouldn't the verifier failure be printing in generic mode already which should skip the verifier?
I'm not sure why we need this?

[sgrechanik]

mlir/lib/IR/AsmPrinter.cpp

1256

Not always. When the error is emitted like op1.emitError() << op2; then op1 will be printed in generic mode (or not at all), and op2 will be printed with default flags. This actually happens when verifying whether a block has a terminator in mlir/lib/IR/Verifier.cpp:154

Operation &terminator = block.back();
if (!terminator.mightHaveTrait<OpTrait::IsTerminator>())
  return block.back().emitError("block with no terminator, has ")
         << terminator;

And also printing an op may require verification of its parent op because AsmState may be constructed for the parent op.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1256

When the error is emitted like op1.emitError() << op2; then op1 will be printed in generic mode (or not at all), and op2 will be printed with default flags. This actually happens when verifying whether a block has a terminator in mlir/lib/IR/Verifier.cpp:154

But isn't this something we should fix in the verifier? Seems like unintended to me.

And also printing an op may require verification of its parent op because AsmState may be constructed for the parent op.

Similarly: if we print an op with the generic flag, then constructing the AsmState for the parent op could just inherit the flags again?

In general I'm wary of mutable global state when we can do differently: it's not clear to me still that this isn't just a workaround instead of something that is fundamentally needed if the system is consistent.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1256

I think we're always gonna hit custom verifier that may print op without propagating flags, I can't see a better option than what you have here right now, so LGTM for now.

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

1–3

I agree with River: we should keep the testing focused.

what about adding a test with invalid IR and check that the printed op is in generic form?

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1256

We should avoid any kind of global or thread_local state. We should just switch how we handle printing operations in diagnostics, e.g. to just use generic form by default (i.e. change this https://github.com/llvm/llvm-project/blob/8e7995884a6525b44fe2f71318883ba2fec2d972/mlir/lib/IR/Diagnostics.cpp#L126).

[sgrechanik]

mlir/lib/IR/AsmPrinter.cpp

1256

We can switch to generic form by default but that will somewhat worsen user experience by making error messages less pretty. I'm ok with it though.

Alternatively we can implement per-thread error suppression in DiagnosticEngine, but that would be a much bigger change.

[sgrechanik]

• Switched to generic printing in diagnostics, removed the thread_local trick
• Added a flag to skip verification
• Implemented more explicit testing of invalid op printing

[mehdi_amini]

Please wait for River to have another look!

Thanks :)

[rriddle]

Nice.

mlir/test/Analysis/test-match-reduction.mlir

10 ↗

(On Diff #411513)

If we wanted, we could allow remarks to use the proper custom format. We would need to add some documentation somewhere about what types of diagnostics are expected from the verifier though.

[sgrechanik]

mlir/test/Analysis/test-match-reduction.mlir

10 ↗

(On Diff #411513)

Makes sense, may be worth doing it in this review to avoid changing tests back and forth.

[sgrechanik]

• Switch to generic form only when severity level is Error

[sgrechanik]

• Standard -> Func

[sgrechanik]

@rriddle Please take another look, if you have any concerns I can revert to the last accepted state and extract the change into a separate patch.

[rriddle]

Yeah, looks good. (Didn't mean to block, when I approved that meant I was okay with submitting/fixing the remaining comments before submitting).

[sgrechanik]

Thanks, I'll merge it on Monday.

mlir/lib/IR/AsmPrinter.cpp

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SubElementInterfaces.h"
+#include "mlir/IR/Verifier.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SaveAndRestore.h"
+#include "llvm/Support/Threading.h"
 
 #include <tuple>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 void OperationName::print(raw_ostream &os) const { os << getStringRef(); }
 
private:
   OpPrintingFlags printerFlags;
 
   /// An optional location map to be populated.
   AsmState::LocationMap *locationMap;
 };
 } // 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) {
+  // Protection from infinite recursion: operation verification may want to
+  // print the operation itself, causing this function to be called recursively.
+  thread_local bool alwaysPrintGenericForm = false;
+  if (alwaysPrintGenericForm)
+    return printerFlags.printGenericOpForm();
+
+  if (!printerFlags.shouldPrintGenericOpForm()) {
+    // 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 &) {
+      return success(parentThreadId == llvm::get_threadid());
+    });
+    alwaysPrintGenericForm = true;
+    if (failed(verify(op)))
+      printerFlags.printGenericOpForm();
+    alwaysPrintGenericForm = false;
+  }
+
+  return printerFlags;
+}
+
 AsmState::AsmState(Operation *op, const OpPrintingFlags &printerFlags,
                    LocationMap *locationMap)
-    : impl(std::make_unique<AsmStateImpl>(op, printerFlags, locationMap)) {}
+    : impl(std::make_unique<AsmStateImpl>(
+          op, verifyOpAndAdjustFlags(op, printerFlags), locationMap)) {}
 AsmState::~AsmState() = default;
 
 const OpPrintingFlags &AsmState::getPrinterFlags() const {
   return impl->getPrinterFlags();
 }
 
 //===----------------------------------------------------------------------===//
 // AsmPrinter::Impl
void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
 
   AsmState state(op, printerFlags);
   print(os, state);
 }
 void Operation::print(raw_ostream &os, AsmState &state) {
   OperationPrinter printer(os, state.getImpl());
   if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope())
     printer.printTopLevelOperation(this);
   else
     printer.print(this);
 }
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());
   llvm::errs() << "\n";
 }
 
 void Block::print(raw_ostream &os) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   // Get the top-level op.
   while (auto *nextOp = parentOp->getParentOp())
     parentOp = nextOp;
 
   AsmState state(parentOp);
   print(os, state);
 }
 void Block::print(raw_ostream &os, AsmState &state) {
   OperationPrinter(os, state.getImpl()).print(this);
 }
 
 void Block::dump() { print(llvm::errs()); }
 
 /// Print out the name of the block without printing its body.
 void Block::printAsOperand(raw_ostream &os, bool printType) {
   Operation *parentOp = getParentOp();
   if (!parentOp) {
     os << "<<UNLINKED BLOCK>>\n";
     return;
   }
   AsmState state(parentOp);
   printAsOperand(os, state);
 }
 void Block::printAsOperand(raw_ostream &os, AsmState &state) {
   OperationPrinter printer(os, state.getImpl());
   printer.printBlockName(this);
 }

mlir/lib/IR/Operation.cpp

while (!pendingRegions.empty()) {
           // we don't assert here as this function is called from the verifier
           // and so could be called on invalid IR.
           if (!operand)
             return op.emitOpError("operation's operand is null");
 
           // Check that any value that is used by an operation is defined in the
           // same region as either an operation result.
           auto *operandRegion = operand.getParentRegion();
+          if (!operandRegion)
+            return op.emitError("operation's operand is unlinked");
           if (!region.isAncestor(operandRegion)) {
             return op.emitOpError("using value defined outside the region")
                        .attachNote(isolatedOp->getLoc())
                    << "required by region isolation constraints";
           }
         }
 
         // Schedule any regions in the operation for further checking.  Don't

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

-// RUN: mlir-opt -lower-affine %s | FileCheck %s
+// `-debug` is used here to check that we don't crash when it's specified.
+// RUN: mlir-opt -lower-affine -debug %s | FileCheck %s
 
 // CHECK-LABEL: func @empty() {
 func @empty() {
   return     // CHECK:  return
 }            // CHECK: }
 
 func private @body(index) -> ()
 
 // Simple loops are properly converted.

mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -convert-scf-to-cf %s | FileCheck %s
+// `-debug` is used here to check that we don't crash when it's specified.
+// RUN: mlir-opt -allow-unregistered-dialect -convert-scf-to-cf -debug %s | FileCheck %s
 
 // CHECK-LABEL: func @simple_std_for_loop(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 //  CHECK-NEXT:  cf.br ^bb1(%{{.*}} : index)
 //  CHECK-NEXT:  ^bb1(%{{.*}}: index):    // 2 preds: ^bb0, ^bb2
 //  CHECK-NEXT:    %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}} : index
 //  CHECK-NEXT:    cf.cond_br %{{.*}}, ^bb2, ^bb3
 //  CHECK-NEXT:  ^bb2:   // pred: ^bb1
 //  CHECK-NEXT:    %{{.*}} = arith.constant 1 : index