[CodeGen][AArch64] Don't split inline asm goto blocks or their targets

[dhoekwater]

Summary

Machine function splitting + branch relaxation currently don't properly
handle inline asm goto blocks that conditional branch to cold goto
labels. While such inline asm is technically invalid, machine
function splitting is the only thing that exposes it as such.

Since machine function splitting doesn't help too much in these
circumstances anyway, disable it for asm goto blocks and their targets.

[dhoekwater]

Rebase onto main and update test function numbers.

[shenhan]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8448

This test (MBB.isInlineAsmBrIndirectTarget()) is quicker and cheaper than acquiring and iterating over jump-tbl , we may rearrange the tests so that the quickest and cheapest tests are placed earlier than those less cheap.

[dhoekwater]

Place cheap checks before expensive ones per reviewer feedback

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::isFunctionSafeToSplit(const MachineFunction &MF) const {
   if (MF.getInfo<AArch64FunctionInfo>()->hasRedZone().value_or(true))
     return false;
 
   return TargetInstrInfo::isFunctionSafeToSplit(MF);
 }
 
 bool AArch64InstrInfo::isMBBSafeToSplitToCold(
     const MachineBasicBlock &MBB) const {
+  // Asm Goto blocks can contain conditional branches to goto labels, which can
+  // get moved out of range of the branch instruction.
+  auto isAsmGoto = [](const MachineInstr &MI) {
+    return MI.getOpcode() == AArch64::INLINEASM_BR;
+  };
+  if (llvm::any_of(MBB, isAsmGoto) || MBB.isInlineAsmBrIndirectTarget())
+    return false;
+
   // Because jump tables are label-relative instead of table-relative, they all
   // must be in the same section or relocation fixup handling will fail.
 
   // Check if MBB is a jump table target
   const MachineJumpTableInfo *MJTI = MBB.getParent()->getJumpTableInfo();
   auto containsMBB = [&MBB](const MachineJumpTableEntry &JTE) {
     return llvm::is_contained(JTE.MBBs, &MBB);
   };
for (const MachineInstr &MI : MBB) {
     }
   }
 
   // MBB isn't a special case, so it's safe to be split to the cold section.
   return true;
 }
 
 std::optional<ParamLoadedValue>
 AArch64InstrInfo::describeLoadedValue(const MachineInstr &MI,
                                       Register Reg) const {
   const MachineFunction *MF = MI.getMF();
   const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
   switch (MI.getOpcode()) {
   case AArch64::MOVZWi:
   case AArch64::MOVZXi: {
     // MOVZWi may be used for producing zero-extended 32-bit immediates in
     // 64-bit parameters, so we need to consider super-registers.

llvm/test/CodeGen/Generic/machine-function-splitter.ll

cold1:                                            ; preds = %0
   %3 = tail call i32 @bam()
   ret i32 %3
 
 cold2:                                            ; preds = %0
   %4 = tail call i32 @qux()
   ret i32 %4
 }
 
+define void @foo20(i1 zeroext %0) !prof !14 !section_prefix !15 {
+;; Check that blocks containing or targeted by asm goto aren't split.
+; MFS-DEFAULTS-LABEL:        foo20
+; MFS-DEFAULTS-AARCH64-NOT:  foo20.cold:
+; MFS-DEFAULTS-X86:          .section        .text.split.foo20
+; MFS-DEFAULTS-X86:          foo20.cold:
+; MFS-DEFAULTS-X86-DAG:      # %cold_asm
+; MFS-DEFAULTS-X86-DAG:      # %cold_asm_target
+
+  br i1 %0, label %hot, label %cold_asm, !prof !17
+
+hot:
+  %2 = call i32 @bar()
+  ret void
+
+cold_asm:
+    callbr void asm sideeffect "nop", "!i"() #3
+          to label %asm.fallthrough [label %cold_asm_target]
+
+asm.fallthrough:
+    br label %cold_asm_target
+
+cold_asm_target:
+  %3 = call i32 @baz()
+  ret void
+}
+
 declare i32 @bar()
 declare i32 @baz()
 declare i32 @bam()
 declare i32 @qux()
 declare void @_Z1fv()
 declare i32 @__gxx_personality_v0(...)
 
 @_ZTIi = external constant ptr