From 65839210c5059acbc82b9cb709ad61d8dfd0a369 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann@nextsilicon.com>
Date: Fri, 6 Dec 2024 07:53:14 +0000
Subject: [PATCH] [MLIR][ControlFlowToSCF] Extend with support for
 FuncOpInterface

This commit ensures that the pass to lift from `cf` to `scf` can be
applied on all operations that implement `FuncOpInterface`.
---
 .../ControlFlowToSCF/ControlFlowToSCF.cpp     |  4 +--
 .../Conversion/ControlFlowToSCF/test.mlir     | 26 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
index 1c592d665f3e4..a54281dfd3375 100644
--- a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
+++ b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
@@ -164,8 +164,8 @@ struct LiftControlFlowToSCF
 
     bool changed = false;
     Operation *op = getOperation();
-    WalkResult result = op->walk([&](func::FuncOp funcOp) {
-      if (funcOp.getBody().empty())
+    WalkResult result = op->walk([&](FunctionOpInterface funcOp) {
+      if (funcOp.getFunctionBody().empty())
         return WalkResult::advance();
 
       auto &domInfo = funcOp != op ? getChildAnalysis<DominanceInfo>(funcOp)
diff --git a/mlir/test/Conversion/ControlFlowToSCF/test.mlir b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
index f6e7fb265790b..3a2ec521aae0a 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -756,3 +756,29 @@ func.func @nested_region_outside_loop_use() {
 
 // CHECK: scf.execute_region
 // CHECK-NEXT: "test.foo"(%[[RES]])
+
+// -----
+
+llvm.func @llvm_func() {
+  %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"() : () -> ()
+  return
+}
+
+// CHECK-LABEL: llvm.func @llvm_func
+// 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"()
+// CHECK-NEXT: return