[BasicBlockSections] Always keep the entry block in the beginning of the function.

[rlavaee]

BasicBlockSections must enforce placing the entry block at the beginning of the function regardless of the basic block sections profile.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Rahman Lavaee (rlavaee)

Changes

BasicBlockSections must enforce placing the entry block at the beginning of the function regardless of the basic block sections profile.

---

Full diff: https://github.com/llvm/llvm-project/pull/74696.diff

5 Files Affected:

• (modified) llvm/include/llvm/CodeGen/MachineBasicBlock.h (+2-2)
• (modified) llvm/lib/CodeGen/BasicBlockSections.cpp (+5-3)
• (modified) llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp (-7)
• (modified) llvm/test/CodeGen/X86/basic-block-sections-clusters-error.ll (+1-5)
• (added) llvm/test/CodeGen/X86/basic-block-sections-entryblock.ll (+43)

diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
index 5812295f73b5a..c84fd281c6a54 100644
--- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h
+++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -74,10 +74,10 @@ struct MBBSectionID {
   MBBSectionID(SectionType T) : Type(T), Number(0) {}
 };
 
-// This structure represents the information for a basic block.
+// This structure represents the information for a basic block pertaining to
+// the basic block sections profile.
 struct UniqueBBID {
   unsigned BaseID;
-  // sections profile).
   unsigned CloneID;
 };
 
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 42997d2287d61..4b7a2698afc8c 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -318,9 +318,8 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   MF.setBBSectionsType(BBSectionsType);
   assignSections(MF, FuncClusterInfo);
 
-  // We make sure that the cluster including the entry basic block precedes all
-  // other clusters.
-  auto EntryBBSectionID = MF.front().getSectionID();
+  const MachineBasicBlock &EntryBB = MF.front();
+  auto EntryBBSectionID = EntryBB.getSectionID();
 
   // Helper function for ordering BB sections as follows:
   //   * Entry section (section including the entry block).
@@ -347,6 +346,9 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
     auto YSectionID = Y.getSectionID();
     if (XSectionID != YSectionID)
       return MBBSectionOrder(XSectionID, YSectionID);
+    // Make sure that the entry block is placed at the beginning.
+    if (&X == &EntryBB || &Y == &EntryBB)
+      return &X == &EntryBB;
     // If the two basic block are in the same section, the order is decided by
     // their position within the section.
     if (XSectionID.Type == MBBSectionID::SectionType::Default)
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index 96662378a8693..c544171e2abb1 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -213,10 +213,6 @@ Error BasicBlockSectionsProfileReader::ReadV1Profile() {
               Twine("duplicate basic block id found '") + BasicBlockIDStr +
               "'");
 
-        if (!BasicBlockID->BaseID && CurrentPosition)
-          return createProfileParseError(
-              "entry BB (0) does not begin a cluster.");
-
         FI->second.ClusterInfo.emplace_back(BBClusterInfo{
             *std::move(BasicBlockID), CurrentCluster, CurrentPosition++});
       }
@@ -287,9 +283,6 @@ Error BasicBlockSectionsProfileReader::ReadV0Profile() {
         if (!FuncBBIDs.insert(BBID).second)
           return createProfileParseError(
               Twine("duplicate basic block id found '") + BBIDStr + "'");
-        if (BBID == 0 && CurrentPosition)
-          return createProfileParseError(
-              "entry BB (0) does not begin a cluster");
 
         FI->second.ClusterInfo.emplace_back(
             BBClusterInfo({{static_cast<unsigned>(BBID), 0},
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-clusters-error.ll b/llvm/test/CodeGen/X86/basic-block-sections-clusters-error.ll
index 597d8f6707ecc..d6f3d5010b556 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-clusters-error.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-clusters-error.ll
@@ -5,10 +5,6 @@
 ; RUN: echo '!!1' >> %t1
 ; RUN: not --crash llc < %s -O0 -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t1 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR1
 ; CHECK-ERROR1: LLVM ERROR: invalid profile {{.*}} at line 3: duplicate basic block id found '1'
-; RUN: echo '!dummy1' > %t2
-; RUN: echo '!!4 0' >> %t2
-; RUN: not --crash llc < %s -O0 -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t2 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR2
-; CHECK-ERROR2: LLVM ERROR: invalid profile {{.*}} at line 2: entry BB (0) does not begin a cluster
 ; RUN: echo '!dummy1' > %t3
 ; RUN: echo '!!-1' >> %t3
 ; RUN: not --crash llc < %s -O0 -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t3 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR3
@@ -48,7 +44,7 @@
 ; RUN: echo 'f dummy1' >> %t11
 ; RUN: echo 'c 0 1.a' >> %t11
 ; RUN: not --crash llc < %s -O0 -mtriple=x86_64 -function-sections -basic-block-sections=%t11 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR11
-; CHECK-ERROR11: LLVM ERROR: invalid profile {{.*}} at line 3: unable to parse clone id: 'a' 
+; CHECK-ERROR11: LLVM ERROR: invalid profile {{.*}} at line 3: unable to parse clone id: 'a'
 ; RUN: echo 'v1' > %t12
 ; RUN: echo 'f dummy1' >> %t12
 ; RUN: echo 'c 0 1' >> %t12
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-entryblock.ll b/llvm/test/CodeGen/X86/basic-block-sections-entryblock.ll
new file mode 100644
index 0000000000000..349015e1403df
--- /dev/null
+++ b/llvm/test/CodeGen/X86/basic-block-sections-entryblock.ll
@@ -0,0 +1,43 @@
+; COM: Tests to verify that the entry basic block is always placed at the beginning of its section.
+; RUN: echo 'v1' > %t1
+; RUN: echo 'f foo' >> %t1
+; RUN: echo 'c2 0' >> %t1
+; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t1 -O0 | FileCheck %s -check-prefix=LINUX-SECTIONS1
+
+; RUN: echo 'v1' > %t2
+; RUN: echo 'f foo' >> %t2
+; RUN: echo 'c2' >> %t2
+; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t2 -O0 | FileCheck %s -check-prefix=LINUX-SECTIONS2
+
+
+define void @foo(i1 %a, i1 %b) {
+b0:
+  br i1 %a, label %b1, label %b2
+
+b1:                                           ; preds = %b0
+  ret void
+
+b2:                                           ; preds = %b0
+  ret void
+}
+
+;; Check that %b0 is emitted at the beginning of the function.
+; LINUX-SECTIONS1:    .section .text.foo,"ax",@progbits
+; LINUX-SECTIONS1:  foo:
+; LINUX-SECTIONS1:  # %bb.0:             # %b0
+; LINUX-SECTIONS1:    jne foo.cold
+; LINUX-SECTIONS1:  # %bb.2:             # %b2
+; LINUX-SECTIONS1:    retq
+; LINUX-SECTIONS1:    .section .text.split.foo,"ax",@progbits
+; LINUX-SECTIONS1:  foo.cold:            # %b1
+; LINUX-SECTIONS1:    retq
+
+; LINUX-SECTIONS2:    .section .text.foo,"ax",@progbits
+; LINUX-SECTIONS2:  foo:
+; LINUX-SECTIONS2:  # %bb.0:             # %b0
+; LINUX-SECTIONS2:    je foo.__part.0
+; LINUX-SECTIONS2:  # %bb.1:             # %b1
+; LINUX-SECTIONS2:    retq
+; LINUX-SECTIONS2:    .section .text.foo,"ax",@progbits
+; LINUX-SECTIONS2:  foo.__part.0:        # %b2
+; LINUX-SECTIONS2:    retq

[tmsri]

This test is pretty large, please split this into multiple tests, can be done in a separate patch.

[rlavaee]

Sure. Will do this in a separate PR.