Support constructing DominanceInfo with an Operation. This computes the dominance...

PiperOrigin-RevId: 251896520

diff --git a/mlir/include/mlir/Analysis/Dominance.h b/mlir/include/mlir/Analysis/Dominance.h
index f22def7..e69756e 100644 (file)
--- a/mlir/include/mlir/Analysis/Dominance.h
+++ b/mlir/include/mlir/Analysis/Dominance.h
@@ -27,6 +27,7 @@ extern template class llvm::DominatorTreeBase<mlir::Block, true>;
 namespace mlir {
 using DominanceInfoNode = llvm::DomTreeNodeBase<Block>;
 class Function;
+class Operation;
 
 namespace detail {
 template <bool IsPostDom> class DominanceInfoBase {
@@ -34,14 +35,16 @@ template <bool IsPostDom> class DominanceInfoBase {
 
 public:
   DominanceInfoBase(Function *function) { recalculate(function); }
+  DominanceInfoBase(Operation *op) { recalculate(op); }
   DominanceInfoBase(DominanceInfoBase &&) = default;
   DominanceInfoBase &operator=(DominanceInfoBase &&) = default;
 
   DominanceInfoBase(const DominanceInfoBase &) = delete;
   DominanceInfoBase &operator=(const DominanceInfoBase &) = delete;
 
-  /// Recalculate the dominance info for the provided function.
+  /// Recalculate the dominance info.
   void recalculate(Function *function);
+  void recalculate(Operation *op);
 
   /// Get the root dominance node of the given region.
   DominanceInfoNode *getRootNode(Region *region) {

diff --git a/mlir/lib/Analysis/Dominance.cpp b/mlir/lib/Analysis/Dominance.cpp
index d914f36..2c9ec00 100644 (file)
--- a/mlir/lib/Analysis/Dominance.cpp
+++ b/mlir/lib/Analysis/Dominance.cpp
@@ -23,6 +23,7 @@
 #include "mlir/Analysis/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/GenericDomTreeConstruction.h"
+
 using namespace mlir;
 using namespace mlir::detail;
 
@@ -34,7 +35,7 @@ template class llvm::DomTreeNodeBase<Block>;
 // DominanceInfoBase
 //===----------------------------------------------------------------------===//
 
-/// Recalculate the dominance info for the provided function.
+/// Recalculate the dominance info.
 template <bool IsPostDom>
 void DominanceInfoBase<IsPostDom>::recalculate(Function *function) {
   dominanceInfos.clear();
@@ -58,6 +59,23 @@ void DominanceInfoBase<IsPostDom>::recalculate(Function *function) {
   });
 }
 
+template <bool IsPostDom>
+void DominanceInfoBase<IsPostDom>::recalculate(Operation *op) {
+  dominanceInfos.clear();
+
+  /// Build the dominance for each of the operation regions.
+  op->walk([&](Operation *op) {
+    for (auto &region : op->getRegions()) {
+      // Don't compute dominance if the region is empty.
+      if (region.empty())
+        continue;
+      auto opDominance = llvm::make_unique<base>();
+      opDominance->recalculate(region);
+      dominanceInfos.try_emplace(&region, std::move(opDominance));
+    }
+  });
+}
+
 /// Return true if the specified block A properly dominates block B.
 template <bool IsPostDom>
 bool DominanceInfoBase<IsPostDom>::properlyDominates(Block *a, Block *b) {