[llvm] [AArch64] Stack probing for dynamic allocas in GlobalISel (PR #67123)

[Oskar Wirga via llvm-commits]

================
@@ -8622,6 +8623,96 @@ bool AArch64InstrInfo::isReallyTriviallyReMaterializable(
   return TargetInstrInfo::isReallyTriviallyReMaterializable(MI);
 }
 
+MachineBasicBlock::iterator
+AArch64InstrInfo::insertStackProbingLoop(MachineBasicBlock::iterator MBBI,
+                                         Register ScratchReg,
+                                         Register TargetReg) const {
+  MachineBasicBlock &MBB = *MBBI->getParent();
+  MachineFunction &MF = *MBB.getParent();
+  const AArch64TargetLowering *TLI =
+      MF.getSubtarget<AArch64Subtarget>().getTargetLowering();
+  const AArch64InstrInfo *TII =
+      MF.getSubtarget<AArch64Subtarget>().getInstrInfo();
+  int64_t ProbeSize = (int64_t)TLI->getStackProbeSize(MF);
+  DebugLoc DL = MBB.findDebugLoc(MBBI);
+
+  MachineFunction::iterator MBBInsertPoint = std::next(MBB.getIterator());
+  MachineBasicBlock *LoopTestMBB =
+      MF.CreateMachineBasicBlock(MBB.getBasicBlock());
+  MF.insert(MBBInsertPoint, LoopTestMBB);
+  MachineBasicBlock *LoopBodyMBB =
+      MF.CreateMachineBasicBlock(MBB.getBasicBlock());
+  MF.insert(MBBInsertPoint, LoopBodyMBB);
+  MachineBasicBlock *ExitMBB = MF.CreateMachineBasicBlock(MBB.getBasicBlock());
+  MF.insert(MBBInsertPoint, ExitMBB);
+
----------------
oskarwirga wrote:

```suggestion

    // Initialize flag indicating whether NZCV needs to be saved
  bool shouldSaveConditionFlag = false;
  bool foundStackAlloc = false;
  bool breakOuterLoop = false;
  Register SavedConditionReg;

  for (auto &currentBlock : MF) {
    for (auto instr = currentBlock.instr_begin(),
              end = currentBlock.instr_end();
         instr != end; ++instr) {
      if (instr->getOpcode() == AArch64::PROBED_STACKALLOC_DYN) {
        foundStackAlloc = true;
      }
      if (foundStackAlloc) {
        if (instr->definesRegister(llvm::AArch64::NZCV)) {
          shouldSaveConditionFlag = false;
          breakOuterLoop = true;
          break;
        }
        if (instr->readsRegister(llvm::AArch64::NZCV)) {
          shouldSaveConditionFlag = true;
          breakOuterLoop = true;
          break;
        }
      }
    }
    if (breakOuterLoop)
      break;
  }

  if (shouldSaveConditionFlag) {
    MachineRegisterInfo &MRI = MF.getRegInfo();

    const TargetRegisterClass *RC = &AArch64::GPR64RegClass;
    SavedConditionReg = MRI.createVirtualRegister(RC);

    // Define SavedConditionReg.
    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrr), SavedConditionReg)
        .addReg(AArch64::XZR)
        .addReg(AArch64::XZR);
  }
  
```

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