[llvm] [DAGCombiner] In mergeTruncStore, make sure we aren't storing shifted in bits. (PR #90939)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/90939

>From 2f0cf4093674070e0f16ec651960464f1be7409d Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 2 May 2024 22:09:59 -0700
Subject: [PATCH 1/2] [AArch64] Add test for #90936. NFC

---
 llvm/test/CodeGen/AArch64/pr90936.ll | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/pr90936.ll

diff --git a/llvm/test/CodeGen/AArch64/pr90936.ll b/llvm/test/CodeGen/AArch64/pr90936.ll
new file mode 100644
index 00000000000000..62c18fa5218765
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/pr90936.ll
@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=aarch64 | FileCheck %s
+
+define void @f(i16 %arg, ptr %arg1) {
+; CHECK-LABEL: f:
+; CHECK:       // %bb.0: // %bb
+; CHECK-NEXT:    strh w0, [x1]
+; CHECK-NEXT:    ret
+bb:
+  %i = trunc i16 %arg to i8
+  %i2 = trunc i16 %arg to i14
+  %i3 = lshr i14 %i2, 8
+  store i8 %i, ptr %arg1, align 1
+  %i4 = getelementptr i8, ptr %arg1, i64 1
+  %i5 = trunc i14 %i3 to i8
+  store i8 %i5, ptr %i4, align 1
+  ret void
+}

>From 2069914d7b235ab0708bc0d209c929af52d37711 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 2 May 2024 22:16:50 -0700
Subject: [PATCH 2/2] [DAGCombiner] In mergeTruncStore, make sure we aren't
 storing shifted in bits.

When looking through a right shift, we need to make sure that all of
the bits we are using from the shift come from the shift input and
not the sign or zero bits that are shifted in.

Fixes #90936.
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 4 ++++
 llvm/test/CodeGen/AArch64/pr90936.ll          | 6 ++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index c0bbea16a64262..fc6bbc119d3c1b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -8840,6 +8840,10 @@ SDValue DAGCombiner::mergeTruncStores(StoreSDNode *N) {
       if (ShiftAmtC % NarrowNumBits != 0)
         return SDValue();
 
+      // Make sure we aren't reading bits that are shifted in.
+      if (ShiftAmtC > WideVal.getScalarValueSizeInBits() - NarrowNumBits)
+        return SDValue();
+
       Offset = ShiftAmtC / NarrowNumBits;
       WideVal = WideVal.getOperand(0);
     }
diff --git a/llvm/test/CodeGen/AArch64/pr90936.ll b/llvm/test/CodeGen/AArch64/pr90936.ll
index 62c18fa5218765..38cda8d388945f 100644
--- a/llvm/test/CodeGen/AArch64/pr90936.ll
+++ b/llvm/test/CodeGen/AArch64/pr90936.ll
@@ -3,8 +3,10 @@
 
 define void @f(i16 %arg, ptr %arg1) {
 ; CHECK-LABEL: f:
-; CHECK:       // %bb.0: // %bb
-; CHECK-NEXT:    strh w0, [x1]
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ubfx w8, w0, #8, #6
+; CHECK-NEXT:    strb w0, [x1]
+; CHECK-NEXT:    strb w8, [x1, #1]
 ; CHECK-NEXT:    ret
 bb:
   %i = trunc i16 %arg to i8