diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp index 4f46ca1538db2..cd976790ebb6f 100644 --- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp +++ b/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(); 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; diff --git a/llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir b/llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir new file mode 100644 index 0000000000000..be828f7f1da05 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir @@ -0,0 +1,142 @@ +# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5 +# RUN: llc -mtriple aarch64-none-elf -mattr=+mte --run-pass=aarch64-ldst-opt %s -o - | FileCheck %s + +## 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. + +--- | + define void @pre_index() { + entry: + ret void + } + define void @pre_index_memtag() sanitize_memtag { + entry: + ret void + } + define void @pre_index_memtag_not_sp() sanitize_memtag { + entry: + ret void + } + define void @post_index() { + entry: + ret void + } + define void @post_index_memtag() sanitize_memtag { + entry: + ret void + } + define void @post_index_memtag_not_sp() sanitize_memtag { + entry: + ret void + } +... +--- +name: pre_index +body: | + bb.0.entry: + liveins: $x0 + + ; CHECK-LABEL: name: pre_index + ; 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: RET undef $lr + $sp = frame-setup SUBXri $sp, 16, 0 + STRXui killed renamable $x0, $sp, 2 + $sp = ADDXri $sp, 16, 0 + RET undef $lr +... +--- +name: pre_index_memtag +body: | + bb.0.entry: + liveins: $x0 + + ; CHECK-LABEL: name: pre_index_memtag + ; CHECK: liveins: $x0 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0 + ; 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 + $sp = ADDXri $sp, 16, 0 + RET undef $lr +... +--- +name: pre_index_memtag_not_sp +body: | + bb.0.entry: + liveins: $x0, $x1 + + ; CHECK-LABEL: name: pre_index_memtag_not_sp + ; CHECK: liveins: $x0, $x1 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $x1 = frame-setup SUBXri $x1, 16, 0 + ; CHECK-NEXT: early-clobber $x1 = STRXpre killed renamable $x0, $x1, 16 + ; CHECK-NEXT: RET undef $lr, implicit $x1 + $x1 = frame-setup SUBXri $x1, 16, 0 + STRXui killed renamable $x0, $x1, 2 + $x1 = ADDXri $x1, 16, 0 + RET undef $lr, implicit $x1 +... +--- +name: post_index +body: | + bb.0.entry: + liveins: $x0 + + ; CHECK-LABEL: name: post_index + ; 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: RET undef $lr + $sp = frame-setup SUBXri $sp, 16, 0 + STRXui killed renamable $x0, $sp, 0 + $sp = ADDXri $sp, 16, 0 + RET undef $lr +... +--- +name: post_index_memtag +body: | + bb.0.entry: + liveins: $x0 + + ; CHECK-LABEL: name: post_index_memtag + ; CHECK: liveins: $x0 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 16, 0 + ; 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 + $sp = ADDXri $sp, 16, 0 + RET undef $lr +... +--- +name: post_index_memtag_not_sp +body: | + bb.0.entry: + liveins: $x0, $x1 + + ; CHECK-LABEL: name: post_index_memtag_not_sp + ; CHECK: liveins: $x0, $x1 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $x1 = frame-setup SUBXri $x1, 16, 0 + ; CHECK-NEXT: early-clobber $x1 = STRXpost killed renamable $x0, $x1, 16 + ; CHECK-NEXT: RET undef $lr, implicit $x1 + $x1 = frame-setup SUBXri $x1, 16, 0 + STRXui killed renamable $x0, $x1, 0 + $x1 = ADDXri $x1, 16, 0 + RET undef $lr, implicit $x1 +...