[llvm] [CodeGen] Fix lpad padding at section start after empty block (PR #112595)

[Fabian Parzefall via llvm-commits]

https://github.com/pzfl updated https://github.com/llvm/llvm-project/pull/112595

>From 2ea993c82c6b01eacd109670abba64948d3ef9b9 Mon Sep 17 00:00:00 2001
From: Fabian Parzefall <parzefall at 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       | 16 +++++++++-
 .../Generic/machine-function-splitter.ll      | 30 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 1eedfc4b259126..b39c43e93cf749 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -265,8 +265,22 @@ 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;
+  bool SectionEmpty = true;
+  auto IsFirstNonEmptyBBInSection = [&](const MachineBasicBlock &MBB) {
+    if (MBB.getSectionID() != CurrentSection) {
+      CurrentSection = MBB.getSectionID();
+      SectionEmpty = true;
+    }
+    if (SectionEmpty && !MBB.empty()) {
+      SectionEmpty = false;
+      return true;
+    }
+    return false;
+  };
+
   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 2097523a61c5f9..5028dbc53ce65c 100644
--- a/llvm/test/CodeGen/Generic/machine-function-splitter.ll
+++ b/llvm/test/CodeGen/Generic/machine-function-splitter.ll
@@ -610,6 +610,36 @@ cold_asm_target:
   ret void
 }
 
+define i32 @foo21(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:         foo21
+; MFS-DEFAULTS-X86-LABEL:     callq   baz
+; MFS-DEFAULTS-X86:           .section        .text.split.foo21,"ax", at progbits
+; MFS-DEFAULTS-X86-NEXT:      foo21.cold:
+; MFS-DEFAULTS-X86:           nop
+; MFS-DEFAULTS-X86:           callq   _Unwind_Resume at 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()