From fa187e0f3bb491d369c6cb4b964727ff3f306571 Mon Sep 17 00:00:00 2001
From: River Riddle <riverriddle@google.com>
Date: Thu, 6 Jun 2019 11:54:14 -0700
Subject: [PATCH] Support constructing DominanceInfo with an Operation. This
 computes the dominance information for any nested regions within the
 operation.

PiperOrigin-RevId: 251896520
---
 mlir/include/mlir/Analysis/Dominance.h |  5 ++++-
 mlir/lib/Analysis/Dominance.cpp        | 20 +++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Dominance.h b/mlir/include/mlir/Analysis/Dominance.h
index f22def7..e69756e 100644
--- 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
--- 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) {
-- 
2.7.4