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

[Daniel Hoekwater via llvm-commits]

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

>From 40d622aa2ff236f2e46f6ae5961401b5d60d5ac0 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 in BBSections

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.
---
 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);
+        }
       }
     }