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

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