From 267c6ed85f62ef521c93982f988e56ab9ed535de Mon Sep 17 00:00:00 2001
From: AdityaK <hiraditya@msn.com>
Date: Sat, 20 Jul 2024 11:53:11 -0700
Subject: [PATCH] Fix 40056: Prevent outlining of blocks with token type
 instructions (#99759)

Hot cold splitting should not outline:
1. Basic blocks with token type instructions
1. Functions with scoped EH personality (As suggested by Vedant in
https://github.com/llvm/llvm-project/issues/40056#issuecomment-981009129)

Fixes: #40056
(cherry picked from commit c59ee7ec627533e00f654d052ade51c9daf7e722)
---
 llvm/lib/Transforms/IPO/HotColdSplitting.cpp | 22 +++++-
 llvm/test/Transforms/HotColdSplit/pr40056.ll | 72 ++++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/HotColdSplit/pr40056.ll

diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
index 599ace9ca79f3..47ab7dde9cc0b 100644
--- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -39,6 +39,7 @@
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/EHPersonalities.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
@@ -136,10 +137,24 @@ 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.
@@ -203,6 +218,11 @@ bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
       F.hasFnAttribute(Attribute::SanitizeMemory))
     return false;
 
+  // Do not outline scoped EH personality functions.
+  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 }