[llvm-branch-commits] [llvm] release/18.x: [DAGCombiner] In mergeTruncStore, make sure we aren't storing shifted in bits. (#90939) (PR #91038)

[via llvm-branch-commits]

https://github.com/AtariDreams created https://github.com/llvm/llvm-project/pull/91038

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.

(cherry picked from commit 3563af6c06ebc92bcaacef0e33285148ef0f75bd)

>From 9212565cd56b6524f05a05bcafb43b458cb82900 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Fri, 3 May 2024 09:59:33 -0700
Subject: [PATCH] [DAGCombiner] In mergeTruncStore, make sure we aren't storing
 shifted in bits. (#90939)

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.

(cherry picked from commit 3563af6c06ebc92bcaacef0e33285148ef0f75bd)
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp |  4 ++++
 llvm/test/CodeGen/AArch64/pr90936.ll          | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/pr90936.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 5038f8a1fc1562..4951e45edb9ed7 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -8952,6 +8952,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
new file mode 100644
index 00000000000000..38cda8d388945f
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/pr90936.ll
@@ -0,0 +1,20 @@
+; 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:
+; 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
+  %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
+}