[AArch64] Fix LDR/STR folding causing memtag failures #118821
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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.
|
@llvm/pr-subscribers-backend-aarch64 Author: Oliver Stannard (ostannard) ChangesWhen 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. Full diff: https://github.com/llvm/llvm-project/pull/118821.diff 2 Files Affected:
diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index 4f46ca1538db2b..cd976790ebb6f8 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<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;
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 00000000000000..be828f7f1da059
--- /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
+...
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.