[llvm] [AArch64] Preserve store width for ptr32 address spaces (PR #229181)

[Daniel Paoliello via llvm-commits]

https://github.com/dpaoliello created https://github.com/llvm/llvm-project/pull/229181

When handling stores for `ptr32` address spaces, we didn't check if they were already truncating stores and emitted an ordinary store. This widened `i8` and `i16` writes to 32 bits and could overwrite adjacent fields.

Fix is to build a truncating store that preserves the memory type.

Fixes #228782

Assisted-by: GPT-6 Sol (via VS Code)

>From a8b7e43d310827973093a3476a9470fa505bbf4a Mon Sep 17 00:00:00 2001
From: Daniel Paoliello <danpao at microsoft.com>
Date: Mon, 5 Oct 2026 11:41:54 -0700
Subject: [PATCH] [AArch64] Preserve store width for ptr32 address spaces

When handling stores for `ptr32` address spaces, we didn't check if they were already truncating stores and emitted an ordinary store.
This widened `i8` and `i16` writes to 32 bits and could overwrite adjacent fields.

Fix is to build a truncating store that preserves the memory type.

Fixes #228782

Assisted-by: GPT-6 Sol (via VS Code)
---
 .../Target/AArch64/AArch64ISelLowering.cpp    |  3 ++
 .../AArch64/aarch64-mixed-ptr-sizes.ll        | 51 +++++++++++++++++--
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index ac8700d09491b..24f34fad6c018 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -27951,6 +27951,9 @@ static SDValue performSTORECombine(SDNode *N,
     MVT PtrVT = TLI.getPointerTy(DAG.getDataLayout());
     if (PtrVT != Ptr.getSimpleValueType()) {
       SDValue Cast = DAG.getAddrSpaceCast(DL, PtrVT, Ptr, AddrSpace, 0);
+      if (ST->isTruncatingStore())
+        return DAG.getTruncStore(Chain, DL, Value, Cast, MemVT,
+                                 ST->getMemOperand());
       return DAG.getStore(Chain, DL, Value, Cast, ST->getPointerInfo(),
                           ST->getBaseAlign(), ST->getMemOperand()->getFlags(),
                           ST->getAAInfo());
diff --git a/llvm/test/CodeGen/AArch64/aarch64-mixed-ptr-sizes.ll b/llvm/test/CodeGen/AArch64/aarch64-mixed-ptr-sizes.ll
index 74e0aaf752bbb..96a72ce726c26 100644
--- a/llvm/test/CodeGen/AArch64/aarch64-mixed-ptr-sizes.ll
+++ b/llvm/test/CodeGen/AArch64/aarch64-mixed-ptr-sizes.ll
@@ -1,7 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s | FileCheck %s
-; RUN: llc --fast-isel < %s | FileCheck %s
-; RUN: llc --global-isel --global-isel-abort=2 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-GI
+; RUN: llc < %s | FileCheck %s --check-prefixes=CHECK,NARROW
+; RUN: llc --fast-isel < %s | FileCheck %s --check-prefixes=CHECK,NARROW
+; RUN: llc --global-isel --global-isel-abort=2 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-GI,NARROW
+; RUN: llc -O0 < %s | FileCheck %s --check-prefix=NARROW
 
 ; CHECK-GI:       warning: Instruction selection used fallback path for test_sign_ext
 ; CHECK-GI-NEXT:  warning: Instruction selection used fallback path for test_sign_ext_store_load
@@ -218,5 +219,49 @@ entry:
   ret void
 }
 
+define void @test_store_i8_sptr(ptr addrspace(270) %p, i8 %v) {
+; NARROW-LABEL: test_store_i8_sptr:
+; NARROW:       strb w1, [x{{[0-9]+}}]
+; NARROW-NEXT:  ret
+  store i8 %v, ptr addrspace(270) %p, align 1
+  ret void
+}
+
+define void @test_store_i16_sptr(ptr addrspace(270) %p, i16 %v) {
+; NARROW-LABEL: test_store_i16_sptr:
+; NARROW:       strh w1, [x{{[0-9]+}}]
+; NARROW-NEXT:  ret
+  store i16 %v, ptr addrspace(270) %p, align 2
+  ret void
+}
+
+define void @test_store_i8_uptr(ptr addrspace(271) %p, i8 %v) {
+; NARROW-LABEL: test_store_i8_uptr:
+; NARROW:       strb w1, [x{{[0-9]+}}]
+; NARROW-NEXT:  ret
+  store i8 %v, ptr addrspace(271) %p, align 1
+  ret void
+}
+
+define void @test_store_i16_uptr(ptr addrspace(271) %p, i16 %v) {
+; NARROW-LABEL: test_store_i16_uptr:
+; NARROW:       strh w1, [x{{[0-9]+}}]
+; NARROW-NEXT:  ret
+  store i16 %v, ptr addrspace(271) %p, align 2
+  ret void
+}
+
+define void @test_copy6_uptr(ptr addrspace(271) %dst, ptr %src) {
+; CHECK-LABEL: test_copy6_uptr:
+; CHECK:       str w{{[0-9]+}}, [x{{[0-9]+}}]
+; CHECK-NEXT:  strh w{{[0-9]+}}, [x{{[0-9]+}}]
+; CHECK-NEXT:  ret
+  call void @llvm.memcpy.p271.p0.i64(ptr addrspace(271) align 2 %dst,
+                                     ptr align 2 %src, i64 6, i1 false)
+  ret void
+}
+
+declare void @llvm.memcpy.p271.p0.i64(ptr addrspace(271), ptr, i64, i1 immarg)
+
 attributes #0 = { mustprogress uwtable "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+fp-armv8,+neon,+v8a,-fmv" }
 attributes #1 = { "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+fp-armv8,+neon,+v8a,-fmv" }