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

[llvmlistbot at llvm.org]

https://github.com/junfengd-nv created https://github.com/llvm/llvm-project/pull/124026

Issue #82401 shows a problem where the MLIR inliner enters an infinite loop when call graph contains multiple indirect cycles. For example, the test case inlining-recursive-2.mlir contains cycles: parent1 -> child ->parent1 and parent2->child ->parent2.  Current check in the shouldInline function fails to detect these cases.  The fix is stop inlining when a pattern like A->B ->A is identified. 


>From cec54b5cf690a5a91acf52d1aa6b1c202319ca5b Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd at nvidia.com>
Date: Wed, 22 Jan 2025 14:27:44 -0800
Subject: [PATCH] Fix Issue#82401: Infinite loop in MLIR inliner for recusive
 call.

---
 mlir/lib/Transforms/Utils/Inliner.cpp         |  4 +-
 .../test/Transforms/inlining-recursive-2.mlir | 37 +++++++++++++++++++
 mlir/test/Transforms/inlining-recursive.mlir  |  2 +-
 3 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Transforms/inlining-recursive-2.mlir

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..7805c0ba348c2c
--- /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 a02fe69133ad87..18b0f28a7cf37d 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