[Mlir-commits] [mlir] [MLIR] Make ControlFlowToSCF work on regions, not func::FuncOp (PR #96725)

[Ivan R. Ivanov]

https://github.com/ivanradanov created https://github.com/llvm/llvm-project/pull/96725

None

>From 1c695ddeb3fe557a46217817a8732eada6874122 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <ivanov.i.aa at m.titech.ac.jp>
Date: Wed, 26 Jun 2024 11:45:01 +0900
Subject: [PATCH] [MLIR] Make ControlFlowToSCF work on regions, not
 func::FuncOp

---
 .../ControlFlowToSCF/ControlFlowToSCF.cpp     | 12 ++++++----
 .../Conversion/ControlFlowToSCF/test.mlir     | 24 +++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
index d3ee89743da9d..209c441cf21d8 100644
--- a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
+++ b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
@@ -165,12 +165,14 @@ struct LiftControlFlowToSCF
 
     bool changed = false;
     Operation *op = getOperation();
-    WalkResult result = op->walk([&](func::FuncOp funcOp) {
-      if (funcOp.getBody().empty())
+    WalkResult result = op->walk([&](Region *region) {
+      if (region->empty())
         return WalkResult::advance();
 
-      auto &domInfo = funcOp != op ? getChildAnalysis<DominanceInfo>(funcOp)
-                                   : getAnalysis<DominanceInfo>();
+      Operation *regionParent = region->getParentOp();
+      auto &domInfo = regionParent != op
+                          ? getChildAnalysis<DominanceInfo>(regionParent)
+                          : getAnalysis<DominanceInfo>();
 
       auto visitor = [&](Operation *innerOp) -> WalkResult {
         for (Region &reg : innerOp->getRegions()) {
@@ -184,7 +186,7 @@ struct LiftControlFlowToSCF
         return WalkResult::advance();
       };
 
-      if (funcOp->walk<WalkOrder::PostOrder>(visitor).wasInterrupted())
+      if (region->walk<WalkOrder::PostOrder>(visitor).wasInterrupted())
         return WalkResult::interrupt();
 
       return WalkResult::advance();
diff --git a/mlir/test/Conversion/ControlFlowToSCF/test.mlir b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
index f6e7fb265790b..40b157da3667a 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -756,3 +756,27 @@ func.func @nested_region_outside_loop_use() {
 
 // CHECK: scf.execute_region
 // CHECK-NEXT: "test.foo"(%[[RES]])
+
+// -----
+
+"test.test5" () ({
+  %cond = "test.test1"() : () -> i1
+  cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+  "test.test2"() : () -> ()
+  cf.br ^bb3
+^bb2:
+  "test.test3"() : () -> ()
+  cf.br ^bb3
+^bb3:
+  "test.test4"() : () -> ()
+}) : () -> ()
+
+// CHECK-LABEL: test.test5
+// CHECK:         %[[COND:.*]] = "test.test1"
+// CHECK-NEXT:    scf.if %[[COND]]
+// CHECK-NEXT:      "test.test2"
+// CHECK-NEXT:    else
+// CHECK-NEXT:      "test.test3"
+// CHECK-NEXT:    }
+// CHECK-NEXT:    "test.test4"