[llvm] [AArch64] Fix LDR/STR folding causing memtag failures (PR #118821)
Oliver Stannard via llvm-commits
llvm-commits at lists.llvm.org
Thu Dec 5 07:48:12 PST 2024
https://github.com/ostannard created https://github.com/llvm/llvm-project/pull/118821
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.
>From 7a9d84ec15bc393bdaeccc79c4a3686454a79f33 Mon Sep 17 00:00:00 2001
From: Oliver Stannard <oliver.stannard at arm.com>
Date: Thu, 5 Dec 2024 10:09:47 +0000
Subject: [PATCH 1/2] Test showing incorrect codegen
---
.../AArch64/memtag-merge-writeback.mir | 140 ++++++++++++++++++
1 file changed, 140 insertions(+)
create mode 100644 llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir
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..91c452e04abab3
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir
@@ -0,0 +1,140 @@
+# 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: 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_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: 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_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
+...
>From 37fd5facb09677a09d9e6012fb471d9096454c54 Mon Sep 17 00:00:00 2001
From: Oliver Stannard <oliver.stannard at arm.com>
Date: Thu, 5 Dec 2024 15:44:36 +0000
Subject: [PATCH 2/2] [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.
---
.../Target/AArch64/AArch64LoadStoreOptimizer.cpp | 14 ++++++++++++--
.../CodeGen/AArch64/memtag-merge-writeback.mir | 6 ++++--
2 files changed, 16 insertions(+), 4 deletions(-)
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
index 91c452e04abab3..be828f7f1da059 100644
--- a/llvm/test/CodeGen/AArch64/memtag-merge-writeback.mir
+++ b/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
More information about the llvm-commits
mailing list