[llvm] [Transforms][CodeExtraction] bug fix regions with stackrestore (PR #118564)

[via llvm-commits]

https://github.com/AbhayKanhere created https://github.com/llvm/llvm-project/pull/118564

   Ensure code extraction for outlining to a function does not create a function with stacksave of caller 
 to restore stack (e.g. tail call) .

>From 272979d4d20741b6c28bac9253b9ef4931270560 Mon Sep 17 00:00:00 2001
From: Abhay Kanhere <abhay at kanhere.net>
Date: Tue, 3 Dec 2024 16:13:45 -0800
Subject: [PATCH] [Transforms][CodeExtraction] bug fix regions with
 stackrestore    Ensure code extraction for outlining to a function does not  
 create a function with tail-call i.e. using stacksave of caller   to restore
 stack. Code Extractor can include all other   basic blocks without such stack
 restore blocks.

---
 llvm/lib/Transforms/Utils/CodeExtractor.cpp   | 16 ++++++
 .../outline-inner-region-stacktoocomplex.ll   | 51 +++++++++++++++++++
 2 files changed, 67 insertions(+)
 create mode 100644 llvm/test/Transforms/HotColdSplit/outline-inner-region-stacktoocomplex.ll

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 6539f924c2edf4..c4efb597e615d1 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -627,6 +627,22 @@ bool CodeExtractor::isEligible() const {
         return false;
     }
   }
+  // stacksave as input implies stackrestore in the outlined function
+  // This can confuse prologue epilogue insertion phase
+  for (BasicBlock *BB : Blocks) {
+    for (Instruction &II : *BB) {
+      for (auto &OI : II.operands()) {
+        Value *V = OI;
+        if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(V)) {
+          if (II->getIntrinsicID() == Intrinsic::stacksave) {
+            if (definedInCaller(Blocks, V)) {
+              return false;
+            }
+          }
+        }
+      }
+    }
+  }
   return true;
 }
 
diff --git a/llvm/test/Transforms/HotColdSplit/outline-inner-region-stacktoocomplex.ll b/llvm/test/Transforms/HotColdSplit/outline-inner-region-stacktoocomplex.ll
new file mode 100644
index 00000000000000..eee3b16c5e0060
--- /dev/null
+++ b/llvm/test/Transforms/HotColdSplit/outline-inner-region-stacktoocomplex.ll
@@ -0,0 +1,51 @@
+; RUN: opt -S -passes=hotcoldsplit -hotcoldsplit-max-params=1 < %s | FileCheck %s
+
+target datalayout = "E-m:a-p:32:32-i64:64-n32"
+target triple = "powerpc64-ibm-aix7.2.0.0"
+
+define void @foo(i32 %cond) {
+; CHECK-LABEL: define {{.*}}@foo(
+; CHECK:       if.then:
+; CHECK:       if.then1:
+; CHECK:       if.else:
+; CHECK:       call {{.*}}@sink
+; CHECK:       if.end:
+; CHECK:       if.end2:
+;
+entry:
+  %cond.addr = alloca i32
+  store i32 %cond, ptr %cond.addr
+  %0 = load i32, ptr %cond.addr
+  %stks =  tail call ptr @llvm.stacksave.p0()
+  %tobool = icmp ne i32 %0, 0
+  br i1 %tobool, label %if.then, label %if.end2
+
+if.then:                                          ; preds = %entry
+  %1 = load i32, ptr %cond.addr
+  call void @sink(i32 %0)
+  %cmp = icmp sgt i32 %1, 10
+  br i1 %cmp, label %if.then1, label %if.else
+
+if.then1:                                         ; preds = %if.then
+  call void @sideeffect(i32 2)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  call void @sink(i32 0)
+  call void @sideeffect(i32 0)
+  call void @llvm.stackrestore.p0(ptr %stks)
+  br label %if.end
+
+
+if.end:                                           ; preds = %if.else, %if.then1
+  br label %if.end2
+
+if.end2:                                          ; preds = %entry
+  call void @sideeffect(i32 1)
+  ret void
+}
+
+
+declare void @sideeffect(i32)
+
+declare void @sink(i32) cold