[BasicBlockSections] Always keep the entry block in the beginning of the function. (#74696)

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

Author: rlavaee

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;
 };
 

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).
@@ -341,12 +340,17 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   // contiguous and ordered accordingly. Furthermore, clusters are ordered in
   // increasing order of their section IDs, with the exception and the
   // cold section placed at the end of the function.
+  // Also, we force the entry block of the function to be placed at the
+  // beginning of the function, regardless of the requested order.
   auto Comparator = [&](const MachineBasicBlock &X,
                         const MachineBasicBlock &Y) {
     auto XSectionID = X.getSectionID();
     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)

llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp

@@ -214,10 +214,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++});
       }
@@ -288,9 +284,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},

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

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