[llvm] Win x64 Unwind V2 1/n: Mark beginning and end of epilogs

llvm/include/llvm/MC/MCStreamer.h

@@ -252,6 +252,12 @@ class MCStreamer {
   bool AllowAutoPadding = false;
 
 protected:
+  // True if we are processing SEH directives in an epilogue.
+  bool InEpilogCFI = false;
+
+  // Symbol of the current epilog for which we are processing SEH directives.
+  MCSymbol *CurrentEpilog = nullptr;
+
   MCFragment *CurFrag = nullptr;
 
   MCStreamer(MCContext &Ctx);
@@ -333,6 +339,10 @@ class MCStreamer {
     return WinFrameInfos;
   }
 
+  MCSymbol *getCurrentEpilog() const { return CurrentEpilog; }
+
+  bool isInEpilogCFI() const { return InEpilogCFI; }
+
   void generateCompactUnwindEncodings(MCAsmBackend *MAB);
 
   /// \name Assembly File Formatting.
@@ -1056,6 +1066,8 @@ class MCStreamer {
                                  SMLoc Loc = SMLoc());
   virtual void emitWinCFIPushFrame(bool Code, SMLoc Loc = SMLoc());
   virtual void emitWinCFIEndProlog(SMLoc Loc = SMLoc());
+  virtual void emitWinCFIBeginEpilogue(SMLoc Loc = SMLoc());
+  virtual void emitWinCFIEndEpilogue(SMLoc Loc = SMLoc());
   virtual void emitWinEHHandler(const MCSymbol *Sym, bool Unwind, bool Except,
                                 SMLoc Loc = SMLoc());
   virtual void emitWinEHHandlerData(SMLoc Loc = SMLoc());

llvm/lib/MC/MCAsmStreamer.cpp

@@ -398,6 +398,8 @@ class MCAsmStreamer final : public MCStreamer {
                          SMLoc Loc) override;
   void emitWinCFIPushFrame(bool Code, SMLoc Loc) override;
   void emitWinCFIEndProlog(SMLoc Loc) override;
+  void emitWinCFIBeginEpilogue(SMLoc Loc) override;
+  void emitWinCFIEndEpilogue(SMLoc Loc) override;
 
   void emitWinEHHandler(const MCSymbol *Sym, bool Unwind, bool Except,
                         SMLoc Loc) override;
@@ -2348,6 +2350,20 @@ void MCAsmStreamer::emitWinCFIEndProlog(SMLoc Loc) {
   EmitEOL();
 }
 
+void MCAsmStreamer::emitWinCFIBeginEpilogue(SMLoc Loc) {
+  MCStreamer::emitWinCFIBeginEpilogue(Loc);
+
+  OS << "\t.seh_startepilogue";
+  EmitEOL();
+}
+
+void MCAsmStreamer::emitWinCFIEndEpilogue(SMLoc Loc) {
+  MCStreamer::emitWinCFIEndEpilogue(Loc);
+
+  OS << "\t.seh_endepilogue";
+  EmitEOL();
+}
+
 void MCAsmStreamer::emitCGProfileEntry(const MCSymbolRefExpr *From,
                                        const MCSymbolRefExpr *To,
                                        uint64_t Count) {

llvm/lib/MC/MCParser/COFFAsmParser.cpp

@@ -92,6 +92,10 @@ class COFFAsmParser : public MCAsmParserExtension {
         ".seh_stackalloc");
     addDirectiveHandler<&COFFAsmParser::parseSEHDirectiveEndProlog>(
         ".seh_endprologue");
+    addDirectiveHandler<&COFFAsmParser::ParseSEHDirectiveBeginEpilog>(
+        ".seh_startepilogue");
+    addDirectiveHandler<&COFFAsmParser::ParseSEHDirectiveEndEpilog>(
+        ".seh_endepilogue");
   }
 
   bool parseSectionDirectiveText(StringRef, SMLoc) {
@@ -141,6 +145,8 @@ class COFFAsmParser : public MCAsmParserExtension {
   bool parseSEHDirectiveHandlerData(StringRef, SMLoc);
   bool parseSEHDirectiveAllocStack(StringRef, SMLoc);
   bool parseSEHDirectiveEndProlog(StringRef, SMLoc);
+  bool ParseSEHDirectiveBeginEpilog(StringRef, SMLoc);
+  bool ParseSEHDirectiveEndEpilog(StringRef, SMLoc);
 
   bool parseAtUnwindOrAtExcept(bool &unwind, bool &except);
   bool parseDirectiveSymbolAttribute(StringRef Directive, SMLoc);
@@ -749,6 +755,18 @@ bool COFFAsmParser::parseSEHDirectiveEndProlog(StringRef, SMLoc Loc) {
   return false;
 }
 
+bool COFFAsmParser::ParseSEHDirectiveBeginEpilog(StringRef, SMLoc Loc) {
+  Lex();
+  getStreamer().emitWinCFIBeginEpilogue(Loc);
+  return false;
+}
+
+bool COFFAsmParser::ParseSEHDirectiveEndEpilog(StringRef, SMLoc Loc) {
+  Lex();
+  getStreamer().emitWinCFIEndEpilogue(Loc);
+  return false;
+}
+
 bool COFFAsmParser::parseAtUnwindOrAtExcept(bool &unwind, bool &except) {
   StringRef identifier;
   if (getLexer().isNot(AsmToken::At) && getLexer().isNot(AsmToken::Percent))

llvm/lib/MC/MCStreamer.cpp

@@ -1013,6 +1013,36 @@ void MCStreamer::emitWinCFIEndProlog(SMLoc Loc) {
   CurFrame->PrologEnd = Label;
 }
 
+void MCStreamer::emitWinCFIBeginEpilogue(SMLoc Loc) {
+  WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
+  if (!CurFrame)
+    return;
+
+  if (!CurFrame->PrologEnd)
+    return getContext().reportError(
+        Loc, "starting epilogue (.seh_startepilogue) before prologue has ended "
+             "(.seh_endprologue) in " +
+                 CurFrame->Function->getName());
+
+  InEpilogCFI = true;
+  CurrentEpilog = emitCFILabel();
+}
+
+void MCStreamer::emitWinCFIEndEpilogue(SMLoc Loc) {
+  WinEH::FrameInfo *CurFrame = EnsureValidWinFrameInfo(Loc);
+  if (!CurFrame)
+    return;
+
+  if (!InEpilogCFI)
+    return getContext().reportError(Loc, "Stray .seh_endepilogue in " +
+                                             CurFrame->Function->getName());
+
+  InEpilogCFI = false;
+  MCSymbol *Label = emitCFILabel();
+  CurFrame->EpilogMap[CurrentEpilog].End = Label;
+  CurrentEpilog = nullptr;
+}
+
 void MCStreamer::emitCOFFSafeSEH(MCSymbol const *Symbol) {}
 
 void MCStreamer::emitCOFFSymbolIndex(MCSymbol const *Symbol) {}

llvm/lib/Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h

@@ -138,12 +138,6 @@ class AArch64TargetELFStreamer : public AArch64TargetStreamer {
 };
 
 class AArch64TargetWinCOFFStreamer : public llvm::AArch64TargetStreamer {
-private:
-  // True if we are processing SEH directives in an epilogue.
-  bool InEpilogCFI = false;
-
-  // Symbol of the current epilog for which we are processing SEH directives.
-  MCSymbol *CurrentEpilog = nullptr;
 public:
   AArch64TargetWinCOFFStreamer(llvm::MCStreamer &S)
     : AArch64TargetStreamer(S) {}

llvm/lib/Target/AArch64/MCTargetDesc/AArch64WinCOFFStreamer.cpp

@@ -73,8 +73,8 @@ void AArch64TargetWinCOFFStreamer::emitARM64WinUnwindCode(unsigned UnwindCode,
   if (!CurFrame)
     return;
   auto Inst = WinEH::Instruction(UnwindCode, /*Label=*/nullptr, Reg, Offset);
-  if (InEpilogCFI)
-    CurFrame->EpilogMap[CurrentEpilog].Instructions.push_back(Inst);
+  if (S.isInEpilogCFI())
+    CurFrame->EpilogMap[S.getCurrentEpilog()].Instructions.push_back(Inst);
   else
     CurFrame->Instructions.push_back(Inst);
 }
@@ -183,13 +183,7 @@ void AArch64TargetWinCOFFStreamer::emitARM64WinCFIPrologEnd() {
 }
 
 void AArch64TargetWinCOFFStreamer::emitARM64WinCFIEpilogStart() {
-  auto &S = getStreamer();
-  WinEH::FrameInfo *CurFrame = S.EnsureValidWinFrameInfo(SMLoc());
-  if (!CurFrame)
-    return;
-
-  InEpilogCFI = true;
-  CurrentEpilog = S.emitCFILabel();
+  getStreamer().emitWinCFIBeginEpilogue();
 }
 
 void AArch64TargetWinCOFFStreamer::emitARM64WinCFIEpilogEnd() {
@@ -198,13 +192,12 @@ void AArch64TargetWinCOFFStreamer::emitARM64WinCFIEpilogEnd() {
   if (!CurFrame)
     return;
 
-  InEpilogCFI = false;
-  WinEH::Instruction Inst =
-      WinEH::Instruction(Win64EH::UOP_End, /*Label=*/nullptr, -1, 0);
-  CurFrame->EpilogMap[CurrentEpilog].Instructions.push_back(Inst);
-  MCSymbol *Label = S.emitCFILabel();
-  CurFrame->EpilogMap[CurrentEpilog].End = Label;
-  CurrentEpilog = nullptr;
+  if (S.isInEpilogCFI()) {
+    WinEH::Instruction Inst =
+        WinEH::Instruction(Win64EH::UOP_End, /*Label=*/nullptr, -1, 0);
+    CurFrame->EpilogMap[S.getCurrentEpilog()].Instructions.push_back(Inst);
+  }
+  S.emitWinCFIEndEpilogue();
 }
 
 void AArch64TargetWinCOFFStreamer::emitARM64WinCFITrapFrame() {

llvm/lib/Target/ARM/MCTargetDesc/ARMWinCOFFStreamer.cpp

@@ -77,13 +77,6 @@ llvm::createARMWinCOFFStreamer(MCContext &Context,
 
 namespace {
 class ARMTargetWinCOFFStreamer : public llvm::ARMTargetStreamer {
-private:
-  // True if we are processing SEH directives in an epilogue.
-  bool InEpilogCFI = false;
-
-  // Symbol of the current epilog for which we are processing SEH directives.
-  MCSymbol *CurrentEpilog = nullptr;
-
 public:
   ARMTargetWinCOFFStreamer(llvm::MCStreamer &S) : ARMTargetStreamer(S) {}
 
@@ -114,8 +107,8 @@ void ARMTargetWinCOFFStreamer::emitARMWinUnwindCode(unsigned UnwindCode,
     return;
   MCSymbol *Label = S.emitCFILabel();
   auto Inst = WinEH::Instruction(UnwindCode, Label, Reg, Offset);
-  if (InEpilogCFI)
-    CurFrame->EpilogMap[CurrentEpilog].Instructions.push_back(Inst);
+  if (S.isInEpilogCFI())
+    CurFrame->EpilogMap[S.getCurrentEpilog()].Instructions.push_back(Inst);
   else
     CurFrame->Instructions.push_back(Inst);
 }
@@ -224,9 +217,10 @@ void ARMTargetWinCOFFStreamer::emitARMWinCFIEpilogStart(unsigned Condition) {
   if (!CurFrame)
     return;
 
-  InEpilogCFI = true;
-  CurrentEpilog = S.emitCFILabel();
-  CurFrame->EpilogMap[CurrentEpilog].Condition = Condition;
+  S.emitWinCFIBeginEpilogue();
+  if (S.isInEpilogCFI()) {
+    CurFrame->EpilogMap[S.getCurrentEpilog()].Condition = Condition;
+  }
 }
 
 void ARMTargetWinCOFFStreamer::emitARMWinCFIEpilogEnd() {
@@ -235,33 +229,26 @@ void ARMTargetWinCOFFStreamer::emitARMWinCFIEpilogEnd() {
   if (!CurFrame)
     return;
 
-  if (!CurrentEpilog) {
-    S.getContext().reportError(SMLoc(), "Stray .seh_endepilogue in " +
-                                            CurFrame->Function->getName());
-    return;
-  }
-
-  std::vector<WinEH::Instruction> &Epilog =
-      CurFrame->EpilogMap[CurrentEpilog].Instructions;
-
-  unsigned UnwindCode = Win64EH::UOP_End;
-  if (!Epilog.empty()) {
-    WinEH::Instruction EndInstr = Epilog.back();
-    if (EndInstr.Operation == Win64EH::UOP_Nop) {
-      UnwindCode = Win64EH::UOP_EndNop;
-      Epilog.pop_back();
-    } else if (EndInstr.Operation == Win64EH::UOP_WideNop) {
-      UnwindCode = Win64EH::UOP_WideEndNop;
-      Epilog.pop_back();
+  if (S.isInEpilogCFI()) {
+    std::vector<WinEH::Instruction> &Epilog =
+        CurFrame->EpilogMap[S.getCurrentEpilog()].Instructions;
+
+    unsigned UnwindCode = Win64EH::UOP_End;
+    if (!Epilog.empty()) {
+      WinEH::Instruction EndInstr = Epilog.back();
+      if (EndInstr.Operation == Win64EH::UOP_Nop) {
+        UnwindCode = Win64EH::UOP_EndNop;
+        Epilog.pop_back();
+      } else if (EndInstr.Operation == Win64EH::UOP_WideNop) {
+        UnwindCode = Win64EH::UOP_WideEndNop;
+        Epilog.pop_back();
+      }
     }
-  }
 
-  InEpilogCFI = false;
-  WinEH::Instruction Inst = WinEH::Instruction(UnwindCode, nullptr, -1, 0);
-  CurFrame->EpilogMap[CurrentEpilog].Instructions.push_back(Inst);
-  MCSymbol *Label = S.emitCFILabel();
-  CurFrame->EpilogMap[CurrentEpilog].End = Label;
-  CurrentEpilog = nullptr;
+    WinEH::Instruction Inst = WinEH::Instruction(UnwindCode, nullptr, -1, 0);
+    CurFrame->EpilogMap[S.getCurrentEpilog()].Instructions.push_back(Inst);
+  }
+  S.emitWinCFIEndEpilogue();
 }
 
 void ARMTargetWinCOFFStreamer::emitARMWinCFICustom(unsigned Opcode) {

llvm/lib/Target/X86/X86FrameLowering.cpp

@@ -2573,14 +2573,8 @@ void X86FrameLowering::emitEpilogue(MachineFunction &MF,
     --MBBI;
   }
 
-  // Windows unwinder will not invoke function's exception handler if IP is
-  // either in prologue or in epilogue.  This behavior causes a problem when a
-  // call immediately precedes an epilogue, because the return address points
-  // into the epilogue.  To cope with that, we insert an epilogue marker here,
-  // then replace it with a 'nop' if it ends up immediately after a CALL in the
-  // final emitted code.
   if (NeedsWin64CFI && MF.hasWinCFI())
-    BuildMI(MBB, MBBI, DL, TII.get(X86::SEH_Epilogue));
+    BuildMI(MBB, MBBI, DL, TII.get(X86::SEH_BeginEpilogue));
 
   if (!HasFP && NeedsDwarfCFI) {
     MBBI = FirstCSPop;
@@ -2625,6 +2619,9 @@ void X86FrameLowering::emitEpilogue(MachineFunction &MF,
   // Emit tilerelease for AMX kernel.
   if (X86FI->getAMXProgModel() == AMXProgModelEnum::ManagedRA)
     BuildMI(MBB, Terminator, DL, TII.get(X86::TILERELEASE));
+
+  if (NeedsWin64CFI && MF.hasWinCFI())
+    BuildMI(MBB, Terminator, DL, TII.get(X86::SEH_EndEpilogue));
 }
 
 StackOffset X86FrameLowering::getFrameIndexReference(const MachineFunction &MF,

llvm/lib/Target/X86/X86InstrCompiler.td

@@ -262,8 +262,10 @@ let isPseudo = 1, isMeta = 1, isNotDuplicable = 1, SchedRW = [WriteSystem] in {
 
 // Epilog instructions:
 let isPseudo = 1, isMeta = 1, SchedRW = [WriteSystem] in {
-  def SEH_Epilogue : I<0, Pseudo, (outs), (ins),
-                            "#SEH_Epilogue", []>;
+  def SEH_BeginEpilogue : I<0, Pseudo, (outs), (ins),
+                            "#SEH_BeginEpilogue", []>;
+  def SEH_EndEpilogue : I<0, Pseudo, (outs), (ins),
+                            "#SEH_EndEpilogue", []>;
 }
 
 //===----------------------------------------------------------------------===//

llvm/lib/Target/X86/X86MCInstLower.cpp

@@ -1779,6 +1779,14 @@ void X86AsmPrinter::EmitSEHInstruction(const MachineInstr *MI) {
     OutStreamer->emitWinCFIEndProlog();
     break;
 
+  case X86::SEH_BeginEpilogue:
+    OutStreamer->emitWinCFIBeginEpilogue();
+    break;
+
+  case X86::SEH_EndEpilogue:
+    OutStreamer->emitWinCFIEndEpilogue();
+    break;
+
   default:
     llvm_unreachable("expected SEH_ instruction");
   }
@@ -2420,11 +2428,17 @@ void X86AsmPrinter::emitInstruction(const MachineInstr *MI) {
   case X86::SEH_SetFrame:
   case X86::SEH_PushFrame:
   case X86::SEH_EndPrologue:
+  case X86::SEH_EndEpilogue:
     EmitSEHInstruction(MI);
     return;
 
-  case X86::SEH_Epilogue: {
+  case X86::SEH_BeginEpilogue: {
     assert(MF->hasWinCFI() && "SEH_ instruction in function without WinCFI?");
+    // Windows unwinder will not invoke function's exception handler if IP is
+    // either in prologue or in epilogue.  This behavior causes a problem when a
+    // call immediately precedes an epilogue, because the return address points
+    // into the epilogue.  To cope with that, we insert a 'nop' if it ends up
+    // immediately after a CALL in the final emitted code.
     MachineBasicBlock::const_iterator MBBI(MI);
     // Check if preceded by a call and emit nop if so.
     for (MBBI = PrevCrossBBInst(MBBI);
@@ -2439,6 +2453,8 @@ void X86AsmPrinter::emitInstruction(const MachineInstr *MI) {
         break;
       }
     }
+
+    EmitSEHInstruction(MI);
     return;
   }
   case X86::UBSAN_UD1: