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

Author: junfengd-nv

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;
 

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
+  }
+}

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