[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

343–345

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

343–345

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

374–381

433

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

433

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

[rriddle]

mlir/lib/IR/Verifier.cpp

439

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

439

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

414

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

 extern template class llvm::DominatorTreeBase<mlir::Block, false>;
 extern template class llvm::DominatorTreeBase<mlir::Block, true>;
 
 namespace mlir {
 using DominanceInfoNode = llvm::DomTreeNodeBase<Block>;
 class Operation;
 
 namespace detail {
-template <bool IsPostDom> class DominanceInfoBase {
+template <bool IsPostDom>
+class DominanceInfoBase {
   using base = llvm::DominatorTreeBase<Block, IsPostDom>;
 
 public:
-  DominanceInfoBase(Operation *op) { recalculate(op); }
+  /// Build dominance info for the regions contained in the specified op.  If
+  /// `ignoreIsolatedRegions` is true, then this doesn't build dominance
+  /// information for the bodies of nested `IsolatedFromAbove` operations.
+  DominanceInfoBase(Operation *op, bool ignoreIsolatedRegions = false) {
+    recalculate(op, ignoreIsolatedRegions);
+  }
   DominanceInfoBase(DominanceInfoBase &&) = default;
   DominanceInfoBase &operator=(DominanceInfoBase &&) = default;
 
   DominanceInfoBase(const DominanceInfoBase &) = delete;
   DominanceInfoBase &operator=(const DominanceInfoBase &) = delete;
 
-  /// Recalculate the dominance info.
-  void recalculate(Operation *op);
+  /// Recalculate the dominance info.  If `ignoreIsolatedRegions` is true, then
+  /// this doesn't build dominance information for the bodies of nested
+  /// `IsolatedFromAbove` operations.
+  void recalculate(Operation *op, bool ignoreIsolatedRegions = false);
 
   /// Finds the nearest common dominator block for the two given blocks a
   /// and b. If no common dominator can be found, this function will return
   /// nullptr.
   Block *findNearestCommonDominator(Block *a, Block *b) const;
 
   /// Return true if there is dominanceInfo for the given region.
   bool hasDominanceInfo(Region *region) {

 };
 
 } //  end namespace mlir
 
 namespace llvm {
 
 /// DominatorTree GraphTraits specialization so the DominatorTree can be
 /// iterated by generic graph iterators.
-template <> struct GraphTraits<mlir::DominanceInfoNode *> {
+template <>
+struct GraphTraits<mlir::DominanceInfoNode *> {
   using ChildIteratorType = mlir::DominanceInfoNode::const_iterator;
   using NodeRef = mlir::DominanceInfoNode *;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
   static inline ChildIteratorType child_begin(NodeRef N) { return N->begin(); }
   static inline ChildIteratorType child_end(NodeRef N) { return N->end(); }
 };
 
-template <> struct GraphTraits<const mlir::DominanceInfoNode *> {
+template <>
+struct GraphTraits<const mlir::DominanceInfoNode *> {
   using ChildIteratorType = mlir::DominanceInfoNode::const_iterator;
   using NodeRef = const mlir::DominanceInfoNode *;
 
   static NodeRef getEntryNode(NodeRef N) { return N; }
   static inline ChildIteratorType child_begin(NodeRef N) { return N->begin(); }
   static inline ChildIteratorType child_end(NodeRef N) { return N->end(); }
 };
 
 } // end namespace llvm
 #endif

mlir/lib/IR/Dominance.cpp

static bool hasSSADominance(Operation *op, unsigned index) {
   return !kindInterface || kindInterface.hasSSADominance(index);
 }
 
 //===----------------------------------------------------------------------===//
 // DominanceInfoBase
 //===----------------------------------------------------------------------===//
 
 template <bool IsPostDom>
-void DominanceInfoBase<IsPostDom>::recalculate(Operation *op) {
+void DominanceInfoBase<IsPostDom>::recalculate(Operation *topOperation,
+                                               bool ignoreIsolatedRegions) {
   dominanceInfos.clear();
 
   // Build the dominance for each of the operation regions.
-  op->walk([&](Operation *op) {
+  topOperation->walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
     unsigned numRegions = op->getNumRegions();
     if (numRegions == 0)
-      return;
+      return WalkResult::advance();
+
+    // If we are supposed to ignore IsolatedFromAbove regions and this operation
+    // is one of them, then don't recurse into it.
+    if (ignoreIsolatedRegions && op->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+        op != topOperation)
+      return WalkResult::skip();
 
     auto kindInterface = dyn_cast<RegionKindInterface>(op);
     for (unsigned i = 0; i < numRegions; i++) {
       Region &region = op->getRegion(i);
       // Don't compute dominance if the region is empty.
       if (region.empty())
         continue;
 
       // Dominance changes based on the region type. Avoid the helper
       // function here so we don't do the region cast repeatedly.
       bool hasSSADominance =
           op->isRegistered() &&
           (!kindInterface || kindInterface.hasSSADominance(i));
       // If a region has SSADominance, then compute detailed dominance
       // info.  Otherwise, all values in the region are live anywhere
       // in the region, which is represented as an empty entry in the
       // dominanceInfos map.
       if (hasSSADominance) {
         auto opDominance = std::make_unique<base>();
         opDominance->recalculate(region);
         dominanceInfos.try_emplace(&region, std::move(opDominance));
       }
     }
+    return WalkResult::advance();
   });
 }
 
 /// Walks up the list of containers of the given block and calls the
 /// user-defined traversal function for every pair of a region and block that
 /// could be found during traversal. If the user-defined function returns true
 /// for a given pair, traverseAncestors will return the current block. Nullptr
 /// otherwise.

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 supports deferred computation of dominance information from the
+/// specified operation.
+///
+/// Dominance info is expensive to create and completely unnecessary for top
+/// levels of the IR structure like modules.  Deferring this allows us to
+/// avoid work in many cases as well as parallelizing the work when it is
+/// required.
+class LazyDomInfo {
+public:
+  explicit LazyDomInfo(Operation *topOperation) : topOperation(topOperation) {}
+
+  /// Get the dominance information for this operation, lazily constructing it
+  /// if necessary.
+  DominanceInfo &get() {
+    if (!domInfo.hasValue())
+      domInfo.emplace(topOperation, /*ignoreIsolatedRegions*/ true);
+    return domInfo.getValue();
+  }
+
+private:
+  void operator=(const LazyDomInfo &) = delete;
+  LazyDomInfo(const LazyDomInfo &) = delete;
+
+  Operation *topOperation;
+  Optional<DominanceInfo> domInfo;
+};
+} // end anonymous namespace
+
+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,
+                                                  LazyDomInfo &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) {
+    LazyDomInfo domInfo(&op);
+    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,
+                                                     LazyDomInfo &domInfo) {
+  // This vector keeps track of ops that have regions which should be checked
+  // in parallel.
+  SmallVector<Operation *> opsWithRegionsToCheckInParallel;
+
+  // This keeps track of operations seen in this block, which avoids dense
+  // dominance checks.
+  SmallPtrSet<Operation *, 8> opsInThisBlock;
+
+  // Get information about the requirements on the regions in this op.
+  auto kindInterface = dyn_cast<RegionKindInterface>(opWithRegions);
+  for (size_t regionNo = 0, e = opWithRegions.getNumRegions(); regionNo != e;
+       ++regionNo) {
+    Region &region = opWithRegions.getRegion(regionNo);
+    bool hasSSADominance =
+        opWithRegions.isRegistered() &&
+        (!kindInterface || kindInterface.hasSSADominance(regionNo));
+
+    bool isFirstBlockInRegion = true;
     for (Block &block : region) {
-      // Dominance is only meaningful inside reachable blocks.
-      bool isReachable = domInfo->isReachableFromEntry(&block);
+      // Dominance is only meaningful inside reachable blocks.  We know that
+      // the first block of every region is reachable.
+      bool isReachable =
+          isFirstBlockInRegion || domInfo.get().isReachableFromEntry(&block);
+
+      // This predicate checks to see if the specified operand is known to
+      // dominate its user.  This is equivalent to but more efficient than
+      // building and querying dominance because we're doing dense queries and
+      // it is better equipped to do sparse lookups.
+      auto doesOperandProperlyDominateUser = [&](Value operand,
+                                                 Operation *user) {
+        if (auto *operandOp = operand.getDefiningOp()) {
+          // If the operand is an operation in the current block, then we can
+          // efficiently know its dominance properties without checking
+          // dominators.
+          if (operandOp->getBlock() == &block)
+            return !hasSSADominance || opsInThisBlock.count(operandOp);
+        } else {
+          // Arguments for the current block dominate all operations in the
+          // block.
+          auto *argumentBlock = operand.cast<BlockArgument>().getOwner();
+          if (argumentBlock == &block)
+            return true;
+        }
+
+        // If we don't otherwise know, fall back to checking with dominance.
+        return domInfo.get().properlyDominates(operand, user);
+      };
+
+      // Check each operation in this block, and any operations in regions
+      // that these operations contain.
+      opsWithRegionsToCheckInParallel.clear();
+      opsInThisBlock.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 (!doesOperandProperlyDominateUser(operand.get(), &op)) {
+              diagnoseInvalidOperandDominance(op, operand.getOperandNumber());
+              return failure();
+            }
+          }
+        }
+
+        // If this region respects SSA dominance, then keep track of this
+        // instruction to accelerate later checks in this block.
+        if (hasSSADominance)
+          opsInThisBlock.insert(&op);
+
+        // 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.
+          LazyDomInfo localDomInfo(&op);
+          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're doing sub-region checking in parallel, do it now.
+      if (!opsWithRegionsToCheckInParallel.empty()) {
+        std::atomic<bool> passFailed(false);
+        llvm::parallelForEach(
+            opsWithRegionsToCheckInParallel, [&](Operation *op) {
+              // Each isolated operation gets its own dom info.
+              LazyDomInfo localDomInfo(op);
+              if (failed(verifyDominanceOfContainedRegions(*op, localDomInfo)))
+                passFailed = true;
+            });
+        if (passFailed)
       }
+
+      isFirstBlockInRegion = false;
     }
   }
   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);
 }