[Mlir-commits] [mlir] [MLIR][SCF] Speed up ConditionPropagation (PR #166080)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: William Moses (wsmoses)

<details>
<summary>Changes</summary>

Fixes https://github.com/llvm/llvm-project/issues/166039

---
Full diff: https://github.com/llvm/llvm-project/pull/166080.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+45-3) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 2946b53c8cb36..2d5a0052d4c53 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2565,6 +2565,41 @@ struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
 struct ConditionPropagation : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
+  enum class Parent {
+    Then,
+    Else,
+    None
+  };
+
+  static Parent getParentType(Region *toCheck, IfOp op, DenseMap<Region*, Parent> &cache) {
+    SmallVector<Region*> seen;
+    while (toCheck) {
+      auto found = cache.find(toCheck);
+      if (found != cache.end()) {
+        return found->second;
+      }
+      seen.push_back(toCheck);
+      if (&op.getThenRegion() == toCheck) {
+        for (auto v : seen) {
+          cache[v] = Parent::Then;
+        }
+        return Parent::Then;
+      }
+      if (&op.getElseRegion() == toCheck) {
+        for (auto v : seen) {
+          cache[v] = Parent::Else;
+        }
+        return Parent::Else;
+      }
+      toCheck = toCheck->getParentRegion();
+    }
+
+    for (auto v : seen) {
+      cache[v] = Parent::None;
+    }
+    return Parent::None;
+  }
+
   LogicalResult matchAndRewrite(IfOp op,
                                 PatternRewriter &rewriter) const override {
     // Early exit if the condition is constant since replacing a constant
@@ -2580,9 +2615,11 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
     Value constantTrue = nullptr;
     Value constantFalse = nullptr;
 
+    DenseMap<Region*, Parent> regionCache;
     for (OpOperand &use :
          llvm::make_early_inc_range(op.getCondition().getUses())) {
-      if (op.getThenRegion().isAncestor(use.getOwner()->getParentRegion())) {
+      switch(getParentType(use.getOwner()->getParentRegion(), op, cache)) {
+      case Parent::Then:{
         changed = true;
 
         if (!constantTrue)
@@ -2591,8 +2628,9 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
 
         rewriter.modifyOpInPlace(use.getOwner(),
                                  [&]() { use.set(constantTrue); });
-      } else if (op.getElseRegion().isAncestor(
-                     use.getOwner()->getParentRegion())) {
+        break;
+      }
+      case Parent::Else:{
         changed = true;
 
         if (!constantFalse)
@@ -2601,6 +2639,10 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
 
         rewriter.modifyOpInPlace(use.getOwner(),
                                  [&]() { use.set(constantFalse); });
+        break;
+      }
+      case Parent::None:
+        break;
       }
     }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/166080