From 7bfd62ed044dd7b8f69215b97f9c18d69b2f8116 Mon Sep 17 00:00:00 2001
From: Fabian Parzefall <parzefall@meta.com>
Date: Wed, 9 Oct 2024 19:03:44 -0700
Subject: [PATCH] [CodeGen] Fix lpad padding at section start after empty block

If a landing pad is at the very start of a split section, it has to be
padded by a nop instruction. Otherwise its offset is marked as zero in
the LSDA, which means no landing pad (leading it to be skipped).

LLVM already handles this. If a landing pad is the first machine block
in a section, a nop is inserted to ensure a non-zero offset. However, if
the landing pad is preceeded by an empty block, the nop would be
omitted.

To fix this, this patch checks not whether a block is the first in
section, but instead whether it is the first *non-empty* block in the
section.
---
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 10 ++++++-
 .../Generic/machine-function-splitter.ll      | 30 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 1eedfc4b25912..a7ffc2e03b838 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -265,8 +265,16 @@ void llvm::sortBasicBlocksAndUpdateBranches(
 // zero implies "no landing pad." This function inserts a NOP just before the EH
 // pad label to ensure a nonzero offset.
 void llvm::avoidZeroOffsetLandingPad(MachineFunction &MF) {
+  std::optional<MBBSectionID> CurrentSection;
+  auto IsFirstNonEmptyBBInSection = [&](const MachineBasicBlock &MBB) {
+    if (MBB.empty() || MBB.getSectionID() == CurrentSection)
+      return false;
+    CurrentSection = MBB.getSectionID();
+    return true;
+  };
+
   for (auto &MBB : MF) {
-    if (MBB.isBeginSection() && MBB.isEHPad()) {
+    if (IsFirstNonEmptyBBInSection(MBB) && MBB.isEHPad()) {
       MachineBasicBlock::iterator MI = MBB.begin();
       while (!MI->isEHLabel())
         ++MI;
diff --git a/llvm/test/CodeGen/Generic/machine-function-splitter.ll b/llvm/test/CodeGen/Generic/machine-function-splitter.ll
index 1a8c9ede8f8b7..d798b2875645b 100644
--- a/llvm/test/CodeGen/Generic/machine-function-splitter.ll
+++ b/llvm/test/CodeGen/Generic/machine-function-splitter.ll
@@ -674,6 +674,36 @@ define void @foo22(i1 zeroext %0) nounwind !prof !14 !section_prefix !15 {
   ret void
 }
 
+define i32 @foo23(i1 zeroext %0) personality ptr @__gxx_personality_v0 !prof !14 {
+;; Check that nop is inserted just before the EH pad if it is the first
+;; instruction in a section (but is preceeded by another empty block).
+; MFS-DEFAULTS-LABEL:         foo23
+; MFS-DEFAULTS-X86-LABEL:     callq   baz
+; MFS-DEFAULTS-X86:           .section        .text.split.foo23,"ax",@progbits
+; MFS-DEFAULTS-X86-NEXT:      foo23.cold:
+; MFS-DEFAULTS-X86:           nop
+; MFS-DEFAULTS-X86:           callq   _Unwind_Resume@PLT
+entry:
+  br i1 %0, label %try, label %unreachable, !prof !17
+
+try:
+  invoke void @_Z1fv()
+          to label %try.cont unwind label %lpad, !prof !17
+
+try.cont:
+  %1 = call i32 @baz()
+  ret i32 %1
+
+unreachable:
+  unreachable
+
+lpad:
+  %2 = landingpad { ptr, i32 }
+          cleanup
+          catch ptr @_ZTIi
+  resume { ptr, i32 } %2
+}
+
 declare i32 @bar()
 declare i32 @baz()
 declare i32 @bam()