[llvm] [CodeGen][AArch64] Only split safe blocks and functions in BBSections (PR #81553)

[Daniel Hoekwater via llvm-commits]

https://github.com/dhoekwater created https://github.com/llvm/llvm-project/pull/81553

Some types of machine function and machine basic block are unsafe to
split on AArch64: functions with red zone enabled (D157127), 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.


>From feb0f9006042cd2c38d43315172370e83e443a2f Mon Sep 17 00:00:00 2001
From: Daniel Hoekwater <hoekwater at google.com>
Date: Sat, 10 Feb 2024 00:51:45 +0000
Subject: [PATCH] [CodeGen][AArch64] Only split safe blocks and functions in
 BBSections

Some types of machine function and machine basic block are unsafe to
split on AArch64: functions with red zone enabled (D157127), 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.
---
 llvm/lib/CodeGen/BasicBlockSections.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index eb3f9e7078f1ac..09e45ea5794b76 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);
+        }
       }
     }