[llvm] [HotColdSplit] Don't outline musttail calls. (PR #133177)

[Ahmed Bougacha via llvm-commits]

https://github.com/ahmedbougacha created https://github.com/llvm/llvm-project/pull/133177

musttail calls have several restrictions, generally enforcing matching calling conventions between the caller parent and musttail callee. We can't usually honor them, because the extracted function has a different signature altogether, taking inputs/outputs and returning a control-flow identifier rather than the actual return value.

Don't consider musttail calls as valid for extraction.

>From 58e3ee9a7dbf23dbd94524f06482be3e841a540e Mon Sep 17 00:00:00 2001
From: Ahmed Bougacha <ahmed at bougacha.org>
Date: Wed, 26 Mar 2025 16:13:35 -0700
Subject: [PATCH] [HotColdSplit] Don't outline musttail calls.

musttail calls have several restrictions, generally enforcing matching
calling conventions between the caller parent and musttail callee.
We can't usually honor them, because the extracted function has a
different signature altogether, taking inputs/outputs and returning
a control-flow identifier rather than the actual return value.

Don't consider musttail calls as valid for extraction.
---
 llvm/lib/Transforms/Utils/CodeExtractor.cpp   |  8 +++
 llvm/test/Transforms/HotColdSplit/musttail.ll | 53 +++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 llvm/test/Transforms/HotColdSplit/musttail.ll

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 7277603b3ec2b..216db3a4570f7 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -170,6 +170,14 @@ static bool isBlockValidForExtraction(const BasicBlock &BB,
     }
 
     if (const CallInst *CI = dyn_cast<CallInst>(I)) {
+      // musttail calls have several restrictions, generally enforcing matching
+      // calling conventions between the caller parent and musttail callee.
+      // We can't usually honor them, because the extracted function has a
+      // different signature altogether, taking inputs/outputs and returning
+      // a control-flow identifier rather than the actual return value.
+      if (CI->isMustTailCall())
+        return false;
+
       if (const Function *F = CI->getCalledFunction()) {
         auto IID = F->getIntrinsicID();
         if (IID == Intrinsic::vastart) {
diff --git a/llvm/test/Transforms/HotColdSplit/musttail.ll b/llvm/test/Transforms/HotColdSplit/musttail.ll
new file mode 100644
index 0000000000000..b7108bd6e6cd9
--- /dev/null
+++ b/llvm/test/Transforms/HotColdSplit/musttail.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=hotcoldsplit -hotcoldsplit-threshold=0 -S < %s | FileCheck %s
+
+; musttail calls can't be outlined, because they have several restrictions
+; enforcing matching calling conventions between the parent and callee.
+
+define void @test_musttail(i1 %b) {
+; CHECK-LABEL: define void @test_musttail(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[B]], label %[[IF_THEN:.*]], label %[[IF_ELSE:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[IF_ELSE]]:
+; CHECK-NEXT:    musttail call void @target(i1 [[B]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %b, label %if.then, label %if.else
+
+if.then:
+  ret void
+
+if.else:
+  musttail call void @target(i1 %b)
+  ret void
+}
+
+define void @test_tail(i1 %b) {
+; CHECK-LABEL: define void @test_tail(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br i1 [[B]], label %[[IF_THEN:.*]], label %[[CODEREPL:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    ret void
+; CHECK:       [[CODEREPL]]:
+; CHECK-NEXT:    call void @test_tail.cold.1(i1 [[B]]) #[[ATTR2:[0-9]+]]
+; CHECK-NEXT:    br label %[[IF_ELSE_RET:.*]]
+; CHECK:       [[IF_ELSE_RET]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %b, label %if.then, label %if.else
+
+if.then:
+  ret void
+
+if.else:
+  tail call void @target(i1 %b)
+  ret void
+}
+
+declare void @target() cold