[CodeGen][AArch64] Only split safe blocks in BBSections (#81553)

Some types of machine function and machine basic block are unsafe to
split on AArch64: basic blocks that contain jump table dispatch or
targets (D157124), and blocks that contain inline ASM GOTO blocks
or their targets (D158647) all cause issues and have been excluded
from Machine Function Splitting on AArch64.

These issues are caused by any transformation pass that places
same-function basic blocks in different text sections
(MachineFunctionSplitter and BasicBlockSections) and must be
special-cased in both passes.

3 files changed, 180 insertions, 3 deletions

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index eb3f9e7078f1..09e45ea5794b 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -208,9 +208,14 @@ assignSections(MachineFunction &MF,
       if (I != FuncClusterInfo.end()) {
         MBB.setSectionID(I->second.ClusterID);
       } else {
-        // BB goes into the special cold section if it is not specified in the
-        // cluster info map.
-        MBB.setSectionID(MBBSectionID::ColdSectionID);
+        const TargetInstrInfo &TII =
+            *MBB.getParent()->getSubtarget().getInstrInfo();
+
+        if (TII.isMBBSafeToSplitToCold(MBB)) {
+          // BB goes into the special cold section if it is not specified in the
+          // cluster info map.
+          MBB.setSectionID(MBBSectionID::ColdSectionID);
+        }
       }
     }
 
diff --git a/llvm/test/CodeGen/AArch64/basic-block-sections-cold.ll b/llvm/test/CodeGen/AArch64/basic-block-sections-cold.ll
new file mode 100644
index 000000000000..6641ef6a51c1
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/basic-block-sections-cold.ll
@@ -0,0 +1,51 @@
+;; Check if basic blocks that don't get unique sections are placed in cold sections.
+;; Basic block with id 1 and 2 must be in the cold section.
+;;
+;; Profile for version 0
+; RUN: echo '!_Z3bazb' > %t1
+; RUN: echo '!!0' >> %t1
+;;
+;; Profile for version 1
+; RUN: echo 'v1' > %t2
+; RUN: echo 'f _Z3bazb' >> %t2
+; RUN: echo 'c 0' >> %t2
+;;
+; RUN: llc < %s -mtriple=aarch64 -function-sections -basic-block-sections=%t1 -unique-basic-block-section-names | FileCheck %s -check-prefix=SECTIONS
+; RUN: llc < %s -mtriple=aarch64 -function-sections -basic-block-sections=%t2 -unique-basic-block-section-names | FileCheck %s -check-prefix=SECTIONS
+; 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=SPLIT
+
+define void @_Z3bazb(i1 zeroext %0) nounwind {
+  br i1 %0, label %2, label %4
+
+2:                                                ; preds = %1
+  %3 = call i32 @_Z3barv()
+  br label %6
+
+4:                                                ; preds = %1
+  %5 = call i32 @_Z3foov()
+  br label %6
+
+6:                                                ; preds = %2, %4
+  ret void
+}
+
+declare i32 @_Z3barv() #1
+
+declare i32 @_Z3foov() #1
+
+; SECTIONS: .section        .text.hot._Z3bazb,"ax",@progbits
+; SECTIONS: _Z3bazb:
+; Check that the basic block with id 1 doesn't get a section.
+; SECTIONS-NOT: .section        .text{{.*}}._Z3bazb.1,"ax",@progbits,unique
+; Check that a single cold section is started here and id 1 and 2 blocks are placed here.
+; SECTIONS: .section	.text.split._Z3bazb,"ax",@progbits
+; SECTIONS: _Z3bazb.cold:
+; SECTIONS-NOT: .section        .text.hot._Z3bazb._Z3bazb.2,"ax",@progbits,unique
+; SECTIONS: .LBB0_2:
+; SECTIONS: .size   _Z3bazb, .Lfunc_end{{[0-9]}}-_Z3bazb
+
+; SPLIT:      .section	.text.unlikely._Z3bazb,"ax",@progbits
+; SPLIT-NEXT: _Z3bazb.cold:
+; SPLIT-NEXT:   bl _Z3barv
+; SPLIT:      .LBB0_2:
+; SPLIT:      .LBB_END0_2:
diff --git a/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll b/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll
new file mode 100644
index 000000000000..a83a47c9c129
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/basic-block-sections-unsafe.ll
@@ -0,0 +1,121 @@
+;; Check if basic blocks without unique sections are only placed in cold sections if it is safe
+;; to do so.
+;;
+;; Profile for version 0.
+; RUN: echo 'v1' > %t1
+; RUN: echo 'f _Z3asm_goto' >> %t1
+; RUN: echo 'c 0' >> %t1
+; RUN: echo 'f _Z3jump_table' >> %t1
+; RUN: echo 'c 0' >> %t1
+; RUN: echo 'f _Z3red_zone' >> %t1
+; RUN: echo 'c 0' >> %t1
+;;
+; RUN: llc < %s -mtriple=aarch64 -function-sections -basic-block-sections=%t1 -unique-basic-block-section-names -bbsections-cold-text-prefix=".text.unlikely." | FileCheck %s
+; RUN: llc < %s -mtriple=aarch64 -function-sections -aarch64-min-jump-table-entries=4 -basic-block-sections=%t1 -unique-basic-block-section-names -bbsections-cold-text-prefix=".text.unlikely." | FileCheck %s -check-prefix=JUMP-TABLES
+; 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:        .section	.text.unlikely._Z3asm_goto,"ax",@progbits
+  ; CHECK-NEXT:     _Z3asm_goto.cold:
+  ; CHECK-NEXT:       bl bam
+  ; CHECK:          .LBB0_4:
+  ; CHECK:            ret
+  ; CHECK:          .LBB_END0_4:
+
+  br i1 %0, label %3, label %5
+
+3:                                                ; preds = %2
+  %4 = call i32 @bar()
+  callbr void asm sideeffect "nop", "!i"() #3
+          to label %asm.fallthrough [label %5]
+
+
+asm.fallthrough:                                  ; preds = %3
+    br label %5
+
+5:                                                ; preds = %2, %asm.fallthrough
+  %6 = call i32 @bar()
+  br i1 %1, label %7, label %9
+
+7:
+  %8 = call i32 @bam()
+  br label %9
+
+9:                                                ; preds = %7
+  ret void
+}
+
+define i32 @_Z3jump_table(i32 %in) nounwind {
+  ;; Check that a cold block that contains a jump table dispatch or
+  ;; that is targeted by a jump table is not split.
+  ; JUMP-TABLES-LABEL:  _Z3jump_table
+  ; JUMP-TABLES:        .section	.text.unlikely._Z3jump_table,"ax",@progbits
+  ; JUMP-TABLES-NEXT:     _Z3jump_table.cold:
+  ; JUMP-TABLES-SAME:                         %common.ret
+  ; JUMP-TABLES-NOT:        b       bar
+  ; JUMP-TABLES-NOT:        b       baz
+  ; JUMP-TABLES-NOT:        b       qux
+  ; JUMP-TABLES-NOT:        b       bam
+
+  switch i32 %in, label %common.ret [
+    i32 0, label %cold1
+    i32 1, label %cold2
+    i32 2, label %cold3
+    i32 3, label %cold4
+  ]
+
+  common.ret:                                       ; preds = %0
+    ret i32 0
+
+  cold1:                                            ; preds = %0
+    %1 = tail call i32 @bar()
+    ret i32 %1
+
+  cold2:                                            ; preds = %0
+    %2 = tail call i32 @baz()
+    ret i32 %2
+
+  cold3:                                            ; preds = %0
+    %3 = tail call i32 @bam()
+    ret i32 %3
+
+  cold4:                                            ; preds = %0
+    %4 = tail call i32 @qux()
+    ret i32 %4
+}
+
+define i32 @_Z3red_zone(i1 zeroext %0, i32 %a, i32 %b) nounwind {
+;; Check that cold blocks in functions with red zones aren't split.
+; RED-ZONE-LABEL:        _Z3red_zone
+; MFS-REDZONE-AARCH64-NOT:   _Z3red_zone.cold:
+  %a.addr = alloca i32, align 4
+  %b.addr = alloca i32, align 4
+  %x = alloca i32, align 4
+
+  br i1 %0, label %2, label %3
+
+2:                                                ; preds = %1
+  store i32 %a, ptr %a.addr, align 4
+  store i32 %b, ptr %b.addr, align 4
+  br label %4
+
+3:                                                ; preds = %1
+  store i32 %a, ptr %b.addr, align 4
+  store i32 %b, ptr %a.addr, align 4
+  br label %4
+
+4:                                                ; preds = %3, %2
+  %tmp = load i32, ptr %a.addr, align 4
+  %tmp1 = load i32, ptr %b.addr, align 4
+  %add = add nsw i32 %tmp, %tmp1
+  store i32 %add, ptr %x, align 4
+  %tmp2 = load i32, ptr %x, align 4
+  ret i32 %tmp2
+}
+
+declare i32 @bar()
+declare i32 @baz()
+declare i32 @bam()
+declare i32 @qux()