[Verifier] Speed up and parallelize dominance checking. NFC

[lattner]

Summary

One of the key algorithms used in the "mlir::verify(op)" method is the
dominance checker, which ensure that operand values properly dominate
the operations that use them. This is currently slower than it needs
to be, so let's parallelize it:

1. It splits dominance checking into "IsolatedFromAbove" units. Instead of building a monolithic DominanceInfo for everything in a module, for example, it checks dominance for the module to all the functions within it (noop, since there are no operands at this level) then each function gets their own DominanceInfo for each of their scope.
2. It parallelizes analysis of collections IsolatedFromAbove operations, e.g. each of the functions within a Module.

All together this is more than a 10% speedup on firtool in circt on a
large design when run in -verify-each mode (our default) since the verifier
is invoked after each pass.

Still todo is to parallelize the main verifier pass. I decided to split
this out to its own thing since this patch is already large-ish.

[lattner]

Improve comments.

[rriddle]

mlir/lib/IR/Verifier.cpp

312–314

Can you just use Operation::isBeforeInBlock? I don't think there needs to be a look up table here.

[lattner]

mlir/lib/IR/Verifier.cpp

312–314

yes, I should have remembered that. That provides an additional speedup, and allows generalizing this a bit more. I'll incorporate that in my next revision, thanks!

[lattner]

Use isBeforeInBlock instead of maintaining a SmallPtrSet.

[lattner]

Actually, instead of making this more complicated, I'd like to keep this simple in this patch as it is.

In a different work thread, I feel compelled to completely replace the Dominator implementation: there is no good reason for it to be as slow as it is, and it is being horribly misused by some clients (e.g. Linalg/.../Fusion.cpp). Making it lazy and not making an llvm::domtree for single block regions will solve this problem and define away problems for other clients.

That said, I'd like to take this step and do that in parallel, then re-simplify the verifier if it works out.

[lattner]

The dominators rewrite is here: https://reviews.llvm.org/D103384. Once that goes into tree I'll merge it into this patch to simplify it - we can drop the "don't recurse into IsolatedFromAbove regions" change clearly and the manual laziness isn't needed, but the parallel constructs are.

[rriddle]

LGTM, especially if the DominanceInfo stuff is going to be cleaned up in a followup. Thanks for cleaning this up!

mlir/lib/IR/Verifier.cpp

332–339

383

The verifier isn't guaranteed to be called in a situation that can handle diagnostics being emitted in parallel. I think we need to use ParallelDiagnosticHandler here.

(Honestly we should just add some MLIR specific threading utilities that wrap all of this automatically)

[lattner]

SGTM, I'll revise this after the other patch lands, thank you for the review!

[lattner]

Updates from River's review

[lattner]

mlir/lib/IR/Verifier.cpp

383

Great catch. I'm not super familiar with ParallelDiagnosticHandler. Does this look like the right thing?

[rriddle]

mlir/lib/IR/Verifier.cpp

389

When using ParallelDiagnosticHandler you need to set an order for each of the child threads that get created (this is what the handler uses to make sure diagnostics get ordered properly):

https://mlir.llvm.org/docs/Diagnostics/#parallel-diagnostic-handler
https://github.com/llvm/llvm-project/blob/e9f82e9e1b6851ff66010a0e500346379b1ef572/mlir/include/mlir/IR/Diagnostics.h#L628

[lattner]

Use setOrderIDForThread correctly.

[lattner]

mlir/lib/IR/Verifier.cpp

389

Got it, thanks. This is actually a funny situation where we don't care what order the diagnostics come out in. We only care about "no errors or some errors". That said, I incorporated this since I expect there is little cost to doing so.

I'm going to land the other patch and then merge it into this one, I'll ping this when I have done so. Thanks!

[lattner]

Dramatically simplify this patch by merging in the dominator improvements.

[lattner]

Final grooming, remove commented out #include, use range base for loop again.

[lattner]

Ok, I think this patch is ready for a re-check. Thanks!

[lattner]

Update commit message

[lattner]

update commit message for real this time.

[lattner]

I'm abandoning this patch. With the other general improvements to dominators and with the inefficiency in the threading stuff, this is no longer a win.

[rriddle]

I thought you abandoned this change? Looks like you landed it.

[lattner]

I'm not sure how this landed, it looks like it got pushed as a consequence of the twine patch. Maybe I put them both into the same directory accidentally (my git fu is not great).

I'll clean this up when I have a chance soon, I do have a plan to parallelize the verifier, I will merge the dom checking into that flow and make the whole thing fast. I'll take care of it then. Until then I don't think this patch causes any significant harm.

I'm sorry for flubbing this though!!! If you'd like to revert this, please feel free.

[mehdi_amini]

Run ln -sf ../../llvm/utils/git/pre-push.py .git/hooks/pre-push in your local repo: it'll protect you against accidentally pushing two commits by asking confirmation when this happens

[GMNGeoffrey]

We are seeing a mysterious hang in IREE's python tests after this change that happens about 90% of the time (making the test flakey). Entirely possible that it's something broken in our python bindings, especially given that I haven't yet seen it in anything that doesn't go through the bindings. But given that it was apparently landed accidentally, perhaps you'd consider a revert :-)

mlir/lib/IR/Verifier.cpp

344

Changing this to else if(false) gets rid of the hanging

[GMNGeoffrey]

(notably the python tests are just shelling out to IREE binaries, so "bindings" is a strong word here 😁)

[GMNGeoffrey]

Oh and I also get a build failure because of a missing include of atomic:

/usr/local/google/home/gcmn/src/iree/third_party/llvm-project/mlir/lib/IR/Verifier.cpp:365:27: error: implicit instantiation of undefined template 'std::atomic<bool>'
        std::atomic<bool> passFailed(false);
                          ^
/usr/bin/../lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/bits/atomic_base.h:152:12: note: template is declared here
    struct atomic;

which is also failing on the premerge checks.

[GMNGeoffrey]

Update: we suspect this is something in our python logic for piping subprocesses together, so likely this change is just a bystander :-)

[lattner]

I don't know how to revert this, but feel free if you see it as a smoking gun. I only abandoned this because "it didn't show a big speedup and I have a hopefully better idea of how to improve things", not because I think this is harmful

[lattner]

The new patch is available here: https://reviews.llvm.org/D103373/new/

The difference is that it parallelizes operation verification *and* dominance checks, pulling them into a common framework. This helps the granularity issues that made this patch be not a significant win.

[mehdi_amini]

I suspect you meant https://reviews.llvm.org/D104207 ?

[lattner]

Yes, thank you! Too many similar things.

mlir/include/mlir/IR/Dominance.h

 class Operation;
 
 namespace detail {
 template <bool IsPostDom>
 class DominanceInfoBase {
   using DomTree = llvm::DominatorTreeBase<Block, IsPostDom>;
 
 public:
-  DominanceInfoBase(Operation *op) {}
+  DominanceInfoBase(Operation *op = nullptr) {}
   DominanceInfoBase(DominanceInfoBase &&) = default;
   DominanceInfoBase &operator=(DominanceInfoBase &&) = default;
   ~DominanceInfoBase();
 
   DominanceInfoBase(const DominanceInfoBase &) = delete;
   DominanceInfoBase &operator=(const DominanceInfoBase &) = delete;
 
   /// Invalidate dominance info. This can be used by clients that make major

mlir/lib/IR/Verifier.cpp

 
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "llvm/ADT/StringMap.h"
-#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/Parallel.h"
 #include "llvm/Support/PrettyStackTrace.h"
-#include "llvm/Support/Regex.h"
 
 using namespace mlir;
 
 namespace {
 /// This class encapsulates all the state used to verify an operation region.
 class OperationVerifier {
 public:
-  explicit OperationVerifier() {}
+  explicit OperationVerifier(MLIRContext *context)
+      : parallelismEnabled(context->isMultithreadingEnabled()) {}
 
   /// Verify the given operation.
   LogicalResult verifyOpAndDominance(Operation &op);
 
 private:
   /// Verify the given potentially nested region or block.
   LogicalResult verifyRegion(Region &region);
   LogicalResult verifyBlock(Block &block);
   LogicalResult verifyOperation(Operation &op);
 
   /// Verify the dominance property of regions contained within the given
   /// Operation.
-  LogicalResult verifyDominanceOfContainedRegions(Operation &op);
+  LogicalResult verifyDominanceOfContainedRegions(Operation &op,
+                                                  DominanceInfo &domInfo);
 
   /// Emit an error for the given block.
   InFlightDiagnostic emitError(Block &bb, const Twine &message) {
     // Take the location information for the first operation in the block.
     if (!bb.empty())
       return bb.front().emitError(message);
 
     // Worst case, fall back to using the parent's location.
     return mlir::emitError(bb.getParent()->getLoc(), message);
   }
 
-  /// Dominance information for this operation, when checking dominance.
-  DominanceInfo *domInfo = nullptr;
+  /// This is true if parallelism is enabled on the MLIRContext.
+  const bool parallelismEnabled;
 };
 } // end anonymous namespace
 
 /// Verify the given operation.
 LogicalResult OperationVerifier::verifyOpAndDominance(Operation &op) {
   // Verify the operation first.
   if (failed(verifyOperation(op)))
     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 constructure to crash and we want the
   // verifier to be resilient to malformed code.
-  DominanceInfo theDomInfo(&op);
-  domInfo = &theDomInfo;
-  if (failed(verifyDominanceOfContainedRegions(op)))
+  if (op.getNumRegions() != 0) {
+    DominanceInfo domInfo;
+    if (failed(verifyDominanceOfContainedRegions(op, /*domInfo*/ domInfo)))
       return failure();
-
-  domInfo = nullptr;
+  }
   return success();
 }
 
 LogicalResult OperationVerifier::verifyRegion(Region &region) {
   if (region.empty())
     return success();
 
   // Verify the first block has no predecessors.
static void diagnoseInvalidOperandDominance(Operation &op, unsigned operandNo) {
   else if (region2->isProperAncestor(region1))
     note << " in a parent region)";
   else if (region1->isProperAncestor(region2))
     note << " in a child region)";
   else
     note << " neither in a parent nor in a child region)";
 }
 
-/// Verify the dominance of each of the nested blocks within the given operation
+/// Verify the dominance of each of the nested blocks within the given
+/// operation.  domInfo may be present or absent (null), depending on whether
+/// the caller computed it for a higher level.
 LogicalResult
-OperationVerifier::verifyDominanceOfContainedRegions(Operation &op) {
-  for (Region &region : op.getRegions()) {
-    // Verify the dominance of each of the held operations.
+OperationVerifier::verifyDominanceOfContainedRegions(Operation &opWithRegions,
+                                                     DominanceInfo &domInfo) {
+  // This vector keeps track of ops that have regions which should be checked
+  // in parallel.
+  SmallVector<Operation *> opsWithRegionsToCheckInParallel;
+
+  // Get information about the requirements on the regions in this op.
+  for (Region &region : opWithRegions.getRegions()) {
     for (Block &block : region) {
       // Dominance is only meaningful inside reachable blocks.
-      bool isReachable = domInfo->isReachableFromEntry(&block);
+      bool isReachable = domInfo.isReachableFromEntry(&block);
+
+      // Check each operation in this block, and any operations in regions
+      // that these operations contain.
+      opsWithRegionsToCheckInParallel.clear();
 
       for (Operation &op : block) {
         if (isReachable) {
           // Check that operands properly dominate this use.
-          for (unsigned operandNo = 0, e = op.getNumOperands(); operandNo != e;
-               ++operandNo) {
-            if (domInfo->properlyDominates(op.getOperand(operandNo), &op))
-              continue;
+          for (auto &operand : op.getOpOperands()) {
+            // If the operand doesn't dominate the user, then emit an error.
+            if (!domInfo.properlyDominates(operand.get(), &op)) {
+              diagnoseInvalidOperandDominance(op, operand.getOperandNumber());
+              return failure();
+            }
+          }
+        }
+
+        // If this operation has any regions, we need to recursively verify
+        // dominance of the ops within it.
+        if (op.getNumRegions() == 0)
 
-            diagnoseInvalidOperandDominance(op, operandNo);
+        // If this is a non-isolated region (e.g. an affine for loop), pass down
+        // the current dominator information.
+        if (!op.hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+          if (failed(verifyDominanceOfContainedRegions(op, domInfo)))
+            return failure();
+        } else if (parallelismEnabled) {
+          // If this is an IsolatedFromAbove op and parallelism is enabled, then
+          // enqueue this for processing later.
+          opsWithRegionsToCheckInParallel.push_back(&op);
+        } else {
+          // If not, just verify inline with a local dom scope.
+          DominanceInfo localDomInfo;
+          if (failed(verifyDominanceOfContainedRegions(op, localDomInfo)))
             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)
-          if (failed(verifyDominanceOfContainedRegions(op)))
-            return failure();
+      // If we have multiple parallelizable subregions, check them in parallel.
+      if (opsWithRegionsToCheckInParallel.size() == 1) {
+        // Each isolated operation gets its own dom info.
+        Operation *op = opsWithRegionsToCheckInParallel[0];
+        DominanceInfo localDomInfo;
+        if (failed(verifyDominanceOfContainedRegions(*op, localDomInfo)))
+          return failure();
+      } else if (!opsWithRegionsToCheckInParallel.empty()) {
+        ParallelDiagnosticHandler handler(opWithRegions.getContext());
+        std::atomic<bool> passFailed(false);
+        llvm::parallelForEachN(
+            0, opsWithRegionsToCheckInParallel.size(), [&](size_t opIdx) {
+              handler.setOrderIDForThread(opIdx);
+              Operation *op = opsWithRegionsToCheckInParallel[opIdx];
+
+              // Each isolated operation gets its own dom info.
+              DominanceInfo localDomInfo;
+              if (failed(verifyDominanceOfContainedRegions(*op, localDomInfo)))
+                passFailed = true;
+              handler.eraseOrderIDForThread();
+            });
+        if (passFailed)
       }
     }
   }
   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.
+/// compiler bugs.  On error, this reports the error through the MLIRContext
+/// and returns failure.
 LogicalResult mlir::verify(Operation *op) {
-  return OperationVerifier().verifyOpAndDominance(*op);
+  return OperationVerifier(op->getContext()).verifyOpAndDominance(*op);
 }