diff --git a/mlir/include/mlir/IR/Dominance.h b/mlir/include/mlir/IR/Dominance.h --- a/mlir/include/mlir/IR/Dominance.h +++ b/mlir/include/mlir/IR/Dominance.h @@ -20,19 +20,27 @@ class Operation; namespace detail { -template class DominanceInfoBase { +template +class DominanceInfoBase { using base = llvm::DominatorTreeBase; 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 @@ -164,7 +172,8 @@ /// DominatorTree GraphTraits specialization so the DominatorTree can be /// iterated by generic graph iterators. -template <> struct GraphTraits { +template <> +struct GraphTraits { using ChildIteratorType = mlir::DominanceInfoNode::const_iterator; using NodeRef = mlir::DominanceInfoNode *; @@ -173,7 +182,8 @@ static inline ChildIteratorType child_end(NodeRef N) { return N->end(); } }; -template <> struct GraphTraits { +template <> +struct GraphTraits { using ChildIteratorType = mlir::DominanceInfoNode::const_iterator; using NodeRef = const mlir::DominanceInfoNode *; diff --git a/mlir/lib/IR/Dominance.cpp b/mlir/lib/IR/Dominance.cpp --- a/mlir/lib/IR/Dominance.cpp +++ b/mlir/lib/IR/Dominance.cpp @@ -39,14 +39,21 @@ //===----------------------------------------------------------------------===// template -void DominanceInfoBase::recalculate(Operation *op) { +void DominanceInfoBase::recalculate(Operation *topOperation, + bool ignoreIsolatedRegions) { dominanceInfos.clear(); // Build the dominance for each of the operation regions. - op->walk([&](Operation *op) { + topOperation->walk([&](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() && + op != topOperation) + return WalkResult::skip(); auto kindInterface = dyn_cast(op); for (unsigned i = 0; i < numRegions; i++) { @@ -70,6 +77,7 @@ dominanceInfos.try_emplace(®ion, std::move(opDominance)); } } + return WalkResult::advance(); }); } diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp --- a/mlir/lib/IR/Verifier.cpp +++ b/mlir/lib/IR/Verifier.cpp @@ -32,16 +32,47 @@ #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 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); @@ -54,7 +85,8 @@ /// 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) { @@ -66,8 +98,8 @@ 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 @@ -81,12 +113,11 @@ // 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))) - return failure(); - - domInfo = nullptr; + if (op.getNumRegions() != 0) { + LazyDomInfo domInfo(&op); + if (failed(verifyDominanceOfContainedRegions(op, /*domInfo*/ domInfo))) + return failure(); + } return success(); } @@ -299,35 +330,117 @@ 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 ®ion : 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 opsWithRegionsToCheckInParallel; + + // Get information about the requirements on the regions in this op. + auto kindInterface = dyn_cast(opWithRegions); + for (size_t regionNo = 0, e = opWithRegions.getNumRegions(); regionNo != e; + ++regionNo) { + Region ®ion = 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 || operandOp->isBeforeInBlock(user); + } else if (&block == operand.cast().getOwner()) { + // Arguments for the current block dominate all operations in the + // 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(); 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; - - diagnoseInvalidOperandDominance(op, operandNo); - return failure(); + 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(); + } } } - // 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))) + // If this operation has any regions, we need to recursively verify + // dominance of the ops within it. + if (op.getNumRegions() == 0) + continue; + + // If this is a non-isolated region (e.g. an affine for loop), pass down + // the current dominator information. + if (!op.hasTrait()) { + 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(); + } } + + // 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]; + LazyDomInfo localDomInfo(op); + if (failed(verifyDominanceOfContainedRegions(*op, localDomInfo))) + return failure(); + } else if (!opsWithRegionsToCheckInParallel.empty()) { + ParallelDiagnosticHandler handler(opWithRegions.getContext()); + std::atomic 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. + LazyDomInfo localDomInfo(op); + if (failed(verifyDominanceOfContainedRegions(*op, localDomInfo))) + passFailed = true; + handler.eraseOrderIDForThread(); + }); + if (passFailed) + return failure(); + } + + isFirstBlockInRegion = false; } } return success(); @@ -338,8 +451,8 @@ //===----------------------------------------------------------------------===// /// 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); }