[AArch64] Fix LDR/STR folding causing memtag failures

When generating code with sanitize_memtag, we make use of the fact that
the sp+imm forms of many load and store instructions are not
tag-checked, so we can use SP directly instead of needing a register
holding the tagged pointer. However, this isn't true for the writeback
versions of the instructions, so we can't fold ADDs and SUBs into them
in AArch64LoadStoreOptimizer. This would be possible in cases where the
loads/stores only access untagged stack slots, but that information
isn't easily available after frame index elimination.

Author: ostannard

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

@@ -733,7 +733,7 @@ static bool isPromotableLoadFromStore(MachineInstr &MI) {
   }
 }
 
-static bool isMergeableLdStUpdate(MachineInstr &MI) {
+static bool isMergeableLdStUpdate(MachineInstr &MI, AArch64FunctionInfo &AFI) {
   unsigned Opc = MI.getOpcode();
   switch (Opc) {
   default:
@@ -785,6 +785,15 @@ static bool isMergeableLdStUpdate(MachineInstr &MI) {
     if (!AArch64InstrInfo::getLdStOffsetOp(MI).isImm())
       return false;
 
+    // When using stack tagging, simple sp+imm loads and stores are not
+    // tag-checked, but pre- and post-indexed versions of them are, so we can't
+    // replace the former with the latter. This transformation would be valid
+    // if the load/store accesses an untagged stack slot, but we don't have
+    // that information available after frame indices have been eliminated.
+    if (AFI.isMTETagged() &&
+        AArch64InstrInfo::getLdStBaseOp(MI).getReg() == AArch64::SP)
+      return false;
+
     return true;
   }
 }
@@ -2772,6 +2781,7 @@ bool AArch64LoadStoreOpt::tryToMergeIndexLdSt(MachineBasicBlock::iterator &MBBI,
 
 bool AArch64LoadStoreOpt::optimizeBlock(MachineBasicBlock &MBB,
                                         bool EnableNarrowZeroStOpt) {
+  AArch64FunctionInfo &AFI = *MBB.getParent()->getInfo<AArch64FunctionInfo>();
 
   bool Modified = false;
   // Four tranformations to do here:
@@ -2842,7 +2852,7 @@ bool AArch64LoadStoreOpt::optimizeBlock(MachineBasicBlock &MBB,
   //        ldr x0, [x2], #4
   for (MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
        MBBI != E;) {
-    if (isMergeableLdStUpdate(*MBBI) && tryToMergeLdStUpdate(MBBI))
+    if (isMergeableLdStUpdate(*MBBI, AFI) && tryToMergeLdStUpdate(MBBI))
       Modified = true;
     else
       ++MBBI;

llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir

@@ -63,7 +63,8 @@ body:             |
     ; CHECK: liveins: $x0
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0
-    ; CHECK-NEXT: early-clobber $sp = STRXpre killed renamable $x0, $sp, 16
+    ; CHECK-NEXT: STRXui killed renamable $x0, $sp, 2
+    ; CHECK-NEXT: $sp = ADDXri $sp, 16, 0
     ; CHECK-NEXT: RET undef $lr
     $sp = frame-setup SUBXri $sp, 16, 0
     STRXui killed renamable $x0, $sp, 2
@@ -114,7 +115,8 @@ body:             |
     ; CHECK: liveins: $x0
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0
-    ; CHECK-NEXT: early-clobber $sp = STRXpost killed renamable $x0, $sp, 16
+    ; CHECK-NEXT: STRXui killed renamable $x0, $sp, 0
+    ; CHECK-NEXT: $sp = ADDXri $sp, 16, 0
     ; CHECK-NEXT: RET undef $lr
     $sp = frame-setup SUBXri $sp, 16, 0
     STRXui killed renamable $x0, $sp, 0