[llvm] [CodeGen] Place unsplittable blocks after the profiled blocks in BB sections (PR #228173)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Rahman Lavaee (rlavaee)

<details>
<summary>Changes</summary>

With a basic block sections profile, a block which is not in the profile but is not safe to split (TargetInstrInfo::isMBBSafeToSplitToCold) stays in the default section, which is cluster 0's section. The comparator uses FuncClusterInfo.lookup(), which returns PositionInCluster 0 for such blocks, so they are placed right after the entry block, in the middle of the hot cluster.

Use find() in the comparator instead, so that blocks without cluster information are placed after all the profiled blocks of the section, in their original order. This keeps the profiled blocks of the hot cluster contiguous.

Only AArch64 overrides isMBBSafeToSplitToCold (for asm goto and jump tables), so only AArch64 is affected. On a large AArch64 workload built with a Propeller profile, this changed the layout of about 100 functions and was performance-neutral.


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


2 Files Affected:

- (modified) llvm/lib/CodeGen/BasicBlockSections.cpp (+16-6) 
- (modified) llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll (+13-7) 


``````````diff
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 9972d431aae8b..a7af7630bd729 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -284,7 +284,8 @@ assignSections(MachineFunction &MF,
 
         if (TII.isMBBSafeToSplitToCold(MBB)) {
           // BB goes into the special cold section if it is not specified in the
-          // cluster info map.
+          // cluster info map. If it is not safe to split, it is kept in the
+          // default section.
           MBB.setSectionID(MBBSectionID::ColdSectionID);
         }
       }
@@ -441,11 +442,20 @@ bool BasicBlockSections::handleBBSections(MachineFunction &MF) {
     // 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)
-      return FuncClusterInfo.lookup(*X.getBBID()).PositionInCluster <
-             FuncClusterInfo.lookup(*Y.getBBID()).PositionInCluster;
+    // If the two basic blocks are in the same section, the order is decided by
+    // their position within the section. Basic blocks which are not in any
+    // cluster (those which are not safe to split) come after all the profiled
+    // basic blocks of the section, in their original order.
+    if (XSectionID.Type == MBBSectionID::SectionType::Default) {
+      auto XI = FuncClusterInfo.find(*X.getBBID());
+      auto YI = FuncClusterInfo.find(*Y.getBBID());
+      bool XInCluster = XI != FuncClusterInfo.end();
+      bool YInCluster = YI != FuncClusterInfo.end();
+      if (XInCluster != YInCluster)
+        return XInCluster;
+      if (XInCluster)
+        return XI->second.PositionInCluster < YI->second.PositionInCluster;
+    }
     return X.getNumber() < Y.getNumber();
   };
 
diff --git a/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll b/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll
index a83a47c9c129c..e5ad5ee513fea 100644
--- a/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll
+++ b/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll
@@ -4,7 +4,7 @@
 ;; Profile for version 0.
 ; RUN: echo 'v1' > %t1
 ; RUN: echo 'f _Z3asm_goto' >> %t1
-; RUN: echo 'c 0' >> %t1
+; RUN: echo 'c 0 3' >> %t1
 ; RUN: echo 'f _Z3jump_table' >> %t1
 ; RUN: echo 'c 0' >> %t1
 ; RUN: echo 'f _Z3red_zone' >> %t1
@@ -15,14 +15,20 @@
 ; RUN: llc < %s -mtriple=aarch64 -function-sections -basic-block-sections=%t1 -unique-basic-block-section-names -bbsections-cold-text-prefix=".text.unlikely." | FileCheck %s -check-prefix=RED-ZONE
 
 define void @_Z3asm_goto(i1 zeroext %0, i1 zeroext %1) nounwind {
-  ;; Check that blocks containing or targeted by asm goto aren't split.
-  ; CHECK-LABEL:  _Z3asm_goto
+  ;; Check that blocks containing or targeted by asm goto aren't split. They are
+  ;; placed after all the blocks of the hot cluster.
+  ; CHECK-LABEL:  _Z3asm_goto:
+  ; CHECK-NOT:    .section
+  ; CHECK:          bl bam
+  ; CHECK-NOT:    .section
+  ; CHECK:          bl bar
+  ; CHECK:          nop
+  ; CHECK:        // Inline asm indirect target
+  ; CHECK:          bl bar
   ; CHECK:        .section	.text.unlikely._Z3asm_goto,"ax", at progbits
   ; CHECK-NEXT:     _Z3asm_goto.cold:
-  ; CHECK-NEXT:       bl bam
-  ; CHECK:          .LBB0_4:
-  ; CHECK:            ret
-  ; CHECK:          .LBB_END0_4:
+  ; CHECK-NEXT:       ldp
+  ; CHECK-NEXT:       ret
 
   br i1 %0, label %3, label %5
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/228173