[Mlir-commits] [mlir] 701f240 - [mlir] fix crash when scf utils work on llvm.func (#120688)

[llvmlistbot at llvm.org]

Author: donald chen
Date: 2024-12-20T17:03:57+08:00
New Revision: 701f2409befa7d44aea0c31494ac0d3a5d18415e

URL: https://github.com/llvm/llvm-project/commit/701f2409befa7d44aea0c31494ac0d3a5d18415e
DIFF: https://github.com/llvm/llvm-project/commit/701f2409befa7d44aea0c31494ac0d3a5d18415e.diff

LOG: [mlir] fix crash when scf utils work on llvm.func (#120688)

fixed https://github.com/llvm/llvm-project/issues/119378

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/Utils.cpp
    mlir/test/Transforms/scf-if-utils.mlir
    mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index e341c3744f1d8f..41410a0a56aa98 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -130,7 +130,7 @@ FailureOr<func::FuncOp> mlir::outlineSingleBlockRegion(RewriterBase &rewriter,
 
   // Outline before current function.
   OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(region.getParentOfType<func::FuncOp>());
+  rewriter.setInsertionPoint(region.getParentOfType<FunctionOpInterface>());
 
   SetVector<Value> captures;
   getUsedValuesDefinedAbove(region, captures);

diff  --git a/mlir/test/Transforms/scf-if-utils.mlir b/mlir/test/Transforms/scf-if-utils.mlir
index 449be18483741c..fd59f5e9295e6e 100644
--- a/mlir/test/Transforms/scf-if-utils.mlir
+++ b/mlir/test/Transforms/scf-if-utils.mlir
@@ -73,3 +73,30 @@ func.func @outline_empty_if_else(%cond: i1, %a: index, %b: memref<?xf32>, %c: i8
   }
   return
 }
+
+// -----
+
+// This test checks scf utils can work on llvm func.
+
+//      CHECK: func @outlined_then0() {
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+//      CHECK: func @outlined_else0(%{{.*}}: i1, %{{.*}}: i32) {
+// CHECK-NEXT:   "some_op"(%{{.*}}, %{{.*}}) : (i1, i32) -> ()
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+//      CHECK: llvm.func @llvm_func(%{{.*}}: i1, %{{.*}}: i32) {
+// CHECK-NEXT:   scf.if %{{.*}} {
+// CHECK-NEXT:     func.call @outlined_then0() : () -> ()
+// CHECK-NEXT:   } else {
+// CHECK-NEXT:     func.call @outlined_else0(%{{.*}}, %{{.*}}) : (i1, i32) -> ()
+// CHECK-NEXT:   }
+// CHECK-NEXT:   llvm.return
+// CHECK-NEXT: }
+llvm.func @llvm_func(%cond: i1, %a: i32) {
+  scf.if %cond {
+  } else {
+    "some_op"(%cond, %a) : (i1, i32) -> ()
+  }
+  llvm.return
+}

diff  --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index 3ff7f9966e93da..a3be1f94fa28a3 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -79,6 +79,10 @@ struct TestSCFIfUtilsPass
   StringRef getDescription() const final { return "test scf.if utils"; }
   explicit TestSCFIfUtilsPass() = default;
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<func::FuncDialect>();
+  }
+
   void runOnOperation() override {
     int count = 0;
     getOperation().walk([&](scf::IfOp ifOp) {