[Mlir-commits] [mlir] 83df39c - [mlir][inline] Fix Issue#82401: Infinite loop in MLIR inliner for indirect recursive call. (#124026)

[llvmlistbot at llvm.org]

Author: junfengd-nv
Date: 2025-01-24T11:06:37-08:00
New Revision: 83df39c649fe1b1dd556d8f2160999c65ce497eb

URL: https://github.com/llvm/llvm-project/commit/83df39c649fe1b1dd556d8f2160999c65ce497eb
DIFF: https://github.com/llvm/llvm-project/commit/83df39c649fe1b1dd556d8f2160999c65ce497eb.diff

LOG: [mlir][inline] Fix Issue#82401: Infinite loop in MLIR inliner for indirect recursive call. (#124026)

Added: 
    mlir/test/Transforms/inlining-recursive-2.mlir

Modified: 
    mlir/lib/Transforms/Utils/Inliner.cpp
    mlir/test/Transforms/inlining-recursive.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 8acfc96d2b611b..756f5e379e7ddd 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -713,9 +713,11 @@ bool Inliner::Impl::shouldInline(ResolvedCall &resolvedCall) {
     return false;
 
   // Don't allow inlining if the target is a self-recursive function.
+  // Don't allow inlining if the call graph is like A->B->A.
   if (llvm::count_if(*resolvedCall.targetNode,
                      [&](CallGraphNode::Edge const &edge) -> bool {
-                       return edge.getTarget() == resolvedCall.targetNode;
+                       return edge.getTarget() == resolvedCall.targetNode ||
+                              edge.getTarget() == resolvedCall.sourceNode;
                      }) > 0)
     return false;
 

diff  --git a/mlir/test/Transforms/inlining-recursive-2.mlir b/mlir/test/Transforms/inlining-recursive-2.mlir
new file mode 100644
index 00000000000000..e50cf9695c4a3b
--- /dev/null
+++ b/mlir/test/Transforms/inlining-recursive-2.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-opt %s -inline='default-pipeline=' | FileCheck %s
+// RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=' | FileCheck %s
+
+module {
+  // CHECK-LABEL: func.func @parent1
+  func.func @parent1(%arg0: i32) -> i32 {
+    // CHECK: call @child
+    %0 = call @child(%arg0) : (i32) -> i32
+    return %0 : i32
+  }
+
+  // CHECK-LABEL: func.func @parent2
+  func.func @parent2(%arg0: i32) -> i32 {
+    // CHECK: call @child
+    %0 = call @child(%arg0) : (i32) -> i32
+    return %0 : i32
+  }
+
+  // CHECK-LABEL: func.func @child
+  func.func @child(%arg0: i32) -> i32 {
+    %c10_i32 = arith.constant 10 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %0 = arith.cmpi sge, %arg0, %c10_i32 : i32
+    %1 = scf.if %0 -> (i32) {
+      scf.yield %arg0 : i32
+    } else {
+      %2 = arith.addi %arg0, %c1_i32 : i32
+      // CHECK: call @parent1
+      // CHECK: call @parent2
+      %3 = func.call @parent1(%2) : (i32) -> i32
+      %4 = func.call @parent2(%2) : (i32) -> i32
+      %5 = arith.addi %3, %4 : i32
+      scf.yield %5 : i32
+    }
+    return %1 : i32
+  }
+}

diff  --git a/mlir/test/Transforms/inlining-recursive.mlir b/mlir/test/Transforms/inlining-recursive.mlir
index 403accd8b7ee8c..f953935475e1a7 100644
--- a/mlir/test/Transforms/inlining-recursive.mlir
+++ b/mlir/test/Transforms/inlining-recursive.mlir
@@ -17,7 +17,7 @@ func.func @foo0(%arg0 : i32) -> i32 {
 
 // CHECK-LABEL: func.func @foo1
 func.func @foo1(%arg0 : i32) -> i32 {
-  // CHECK:    call @foo1
+  // CHECK:    call @foo0
   %0 = arith.constant 1 : i32
   %1 = arith.subi %arg0, %0 : i32
   %2 = call @foo0(%1) : (i32) -> i32