[AArch64] Async unwind - Fix MTE codegen emitting frame adjustments in a loop

[chill]

Summary

When untagging the stack, the compiler may emit a sequence like:

.LBB0_1:
  st2g sp, [sp], #32
  sub x8, x8, #32
  cbnz x8, .LBB0_1
  stg sp, [sp], #16

These stack adjustments cannot be described by CFI instructions.

This patch disables merging of SP update with untagging, i.e. makes the
compiler use an additional scratch register (there should be plenty
available at this point as we are in the epilogue) and generate:

        mov     x9, sp
        mov     x8, #256
        stg     x9, [x9], #16
.LBB0_1:
        sub     x8, x8, #32
        st2g    x9, [x9], #32
        cbnz    x8, .LBB0_1
        add     sp, sp, #272

Merging is disabled only when we need to generate asynchronous unwind
tables.

[efriedma]

It looks like none of the testcases actually generate the loop using an additional scratch register described in the commit message?

[chill]

In D114548#3152193, @efriedma wrote:

It looks like none of the testcases actually generate the loop using an additional scratch register described in the commit message?

D'oh, that was the whole point of duplicating stg_alloca17() into stg_alloca18(). Fixed.

[eugenis]

Sorry for the delayed review.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

3494

typo: TryMergeSPUpdate

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

609

Is this an unrelated fix? Are we generating broken code because of this?

[chill]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

609

We are changing an "ordinary" instruction to a one with a writeback. Without this part we get

*** Bad machine code: Operand should be tied ***
- function:    stg_alloca18
- basic block: %bb.0 entry (0x70c9038)
- instruction: dead early-clobber renamable $x8, dead early-clobber renamable $x9 = STGloop_wback 272, killed $x9 :: (store (s2176) into %ir.a, align 16)
- operand 3:   killed $x9

when running MachineVerifier.

[efriedma]

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

591

"but not the operands" needs to be updated.

593

Maybe rename this helper? Not a big deal either way.

601

Minor nit: move the definition of "Op" inside the if statement.

[chill]

PIng?

[eugenis]

LGTM

[chill]

Rebased.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

assert((TagStores.empty() ||
            "Non-adjacent tag store instructions.");
     TagStores.push_back(I);
   }
   void clear() { TagStores.clear(); }
   // Emit equivalent code at the given location, and erase the current set of
   // instructions. May skip if the replacement is not profitable. May invalidate
   // the input iterator and replace it with a valid one.
   void emitCode(MachineBasicBlock::iterator &InsertI,
-                const AArch64FrameLowering *TFI, bool IsLast);
+                const AArch64FrameLowering *TFI, bool TryMergeSPUpdate);
 };
 
 void TagStoreEdit::emitUnrolled(MachineBasicBlock::iterator InsertI) {
   const AArch64InstrInfo *TII =
       MF->getSubtarget<AArch64Subtarget>().getInstrInfo();
 
   const int64_t kMinOffset = -256 * 16;
   const int64_t kMaxOffset = 255 * 16;
if (MI->memoperands_empty()) {
       MemRefs.clear();
       return;
     }
     MemRefs.append(MI->memoperands_begin(), MI->memoperands_end());
   }
 }
 
 void TagStoreEdit::emitCode(MachineBasicBlock::iterator &InsertI,
-                            const AArch64FrameLowering *TFI, bool IsLast) {
+                            const AArch64FrameLowering *TFI,
+                            bool TryMergeSPUpdate) {
   if (TagStores.empty())
     return;
   TagStoreInstr &FirstTagStore = TagStores[0];
   TagStoreInstr &LastTagStore = TagStores[TagStores.size() - 1];
   Size = LastTagStore.Offset - FirstTagStore.Offset + LastTagStore.Size;
   DL = TagStores[0].MI->getDebugLoc();
 
   Register Reg;
void TagStoreEdit::emitCode(MachineBasicBlock::iterator &InsertI,
   // tagging instructions.
   const int kSetTagLoopThreshold = 176;
   if (Size < kSetTagLoopThreshold) {
     if (TagStores.size() < 2)
       return;
     emitUnrolled(InsertI);
   } else {
     MachineInstr *UpdateInstr = nullptr;
-    int64_t TotalOffset;
-    if (IsLast) {
+    int64_t TotalOffset = 0;
+    if (TryMergeSPUpdate) {
       // See if we can merge base register update into the STGloop.
       // This is done in AArch64LoadStoreOptimizer for "normal" stores,
       // but STGloop is way too unusual for that, and also it only
       // realistically happens in function epilogue. Also, STGloop is expanded
       // before that pass.
       if (InsertI != MBB->end() &&
           canMergeRegUpdate(InsertI, FrameReg, FrameRegOffset.getFixed() + Size,
                             &TotalOffset)) {
MachineBasicBlock::iterator tryMergeAdjacentSTG(MachineBasicBlock::iterator II,
 
   // Find contiguous runs of tagged memory and emit shorter instruction
   // sequencies for them when possible.
   TagStoreEdit TSE(MBB, FirstZeroData);
   Optional<int64_t> EndOffset;
   for (auto &Instr : Instrs) {
     if (EndOffset && *EndOffset != Instr.Offset) {
       // Found a gap.
-      TSE.emitCode(InsertI, TFI, /*IsLast = */ false);
+      TSE.emitCode(InsertI, TFI, /*TryMergerSPUpdate = */ false);
       TSE.clear();
     }
 
     TSE.addInstruction(Instr);
     EndOffset = Instr.Offset + Instr.Size;
   }
 
-  TSE.emitCode(InsertI, TFI, /*IsLast = */ true);
+  // Multiple FP/SP updates in a loop cannot be described by CFI instructions.
+  TSE.emitCode(InsertI, TFI, /*TryMergeSPUpdate = */
+                             !MBB->getParent()
+                                 ->getInfo<AArch64FunctionInfo>()
+                                 ->needsAsyncUnwindInfo());
 
   return InsertI;
 }
 } // namespace
 
 void AArch64FrameLowering::processFunctionBeforeFrameIndicesReplaced(
     MachineFunction &MF, RegScavenger *RS = nullptr) const {
   if (StackTaggingMergeSetTag)

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

const AArch64InstrInfo *TII =
       MF->getSubtarget<AArch64Subtarget>().getInstrInfo();
   bool Done = rewriteAArch64FrameIndex(MI, i, BaseReg, Off, TII);
   assert(Done && "Unable to resolve frame index!");
   (void)Done;
 }
 
 // Create a scratch register for the frame index elimination in an instruction.
 // This function has special handling of stack tagging loop pseudos, in which
 // case it can also change the instruction opcode (but not the operands).
 static Register
-createScratchRegisterForInstruction(MachineInstr &MI,
+createScratchRegisterForInstruction(MachineInstr &MI, unsigned FIOperandNum,
                                     const AArch64InstrInfo *TII) {
   // ST*Gloop have a reserved scratch register in operand 1. Use it, and also
   // replace the instruction with the writeback variant because it will now
   // satisfy the operand constraints for it.
-  if (MI.getOpcode() == AArch64::STGloop) {
-    MI.setDesc(TII->get(AArch64::STGloop_wback));
-    return MI.getOperand(1).getReg();
-  } else if (MI.getOpcode() == AArch64::STZGloop) {
-    MI.setDesc(TII->get(AArch64::STZGloop_wback));
-    return MI.getOperand(1).getReg();
+  Register ScratchReg;
+  unsigned Op = MI.getOpcode() == AArch64::STGloop    ? AArch64::STGloop_wback
+                : MI.getOpcode() == AArch64::STZGloop ? AArch64::STZGloop_wback
+                                                      : 0 /* unused */;
+  if (MI.getOpcode() == AArch64::STGloop ||
+      MI.getOpcode() == AArch64::STZGloop) {
+    assert(FIOperandNum == 3 &&
+           "Wrong frame index operand for STGloop/STZGloop");
+    ScratchReg = MI.getOperand(1).getReg();
+    MI.getOperand(3).ChangeToRegister(ScratchReg, false, false, true);
+    MI.setDesc(TII->get(Op));
+    MI.tieOperands(1, 3);
   } else {
-    return MI.getMF()->getRegInfo().createVirtualRegister(
-        &AArch64::GPR64RegClass);
+    ScratchReg =
+        MI.getMF()->getRegInfo().createVirtualRegister(&AArch64::GPR64RegClass);
+    MI.getOperand(FIOperandNum)
+        .ChangeToRegister(ScratchReg, false, false, true);
   }
+  return ScratchReg;
 }
 
 void AArch64RegisterInfo::getOffsetOpcodes(
     const StackOffset &Offset, SmallVectorImpl<uint64_t> &Ops) const {
   // The smallest scalable element supported by scaled SVE addressing
   // modes are predicates, which are 2 scalable bytes in size. So the scalable
   // byte offset must always be a multiple of 2.
   assert(Offset.getScalable() % 2 == 0 && "Invalid frame offset");
if (rewriteAArch64FrameIndex(MI, FIOperandNum, FrameReg, Offset, TII))
     return;
 
   assert((!RS || !RS->isScavengingFrameIndex(FrameIndex)) &&
          "Emergency spill slot is out of reach");
 
   // If we get here, the immediate doesn't fit into the instruction.  We folded
   // as much as possible above.  Handle the rest, providing a register that is
   // SP+LargeImm.
-  Register ScratchReg = createScratchRegisterForInstruction(MI, TII);
+  Register ScratchReg = createScratchRegisterForInstruction(MI, FIOperandNum, TII);
   emitFrameOffset(MBB, II, MI.getDebugLoc(), ScratchReg, FrameReg, Offset, TII);
-  MI.getOperand(FIOperandNum).ChangeToRegister(ScratchReg, false, false, true);
 }
 
 unsigned AArch64RegisterInfo::getRegPressureLimit(const TargetRegisterClass *RC,
                                                   MachineFunction &MF) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
 
   switch (RC->getID()) {
   default:

llvm/test/CodeGen/AArch64/settag.ll

 ; CHECK-NEXT:    .cfi_def_cfa_offset 0
 ; CHECK-NEXT:    ret
 entry:
   %a = alloca i8, i32 80, align 16
   call void @llvm.aarch64.settag(i8* %a, i64 80)
   ret void
 }
 
-define void @stg_alloca17() uwtable {
+define void @stg_alloca17() nounwind {
 ; CHECK-LABEL: stg_alloca17:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sub sp, sp, #288
+; CHECK-NEXT:    mov x8, #256
+; CHECK-NEXT:    str x29, [sp, #272] // 8-byte Folded Spill
+; CHECK-NEXT:  .LBB11_1: // %entry
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    st2g sp, [sp], #32
+; CHECK-NEXT:    sub x8, x8, #32
+; CHECK-NEXT:    cbnz x8, .LBB11_1
+; CHECK-NEXT:  // %bb.2: // %entry
+; CHECK-NEXT:    stg sp, [sp], #16
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %a = alloca i8, i32 272, align 16
+  call void @llvm.aarch64.settag(i8* %a, i64 272)
+  ret void
+}
+
+define void @stg_alloca18() {
+; CHECK-LABEL: stg_alloca18:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    sub sp, sp, #288
 ; CHECK-NEXT:    .cfi_def_cfa_offset 288
 ; CHECK-NEXT:    str x29, [sp, #272] // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_offset w29, -16
 ; CHECK-NEXT:    mov x8, #256
-; CHECK-NEXT:  .LBB11_1: // %entry
+; CHECK-NEXT:  .LBB12_1: // %entry
 ; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    st2g sp, [sp], #32
 ; CHECK-NEXT:    sub x8, x8, #32
-; CHECK-NEXT:    cbnz x8, .LBB11_1
+; CHECK-NEXT:    cbnz x8, .LBB12_1
 ; CHECK-NEXT:  // %bb.2: // %entry
 ; CHECK-NEXT:    stg sp, [sp], #16
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
-; CHECK-NEXT:    .cfi_def_cfa_offset 0
-; CHECK-NEXT:    .cfi_restore w29
 ; CHECK-NEXT:    ret
 entry:
   %a = alloca i8, i32 272, align 16
   call void @llvm.aarch64.settag(i8* %a, i64 272)
   ret void
 }
 
 declare void @llvm.aarch64.settag(i8* %p, i64 %a)
 declare void @llvm.aarch64.settag.zero(i8* %p, i64 %a)