[CodeGen][WinEH] Update saved esp for inlined inallocas (llvm#116585)

MuellerMP · web-flow · commit 926a71f0c9ff · 2024-11-21T13:58:43.000-08:00
This fixes issue llvm#116583 When inalloca calls are inlined the static stack pointer saving prolog of X86WinEHState breaks due to dynamic allocas. In this case we need to update the saved esp for every inalloca and for every stackrestore also related to inalloca.
diff --git a/llvm/lib/Target/X86/X86WinEHState.cpp b/llvm/lib/Target/X86/X86WinEHState.cpp
@@ -23,6 +23,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsX86.h"
 #include "llvm/IR/Module.h"
@@ -41,7 +42,7 @@ class WinEHStatePass : public FunctionPass {
 public:
   static char ID; // Pass identification, replacement for typeid.
 
-  WinEHStatePass() : FunctionPass(ID) { }
+  WinEHStatePass() : FunctionPass(ID) {}
 
   bool runOnFunction(Function &Fn) override;
 
@@ -75,6 +76,8 @@ class WinEHStatePass : public FunctionPass {
   int getStateForCall(DenseMap<BasicBlock *, ColorVector> &BlockColors,
                       WinEHFuncInfo &FuncInfo, CallBase &Call);
 
+  void updateEspForInAllocas(Function &F);
+
   // Module-level type getters.
   Type *getEHLinkRegistrationType();
   Type *getSEHRegistrationType();
@@ -100,6 +103,9 @@ class WinEHStatePass : public FunctionPass {
   /// fs:00 chain and the current state.
   AllocaInst *RegNode = nullptr;
 
+  // Struct type of RegNode. Used for GEPing.
+  Type *RegNodeTy = nullptr;
+
   // The allocation containing the EH security guard.
   AllocaInst *EHGuardNode = nullptr;
 
@@ -152,8 +158,7 @@ bool WinEHStatePass::runOnFunction(Function &F) {
   // Check the personality. Do nothing if this personality doesn't use funclets.
   if (!F.hasPersonalityFn())
     return false;
-  PersonalityFn =
-      dyn_cast<Function>(F.getPersonalityFn()->stripPointerCasts());
+  PersonalityFn = dyn_cast<Function>(F.getPersonalityFn()->stripPointerCasts());
   if (!PersonalityFn)
     return false;
   Personality = classifyEHPersonality(PersonalityFn);
@@ -188,11 +193,13 @@ bool WinEHStatePass::runOnFunction(Function &F) {
   // numbers into an immutable analysis pass.
   WinEHFuncInfo FuncInfo;
   addStateStores(F, FuncInfo);
+  updateEspForInAllocas(F);
 
   // Reset per-function state.
   PersonalityFn = nullptr;
   Personality = EHPersonality::Unknown;
   UseStackGuard = false;
+  RegNodeTy = nullptr;
   RegNode = nullptr;
   EHGuardNode = nullptr;
 
@@ -269,9 +276,6 @@ void WinEHStatePass::emitExceptionRegistrationRecord(Function *F) {
   assert(Personality == EHPersonality::MSVC_CXX ||
          Personality == EHPersonality::MSVC_X86SEH);
 
-  // Struct type of RegNode. Used for GEPing.
-  Type *RegNodeTy;
-
   IRBuilder<> Builder(&F->getEntryBlock(), F->getEntryBlock().begin());
   Type *Int8PtrType = Builder.getPtrTy();
   Type *Int32Ty = Builder.getInt32Ty();
@@ -387,11 +391,11 @@ Function *WinEHStatePass::generateLSDAInEAXThunk(Function *ParentFunc) {
   FunctionType *TargetFuncTy =
       FunctionType::get(Int32Ty, ArrayRef(&ArgTys[0], 5),
                         /*isVarArg=*/false);
-  Function *Trampoline =
-      Function::Create(TrampolineTy, GlobalValue::InternalLinkage,
-                       Twine("__ehhandler$") + GlobalValue::dropLLVMManglingEscape(
-                                                   ParentFunc->getName()),
-                       TheModule);
+  Function *Trampoline = Function::Create(
+      TrampolineTy, GlobalValue::InternalLinkage,
+      Twine("__ehhandler$") +
+          GlobalValue::dropLLVMManglingEscape(ParentFunc->getName()),
+      TheModule);
   if (auto *C = ParentFunc->getComdat())
     Trampoline->setComdat(C);
   BasicBlock *EntryBB = BasicBlock::Create(Context, "entry", Trampoline);
@@ -482,8 +486,8 @@ void WinEHStatePass::rewriteSetJmpCall(IRBuilder<> &Builder, Function &F,
     NewCall = NewCI;
   } else {
     auto *II = cast<InvokeInst>(&Call);
-    NewCall = Builder.CreateInvoke(
-        SetJmp3, II->getNormalDest(), II->getUnwindDest(), Args, OpBundles);
+    NewCall = Builder.CreateInvoke(SetJmp3, II->getNormalDest(),
+                                   II->getUnwindDest(), Args, OpBundles);
   }
   NewCall->setCallingConv(Call.getCallingConv());
   NewCall->setAttributes(Call.getAttributes());
@@ -774,3 +778,27 @@ void WinEHStatePass::insertStateNumberStore(Instruction *IP, int State) {
                                               RegNode, StateFieldIndex);
   Builder.CreateStore(Builder.getInt32(State), StateField);
 }
+
+void WinEHStatePass::updateEspForInAllocas(Function &F) {
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      if (auto *Alloca = dyn_cast<AllocaInst>(&I)) {
+        if (Alloca->isStaticAlloca())
+          continue;
+        IRBuilder<> Builder(Alloca->getNextNonDebugInstruction());
+        // SavedESP = llvm.stacksave()
+        Value *SP = Builder.CreateStackSave();
+        Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
+      }
+
+      if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
+        if (II->getIntrinsicID() != Intrinsic::stackrestore)
+          continue;
+        IRBuilder<> Builder(II->getNextNonDebugInstruction());
+        // SavedESP = llvm.stacksave()
+        Value *SP = Builder.CreateStackSave();
+        Builder.CreateStore(SP, Builder.CreateStructGEP(RegNodeTy, RegNode, 0));
+      }
+    }
+  }
+}
diff --git a/llvm/test/CodeGen/WinEH/wineh-dynamic-alloca.ll b/llvm/test/CodeGen/WinEH/wineh-dynamic-alloca.ll
@@ -0,0 +1,95 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s | FileCheck %s
+target datalayout = "e-m:x-p:32:32-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32-a:0:32-S32"
+target triple = "i386-pc-windows-msvc"
+
+%struct.Foo = type { i32, i32 }
+
+@bar = external global i1
+
+define dso_local noundef i32 @foo() local_unnamed_addr #0 personality ptr @__CxxFrameHandler3 {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushl %ebp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    pushl %ebx
+; CHECK-NEXT:    pushl %edi
+; CHECK-NEXT:    pushl %esi
+; CHECK-NEXT:    subl $16, %esp
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $-1, -16(%ebp)
+; CHECK-NEXT:    leal -24(%ebp), %eax
+; CHECK-NEXT:    movl $___ehhandler$foo, -20(%ebp)
+; CHECK-NEXT:    movl %fs:0, %ecx
+; CHECK-NEXT:    movl %ecx, -24(%ebp)
+; CHECK-NEXT:    movl %eax, %fs:0
+; CHECK-NEXT:    cmpb $1, _bar
+; CHECK-NEXT:    je LBB0_1
+; CHECK-NEXT:  LBB0_5: # %exit
+; CHECK-NEXT:  $ehgcr_0_5:
+; CHECK-NEXT:    movl -24(%ebp), %eax
+; CHECK-NEXT:    movl %eax, %fs:0
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    leal -12(%ebp), %esp
+; CHECK-NEXT:    popl %esi
+; CHECK-NEXT:    popl %edi
+; CHECK-NEXT:    popl %ebx
+; CHECK-NEXT:    popl %ebp
+; CHECK-NEXT:    retl
+; CHECK-NEXT:  LBB0_1: # %if.then
+; CHECK-NEXT:    pushl %eax
+; CHECK-NEXT:    pushl %eax
+; CHECK-NEXT:    movl %esp, %eax
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $123, (%eax)
+; CHECK-NEXT:    movl $0, -16(%ebp)
+; CHECK-NEXT:    calll _alwaysthrows
+; CHECK-NEXT:  # %bb.4: # %unreachable.i
+; CHECK-NEXT:  LBB0_3: # Block address taken
+; CHECK-NEXT:    # %catch.i
+; CHECK-NEXT:    addl $12, %ebp
+; CHECK-NEXT:    jmp LBB0_5
+; CHECK-NEXT:    .def "?catch$2@?0?foo@4HA";
+; CHECK-NEXT:    .scl 3;
+; CHECK-NEXT:    .type 32;
+; CHECK-NEXT:    .endef
+; CHECK-NEXT:    .p2align 4
+; CHECK-NEXT:  "?catch$2@?0?foo@4HA":
+; CHECK-NEXT:  LBB0_2: # %catch.i
+; CHECK-NEXT:    pushl %ebp
+; CHECK-NEXT:    addl $12, %ebp
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $LBB0_3, %eax
+; CHECK-NEXT:    popl %ebp
+; CHECK-NEXT:    retl # CATCHRET
+; CHECK-NEXT:  Lfunc_end0:
+entry:
+  %cmp = load i1, ptr @bar
+  br i1 %cmp, label %if.then, label %exit
+
+if.then:                                          ; preds = %entry
+  %foo = alloca <{ %struct.Foo }>, align 4
+  store i32 123, ptr %foo, align 4
+  invoke void @alwaysthrows() #1
+          to label %unreachable.i unwind label %catch.dispatch.i
+
+catch.dispatch.i:                                 ; preds = %if.then
+  %3 = catchswitch within none [label %catch.i] unwind to caller
+
+catch.i:                                          ; preds = %catch.dispatch.i
+  %4 = catchpad within %3 [ptr null, i32 64, ptr null]
+  catchret from %4 to label %exit
+
+unreachable.i:                                    ; preds = %if.then
+  unreachable
+
+exit:                                             ; preds = %entry, %catch.i
+  ret i32 0
+}
+
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+declare dso_local void @alwaysthrows() local_unnamed_addr
+
+attributes #0 = { norecurse "min-legal-vector-width"="0" "target-cpu"="pentium4" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #1 = { noreturn }
diff --git a/llvm/test/CodeGen/WinEH/wineh-inlined-inalloca.ll b/llvm/test/CodeGen/WinEH/wineh-inlined-inalloca.ll
@@ -0,0 +1,90 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s | FileCheck %s
+target datalayout = "e-m:x-p:32:32-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32-a:0:32-S32"
+target triple = "i386-pc-windows-msvc"
+
+%struct.Foo = type { i32, i32 }
+
+define dso_local noundef i32 @foo() local_unnamed_addr #0 personality ptr @__CxxFrameHandler3 {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushl %ebp
+; CHECK-NEXT:    movl %esp, %ebp
+; CHECK-NEXT:    pushl %ebx
+; CHECK-NEXT:    pushl %edi
+; CHECK-NEXT:    pushl %esi
+; CHECK-NEXT:    subl $16, %esp
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $-1, -16(%ebp)
+; CHECK-NEXT:    leal -24(%ebp), %eax
+; CHECK-NEXT:    movl $___ehhandler$foo, -20(%ebp)
+; CHECK-NEXT:    movl %fs:0, %ecx
+; CHECK-NEXT:    movl %ecx, -24(%ebp)
+; CHECK-NEXT:    movl %eax, %fs:0
+; CHECK-NEXT:    pushl %eax
+; CHECK-NEXT:    pushl %eax
+; CHECK-NEXT:    movl %esp, %ecx
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $123, (%ecx)
+; CHECK-NEXT:    calll _bar
+; CHECK-NEXT:    movl $0, -16(%ebp)
+; CHECK-NEXT:    calll _alwaysthrows
+; CHECK-NEXT:  # %bb.3: # %unreachable.i
+; CHECK-NEXT:  LBB0_2: # Block address taken
+; CHECK-NEXT:    # %catch.i
+; CHECK-NEXT:    addl $12, %ebp
+; CHECK-NEXT:    jmp LBB0_4
+; CHECK-NEXT:  LBB0_4: # %exit
+; CHECK-NEXT:  $ehgcr_0_4:
+; CHECK-NEXT:    movl -24(%ebp), %eax
+; CHECK-NEXT:    movl %eax, %fs:0
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:    leal -12(%ebp), %esp
+; CHECK-NEXT:    popl %esi
+; CHECK-NEXT:    popl %edi
+; CHECK-NEXT:    popl %ebx
+; CHECK-NEXT:    popl %ebp
+; CHECK-NEXT:    retl
+; CHECK-NEXT:    .def "?catch$1@?0?foo@4HA";
+; CHECK-NEXT:    .scl 3;
+; CHECK-NEXT:    .type 32;
+; CHECK-NEXT:    .endef
+; CHECK-NEXT:    .p2align 4
+; CHECK-NEXT:  "?catch$1@?0?foo@4HA":
+; CHECK-NEXT:  LBB0_1: # %catch.i
+; CHECK-NEXT:    pushl %ebp
+; CHECK-NEXT:    addl $12, %ebp
+; CHECK-NEXT:    movl %esp, -28(%ebp)
+; CHECK-NEXT:    movl $LBB0_2, %eax
+; CHECK-NEXT:    popl %ebp
+; CHECK-NEXT:    retl # CATCHRET
+; CHECK-NEXT:  Lfunc_end0:
+entry:
+  %argmem = alloca inalloca <{ %struct.Foo }>, align 4
+  store i32 123, ptr %argmem, align 4
+  call x86_thiscallcc void @bar(ptr noundef nonnull align 4 dereferenceable(8) %argmem)
+  invoke void @alwaysthrows() #1
+          to label %unreachable.i unwind label %catch.dispatch.i
+
+catch.dispatch.i:                                 ; preds = %entry
+  %3 = catchswitch within none [label %catch.i] unwind to caller
+
+catch.i:                                          ; preds = %catch.dispatch.i
+  %4 = catchpad within %3 [ptr null, i32 64, ptr null]
+  catchret from %4 to label %exit
+
+unreachable.i:                                    ; preds = %entry
+  unreachable
+
+exit:                                             ; preds = %catch.i
+  ret i32 0
+}
+
+declare dso_local x86_thiscallcc void @bar(ptr noundef nonnull align 4 dereferenceable(8) %this) local_unnamed_addr
+
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+declare dso_local void @alwaysthrows() local_unnamed_addr
+
+attributes #0 = { norecurse "min-legal-vector-width"="0" "target-cpu"="pentium4" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #1 = { noreturn }
