Fix 40056: Prevent outlining of blocks with token type instructions

[hiraditya]

Hot cold splitting should not outline:

1. Basic blocks with token type instructions
2. Functions with scoped EH personality (As suggested by Vedant in Function has token parameter but isn't an intrinsic #40056 (comment))

Fixes: #40056

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: AdityaK (hiraditya)

Changes

Hot cold splitting should not outline:

1. Basic blocks with token type instructions
2. Functions with scoped EH personality

Fixes: #40056

---

Full diff: https://github.com/llvm/llvm-project/pull/99759.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/HotColdSplitting.cpp (+23-1)
• (added) llvm/test/Transforms/HotColdSplit/pr40056.ll (+72)

diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
index 5aefcbf13182c..849bfd3d95caf 100644
--- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -43,6 +43,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/EHPersonalities.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/User.h"
@@ -168,10 +169,25 @@ static bool mayExtractBlock(const BasicBlock &BB) {
   //
   // Resumes that are not reachable from a cleanup landing pad are considered to
   // be unreachable. It’s not safe to split them out either.
+
   if (BB.hasAddressTaken() || BB.isEHPad())
     return false;
   auto Term = BB.getTerminator();
-  return !isa<InvokeInst>(Term) && !isa<ResumeInst>(Term);
+  if (isa<InvokeInst>(Term) || isa<ResumeInst>(Term))
+    return false;
+
+  // Do not outline basic blocks that have token type instructions. e.g.,
+  // exception:
+  // %0 = cleanuppad within none []
+  // call void @"?terminate@@YAXXZ"() [ "funclet"(token %0) ]
+  // br label %continue-exception
+  if (llvm::any_of(BB, [](const Instruction &I){
+    return I.getType()->isTokenTy();
+  })) {
+    return false;
+  }
+
+  return true;
 }
 
 /// Mark \p F cold. Based on this assumption, also optimize it for minimum size.
@@ -258,6 +274,12 @@ bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
       F.hasFnAttribute(Attribute::SanitizeMemory))
     return false;
 
+  // Support for outlining WinEH code has not been tested sufficiently. It may
+  // trigger verifier failures or miscompiles (see llvm.org/PR40710).
+  if (F.hasPersonalityFn())
+    if (isScopedEHPersonality(classifyEHPersonality(F.getPersonalityFn())))
+      return false;
+
   return true;
 }
 
diff --git a/llvm/test/Transforms/HotColdSplit/pr40056.ll b/llvm/test/Transforms/HotColdSplit/pr40056.ll
new file mode 100644
index 0000000000000..950b62c673fbf
--- /dev/null
+++ b/llvm/test/Transforms/HotColdSplit/pr40056.ll
@@ -0,0 +1,72 @@
+; RUN: opt -passes=hotcoldsplit -hotcoldsplit-threshold=-1 -S < %s | FileCheck %s
+; Hot cold splitting should not outline:
+; 1. Basic blocks with token type instructions
+; 2. Functions with scoped EH personality
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.0"
+
+; CHECK-LABEL: define {{.*}}@with_funclet
+; CHECK-NOT: with_funclet.cold
+define void @with_funclet() personality ptr @__CxxFrameHandler3 {
+entry:
+  invoke void @fYAXXZ()
+          to label %normal unwind label %exception
+
+normal:                                           ; preds = %entry
+  ret void
+
+exception:                                        ; preds = %entry
+  %0 = cleanuppad within none []
+  call void @terminateYAXXZ() [ "funclet"(token %0) ]
+  br label %continueexception
+
+continueexception:                                ; preds = %exception
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}}@with_personality
+; CHECK-NOT: with_personality.cold
+define void @with_personality(i32 %cond) personality ptr @__CxxFrameHandler3 {
+entry:
+  %cond.addr = alloca i32
+  store i32 %cond, ptr %cond.addr
+  %0 = load i32, ptr %cond.addr
+  %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
+  %cmp = icmp sgt i32 %1, 10
+  br i1 %cmp, label %if.then1, label %if.else
+
+if.then1:                                         ; preds = %if.then
+  call void @sideeffect(i32 0)
+  br label %if.end
+
+if.else:                                          ; preds = %if.then
+  call void @sideeffect(i32 1)
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then1
+  call void (...) @sink()
+  ret void
+
+if.end2:                                          ; preds = %entry
+  call void @sideeffect(i32 2)
+  ret void
+}
+
+declare i32 @__CxxFrameHandler3(...)
+
+declare void @fYAXXZ()
+
+declare void @bar() #0
+
+declare void @terminateYAXXZ()
+
+declare void @sideeffect(i32)
+
+declare void @sink(...) #0
+
+attributes #0 = { cold }

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[compnerd]

Do you see any way to actually do the outlining where it is profitable? It seems that the comment saying it is untested is somewhat misleading unless there's reason to believe it can be beneficial.

[hiraditya]

Do you see any way to actually do the outlining where it is profitable? It seems that the comment saying it is untested is somewhat misleading unless there's reason to believe it can be beneficial.

Amended the comment. I dont have a good rationale where it might be profitable. If someone finds it profitable, we can always revisit this but yeah there is no reason to indicate that in comment right now.

[compnerd]

Great, thank you!