[AArch64] Fix range check for STGPostIndex

When generating function epilogues using AArch64 stack tagging, we can
fold an SP update into the tag-setting loop. The loop tags 32 bytes at a
time using ST2G, so the final SP update might be done either by a post
indexed STG which tags the final 16 bytes of the tagged region, or by an
ADD/SUB instruction after the loop. However, we were only considering
the range of the ADD/SUB instructions when deciding whether to do this,
and the valid immediate range for STG is slightly lower when the offset
is positive.

Author: ostannard

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -4393,10 +4393,20 @@ bool canMergeRegUpdate(MachineBasicBlock::iterator II, unsigned Reg,
     int64_t Offset = MI.getOperand(2).getImm() << Shift;
     if (MI.getOpcode() == AArch64::SUBXri)
       Offset = -Offset;
-    int64_t AbsPostOffset = std::abs(Offset - Size);
-    const int64_t kMaxOffset =
-        0xFFF; // Max encoding for unshifted ADDXri / SUBXri
-    if (AbsPostOffset <= kMaxOffset && AbsPostOffset % 16 == 0) {
+    int64_t PostOffset = Offset - Size;
+    // TagStoreEdit::emitLoop might emit either an ADD/SUB after the loop, or
+    // an STGPostIndex which does the last 16 bytes of tag write. Which one is
+    // chosen depends on the alignment of the loop size, but the difference
+    // between the valid ranges for the two instructions is small, so we
+    // conservatively assume that it could be either case here.
+    //
+    // Max offset of STGPostIndex, minus the 16 byte tag write folded into that
+    // instruction.
+    const int64_t kMaxOffset = 4080 - 16;
+    // Max offset of SUBXri.
+    const int64_t kMinOffset = -4095;
+    if (PostOffset <= kMaxOffset && PostOffset >= kMinOffset &&
+        PostOffset % 16 == 0) {
       *TotalOffset = Offset;
       return true;
     }

llvm/test/CodeGen/AArch64/stack-tagging-epilogue-fold.mir

@@ -53,8 +53,9 @@ body:             |
     ; CHECK-NEXT: renamable $x10 = COPY renamable $x9
     ; CHECK-NEXT: dead early-clobber renamable $x11, dead early-clobber renamable $x10 = STGloop_wback 4080, killed renamable $x10, implicit-def dead $nzcv
     ; CHECK-NEXT: ST2Gi renamable $x8, renamable $x8, 0
-    ; CHECK-NEXT: dead early-clobber $x8, early-clobber $sp = frame-destroy STGloop_wback 4096, $sp, implicit-def $nzcv
-    ; CHECK-NEXT: $sp = frame-destroy STGPostIndex $sp, $sp, 256
+    ; CHECK-NEXT: $x9 = ADDXri $sp, 0, 0
+    ; CHECK-NEXT: dead early-clobber $x8, dead early-clobber $x9 = STGloop_wback 4112, $x9, implicit-def $nzcv
+    ; CHECK-NEXT: $sp = frame-destroy ADDXri $sp, 2, 12
     ; CHECK-NEXT: early-clobber $sp, $fp = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.5)
     ; CHECK-NEXT: RET_ReallyLR
     /* Prologue */