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

[Fabian Parzefall via llvm-commits]

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

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 adds a field to machine blocks indicating whether this block contains the first instruction in its section. This variable is then used to determine whether to emit the padding.

>From d9c7203d04d9148510e2c9fe13182bc89d9742ee 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 adds a field to machine blocks indicating
whether this block contains the first instruction in its section. This
variable is then used to determine whether to emit the padding.
---
 llvm/include/llvm/CodeGen/MachineBasicBlock.h | 11 +++++++
 llvm/lib/CodeGen/BasicBlockSections.cpp       |  2 +-
 llvm/lib/CodeGen/MachineFunction.cpp          | 17 ++++++++++-
 .../Generic/machine-function-splitter.ll      | 30 +++++++++++++++++++
 4 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
index 6cf151c951b19f..d1ae1b037493bd 100644
--- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h
+++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -224,6 +224,10 @@ class MachineBasicBlock
   // Indicate that this basic block ends a section.
   bool IsEndSection = false;
 
+  // Indicate that this basic block contains the first instruction in the
+  // section.
+  bool IsFirstNonEmptyBBInSection = false;
+
   /// Indicate that this basic block is the indirect dest of an INLINEASM_BR.
   bool IsInlineAsmBrIndirectTarget = false;
 
@@ -673,10 +677,17 @@ class MachineBasicBlock
   /// Returns true if this block ends any section.
   bool isEndSection() const { return IsEndSection; }
 
+  /// Returns true if this block contains the first instruction of its section.
+  bool isFirstNonEmptyBBInSection() const { return IsFirstNonEmptyBBInSection; }
+
   void setIsBeginSection(bool V = true) { IsBeginSection = V; }
 
   void setIsEndSection(bool V = true) { IsEndSection = V; }
 
+  void setIsFirstNonEmptyBBInSection(bool V = true) {
+    IsFirstNonEmptyBBInSection = V;
+  }
+
   std::optional<UniqueBBID> getBBID() const { return BBID; }
 
   /// Returns the section ID of this basic block.
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 1eedfc4b259126..d04d05aa336de9 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -266,7 +266,7 @@ void llvm::sortBasicBlocksAndUpdateBranches(
 // pad label to ensure a nonzero offset.
 void llvm::avoidZeroOffsetLandingPad(MachineFunction &MF) {
   for (auto &MBB : MF) {
-    if (MBB.isBeginSection() && MBB.isEHPad()) {
+    if (MBB.isFirstNonEmptyBBInSection() && MBB.isEHPad()) {
       MachineBasicBlock::iterator MI = MBB.begin();
       while (!MI->isEHLabel())
         ++MI;
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
index b56888a0f71fe6..3dbd59fe1b0a14 100644
--- a/llvm/lib/CodeGen/MachineFunction.cpp
+++ b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -383,13 +383,28 @@ void MachineFunction::RenumberBlocks(MachineBasicBlock *MBB) {
 /// and the SectionID's are assigned to MBBs.
 void MachineFunction::assignBeginEndSections() {
   front().setIsBeginSection();
+  front().setIsFirstNonEmptyBBInSection();
   auto CurrentSectionID = front().getSectionID();
+  bool FirstSectionFirstInstructionEmitted = true;
   for (auto MBBI = std::next(begin()), E = end(); MBBI != E; ++MBBI) {
-    if (MBBI->getSectionID() == CurrentSectionID)
+    if (MBBI->getSectionID() == CurrentSectionID) {
+      if (!FirstSectionFirstInstructionEmitted && !MBBI->empty()) {
+        MBBI->setIsFirstNonEmptyBBInSection();
+        FirstSectionFirstInstructionEmitted = true;
+      }
       continue;
+    }
+
     MBBI->setIsBeginSection();
     std::prev(MBBI)->setIsEndSection();
     CurrentSectionID = MBBI->getSectionID();
+
+    if (MBBI->empty()) {
+      FirstSectionFirstInstructionEmitted = false;
+    } else {
+      MBBI->setIsFirstNonEmptyBBInSection();
+      FirstSectionFirstInstructionEmitted = true;
+    }
   }
   back().setIsEndSection();
 }
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()