[mlir] Fix missing verification after running an OpToOpAdaptorPass

[rriddle]

Summary

The current decision of when to run the verifier is running on the
assumption that nested passes can't affect the validity of the parent
operation, which isn't true. Parent operations may attach any number
of constraints on nested operations, which may not necessarily be
captured (or shouldn't be captured) at a smaller granularity.

This commit rectifies this by properly running the verifier after an
OpToOpAdaptor pass. To avoid an explosive increase in compile time,
we only run verification on the parent operation itself. To do this, a
flag to mlir::verify is added to avoid recursive verification if it isn't
desired.

Fixes #54288

[mehdi_amini]

I'm a bit sad about the "added complexity" of this verifyRecursively, but I guess it makes sense here!

mlir/include/mlir/IR/Verifier.h

 #define MLIR_IR_VERIFIER_H
 
 namespace mlir {
 struct LogicalResult;
 class Operation;
 
 /// Perform (potentially expensive) checks of invariants, used to detect
 /// compiler bugs, on this operation and any nested operations. On error, this
-/// reports the error through the MLIRContext and returns failure.
-LogicalResult verify(Operation *op);
+/// reports the error through the MLIRContext and returns failure. If
+/// `verifyRecursively` is false, this assumes that nested operations have
+/// already been properly verified, and does not recursively invoke the verifier
+/// on nested operations.
+LogicalResult verify(Operation *op, bool verifyRecursively = true);
+
 } // namespace mlir
 
 #endif

mlir/lib/IR/Verifier.cpp

 #include <atomic>
 
 using namespace mlir;
 
 namespace {
 /// This class encapsulates all the state used to verify an operation region.
 class OperationVerifier {
 public:
+  /// If `verifyRecursively` is true, then this will also recursively verify
+  /// nested operations.
+  explicit OperationVerifier(bool verifyRecursively)
+      : verifyRecursively(verifyRecursively) {}
+
   /// Verify the given operation.
   LogicalResult verifyOpAndDominance(Operation &op);
 
 private:
   LogicalResult
   verifyBlock(Block &block,
               SmallVectorImpl<Operation *> &opsWithIsolatedRegions);
   /// Verify the properties and dominance relationships of this operation,
   /// stopping region recursion at any "isolated from above operations".  Any
   /// such ops are returned in the opsWithIsolatedRegions vector.
   LogicalResult
   verifyOperation(Operation &op,
                   SmallVectorImpl<Operation *> &opsWithIsolatedRegions);
 
   /// Verify the dominance property of regions contained within the given
   /// Operation.
   LogicalResult verifyDominanceOfContainedRegions(Operation &op,
                                                   DominanceInfo &domInfo);
+
+  /// A flag indicating if this verifier should recursively verify nested
+  /// operations.
+  bool verifyRecursively;
 };
 } // namespace
 
 LogicalResult OperationVerifier::verifyOpAndDominance(Operation &op) {
   SmallVector<Operation *> opsWithIsolatedRegions;
 
   // Verify the operation first, collecting any IsolatedFromAbove operations.
   if (failed(verifyOperation(op, opsWithIsolatedRegions)))
     return failure();
 
   // Since everything looks structurally ok to this point, we do a dominance
   // check for any nested regions. We do this as a second pass since malformed
   // CFG's can cause dominator analysis construction to crash and we want the
   // verifier to be resilient to malformed code.
   if (op.getNumRegions() != 0) {
     DominanceInfo domInfo;
     if (failed(verifyDominanceOfContainedRegions(op, domInfo)))
       return failure();
   }
 
-  // Check the dominance properties and invariants of any operations in the
-  // regions contained by the 'opsWithIsolatedRegions' operations.
+  // If we aren't verifying nested operations, then we're done.
+  if (!verifyRecursively)
+    return success();
+
+  // Otherwise, check the dominance properties and invariants of any operations
+  // in the regions contained by the 'opsWithIsolatedRegions' operations.
   return failableParallelForEach(
       op.getContext(), opsWithIsolatedRegions,
       [&](Operation *op) { return verifyOpAndDominance(*op); });
 }
 
 /// Returns true if this block may be valid without terminator. That is if:
 /// - it does not have a parent region.
 /// - Or the parent region have a single block and:
if (block.empty()) {
     if (mayBeValidWithoutTerminator(&block))
       return success();
     return emitError(block.getParent()->getLoc(),
                      "empty block: expect at least a terminator");
   }
 
   // Check each operation, and make sure there are no branches out of the
   // middle of this block.
-  for (auto &op : block) {
+  for (Operation &op : block) {
     // Only the last instructions is allowed to have successors.
     if (op.getNumSuccessors() != 0 && &op != &block.back())
       return op.emitError(
           "operation with block successors must terminate its parent block");
 
+    // If we aren't verifying recursievly, there is nothing left to check.
+    if (!verifyRecursively)
+      continue;
+
     // If this operation has regions and is IsolatedFromAbove, we defer
     // checking.  This allows us to parallelize verification better.
     if (op.getNumRegions() != 0 &&
         op.hasTrait<OpTrait::IsIsolatedFromAbove>()) {
       opsWithIsolatedRegions.push_back(&op);
-    } else {
+
       // Otherwise, check the operation inline.
-      if (failed(verifyOperation(op, opsWithIsolatedRegions)))
-        return failure();
+    } else if (failed(verifyOperation(op, opsWithIsolatedRegions))) {
     }
   }
 
   // Verify that this block is not branching to a block of a different
   // region.
   for (Block *successor : block.getSuccessors())
     if (successor->getParent() != block.getParent())
       return block.back().emitOpError(
LogicalResult OperationVerifier::verifyOperation(
   Optional<RegisteredOperationName> registeredInfo = opName.getRegisteredInfo();
   if (registeredInfo && failed(registeredInfo->verifyInvariants(&op)))
     return failure();
 
   if (unsigned numRegions = op.getNumRegions()) {
     auto kindInterface = dyn_cast<RegionKindInterface>(op);
 
     // Verify that all child regions are ok.
+    MutableArrayRef<Region> regions = op.getRegions();
     for (unsigned i = 0; i < numRegions; ++i) {
-      Region &region = op.getRegion(i);
+      Region &region = regions[i];
       RegionKind kind =
           kindInterface ? kindInterface.getRegionKind(i) : RegionKind::SSACFG;
       // Check that Graph Regions only have a single basic block. This is
       // similar to the code in SingleBlockImplicitTerminator, but doesn't
       // require the trait to be specified. This arbitrary limitation is
       // designed to limit the number of cases that have to be handled by
       // transforms and conversions.
       if (op.isRegistered() && kind == RegionKind::Graph) {
         // Non-empty regions must contain a single basic block.
         if (!region.empty() && !region.hasOneBlock())
           return op.emitOpError("expects graph region #")
                  << i << " to have 0 or 1 blocks";
       }
 
       if (region.empty())
         continue;
 
       // Verify the first block has no predecessors.
       Block *firstBB = &region.front();
       if (!firstBB->hasNoPredecessors())
         return emitError(op.getLoc(),
                          "entry block of region may not have predecessors");
 
-      // Verify each of the blocks within the region.
+      // Verify each of the blocks within the region if we are verifying
+      // recursively.
+      if (verifyRecursively) {
         for (Block &block : region)
           if (failed(verifyBlock(block, opsWithIsolatedRegions)))
             return failure();
       }
     }
+  }
 
   // After the region ops are verified, run the verifiers that have additional
   // region invariants need to veirfy.
   if (registeredInfo && failed(registeredInfo->verifyRegionInvariants(&op)))
     return failure();
 
   // If this is a registered operation, there is nothing left to do.
   if (registeredInfo)
for (Block &block : region) {
             if (domInfo.properlyDominates(operand.value(), &op))
               continue;
 
             diagnoseInvalidOperandDominance(op, operand.index());
             return failure();
           }
         }
 
-        // Recursively verify dominance within each operation in the
-        // block, even if the block itself is not reachable, or we are in
-        // a region which doesn't respect dominance.
-        if (op.getNumRegions() != 0) {
+        // Recursively verify dominance within each operation in the block, even
+        // if the block itself is not reachable, or we are in a region which
+        // doesn't respect dominance.
+        if (verifyRecursively && op.getNumRegions() != 0) {
           // If this operation is IsolatedFromAbove, then we'll handle it in the
           // outer verification loop.
           if (op.hasTrait<OpTrait::IsIsolatedFromAbove>())
             continue;
 
           if (failed(verifyDominanceOfContainedRegions(op, domInfo)))
             return failure();
         }
       }
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Entrypoint
 //===----------------------------------------------------------------------===//
 
-/// Perform (potentially expensive) checks of invariants, used to detect
-/// compiler bugs.  On error, this reports the error through the MLIRContext and
-/// returns failure.
-LogicalResult mlir::verify(Operation *op) {
-  return OperationVerifier().verifyOpAndDominance(*op);
+LogicalResult mlir::verify(Operation *op, bool verifyRecursively) {
+  OperationVerifier verifier(verifyRecursively);
+  return verifier.verifyOpAndDominance(*op);
 }

mlir/lib/Pass/Pass.cpp

LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // When verifyPasses is specified, we run the verifier (unless the pass
   // failed).
   if (!passFailed && verifyPasses) {
     bool runVerifierNow = true;
+
+    // If the pass is an adaptor pass, we don't run the verifier recursively
+    // because the nested operations should have already been verified after
+    // nested passes had run.
+    bool runVerifierRecursively = !isa<OpToOpPassAdaptor>(pass);
+
     // Reduce compile time by avoiding running the verifier if the pass didn't
     // change the IR since the last time the verifier was run:
     //
     //  1) If the pass said that it preserved all analyses then it can't have
     //     permuted the IR.
-    //  2) If we just ran an OpToOpPassAdaptor (e.g. to run function passes
-    //     within a module) then each sub-unit will have been verified on the
-    //     subunit (and those passes aren't allowed to modify the parent).
     //
     // We run these checks in EXPENSIVE_CHECKS mode out of caution.
 #ifndef EXPENSIVE_CHECKS
-    runVerifierNow = !isa<OpToOpPassAdaptor>(pass) &&
-                     !pass->passState->preservedAnalyses.isAll();
+    runVerifierNow = !pass->passState->preservedAnalyses.isAll();
 #endif
     if (runVerifierNow)
-      passFailed = failed(verify(op));
+      passFailed = failed(verify(op, runVerifierRecursively));
   }
 
   // Instrument after the pass has run.
   if (pi) {
     if (passFailed)
       pi->runAfterPassFailed(pass, op);
     else
       pi->runAfterPass(pass, op);

mlir/test/Pass/invalid-parent.mlir

+// RUN: mlir-opt %s -pass-pipeline='builtin.func(test-pass-invalid-parent)' -verify-diagnostics
+
+// Test that we properly report errors when the parent becomes invalid after running a pass
+// on a child operation.
+// expected-error@below {{'some_unknown_func' does not reference a valid function}}
+func @TestCreateInvalidCallInPass() {
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 }
 
 void TestDialect::getCanonicalizationPatterns(
     RewritePatternSet &results) const {
   results.add(&dialectCanonicalizationPattern);
 }
 
 //===----------------------------------------------------------------------===//
+// TestCallOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult TestCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // Check that the callee attribute was specified.
+  auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
+  if (!fnAttr)
+    return emitOpError("requires a 'callee' symbol reference attribute");
+  if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))
+    return emitOpError() << "'" << fnAttr.getValue()
+                         << "' does not reference a valid function";
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TestFoldToCallOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
   using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(FoldToCallOp op,

mlir/test/lib/Dialect/Test/TestOps.td

def VariadicNoTerminatorOp : TEST_Op<"variadic_no_terminator_op",
 
   let assemblyFormat = "attr-dict `:` $my_regions";
 }
 
 //===----------------------------------------------------------------------===//
 // Test Call Interfaces
 //===----------------------------------------------------------------------===//
 
+def TestCallOp : TEST_Op<"call", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands);
+  let results = (outs Variadic<AnyType>);
+  let assemblyFormat = [{
+    $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
+  }];
+}
+
 def ConversionCallOp : TEST_Op<"conversion_call_op",
     [CallOpInterface]> {
   let arguments = (ins Variadic<AnyType>:$arg_operands, SymbolRefAttr:$callee);
   let results = (outs Variadic<AnyType>);
 
   let extraClassDeclaration = [{
     /// Return the callee of this operation.
     ::mlir::CallInterfaceCallable getCallableForCallee();

mlir/test/lib/Pass/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestPass
   TestDynamicPipeline.cpp
   TestPassManager.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Pass
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
+  MLIRTestDialect
+  )
+
+target_include_directories(MLIRTestPass
+  PRIVATE
+  ${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test
+  ${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test
   )

mlir/test/lib/Pass/TestPassManager.cpp

 //===- TestPassManager.cpp - Test pass manager functionality --------------===//
 //
 // 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 "TestDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 
 using namespace mlir;
 
 namespace {
 struct TestModulePass

 class TestFailurePass : public PassWrapper<TestFailurePass, OperationPass<>> {
   void runOnOperation() final { signalPassFailure(); }
   StringRef getArgument() const final { return "test-pass-failure"; }
   StringRef getDescription() const final {
     return "Test a pass in the pass manager that always fails";
   }
 };
 
+/// A test pass that always fails to enable testing the failure recovery
+/// mechanisms of the pass manager.
+class TestInvalidParentPass
+    : public PassWrapper<TestInvalidParentPass,
+                         InterfacePass<FunctionOpInterface>> {
+  StringRef getArgument() const final { return "test-pass-invalid-parent"; }
+  StringRef getDescription() const final {
+    return "Test a pass in the pass manager that makes the parent operation "
+           "invalid";
+  }
+  void getDependentDialects(DialectRegistry &registry) const final {
+    registry.insert<test::TestDialect>();
+  }
+  void runOnOperation() final {
+    FunctionOpInterface op = getOperation();
+    OpBuilder b(getOperation().getBody());
+    b.create<test::TestCallOp>(op.getLoc(), TypeRange(), "some_unknown_func",
+                               ValueRange());
+  }
+};
+
 /// A test pass that contains a statistic.
 struct TestStatisticPass
     : public PassWrapper<TestStatisticPass, OperationPass<>> {
   TestStatisticPass() = default;
   TestStatisticPass(const TestStatisticPass &) {}
   StringRef getArgument() const final { return "test-stats-pass"; }
   StringRef getDescription() const final { return "Test pass statistics"; }
 
void registerPassManagerTestPass() {
   PassRegistration<TestModulePass>();
 
   PassRegistration<TestFunctionPass>();
 
   PassRegistration<TestInterfacePass>();
 
   PassRegistration<TestCrashRecoveryPass>();
   PassRegistration<TestFailurePass>();
+  PassRegistration<TestInvalidParentPass>();
 
   PassRegistration<TestStatisticPass>();
 
   PassPipelineRegistration<>("test-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",
                              testNestedPipeline);
   PassPipelineRegistration<>("test-textual-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",