[Mlir-commits] [mlir] ad23127 - [mlir][inline] avoid inline self-recursive function (#83092)

[llvmlistbot at llvm.org]

Author: Congcong Cai
Date: 2024-03-12T06:49:09+08:00
New Revision: ad23127222fe23e28ac3deaa16f3ae64d13b7b6f

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

LOG: [mlir][inline] avoid inline self-recursive function (#83092)

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

Modified: 
    mlir/lib/Transforms/Utils/Inliner.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 74776a73db9aaa..f227cedb269d81 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -21,6 +21,7 @@
 #include "mlir/Support/DebugStringHelper.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SCCIterator.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 
@@ -711,6 +712,13 @@ bool Inliner::Impl::shouldInline(ResolvedCall &resolvedCall) {
   if (resolvedCall.call->hasTrait<OpTrait::IsTerminator>())
     return false;
 
+  // Don't allow inlining if the target is a self-recursive function.
+  if (llvm::count_if(*resolvedCall.targetNode,
+                     [&](CallGraphNode::Edge const &edge) -> bool {
+                       return edge.getTarget() == resolvedCall.targetNode;
+                     }) > 0)
+    return false;
+
   // Don't allow inlining if the target is an ancestor of the call. This
   // prevents inlining recursively.
   Region *callableRegion = resolvedCall.targetNode->getCallableRegion();

diff  --git a/mlir/test/Transforms/inlining-recursive-self.mlir b/mlir/test/Transforms/inlining-recursive-self.mlir
new file mode 100644
index 00000000000000..5cc922db8e9786
--- /dev/null
+++ b/mlir/test/Transforms/inlining-recursive-self.mlir
@@ -0,0 +1,22 @@
+// RUN: mlir-opt %s -inline='default-pipeline=''' | FileCheck %s
+// RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=''' | FileCheck %s
+
+// CHECK-LABEL: func.func @b0
+func.func @b0() {
+  // CHECK:         call @b0
+  // CHECK-NEXT:    call @b1
+  // CHECK-NEXT:    call @b0
+  // CHECK-NEXT:    call @b1
+  // CHECK-NEXT:    call @b0
+  func.call @b0() : () -> ()
+  func.call @b1() : () -> ()
+  func.call @b0() : () -> ()
+  func.call @b1() : () -> ()
+  func.call @b0() : () -> ()
+  return
+}
+func.func @b1() {
+  func.call @b1() : () -> ()
+  func.call @b1() : () -> ()
+  return
+}