[llvm] [SystemZ] Fix bitwidth problem in FindReplicatedImm() (NFC). (PR #115383)

[Jonas Paulsson via llvm-commits]

https://github.com/JonPsson1 created https://github.com/llvm/llvm-project/pull/115383

A test case emerged with an i32 truncating store of an i64 constant operand, where the i64 constant did not fit in 32 bits, which caused FindReplicatedImm() to crash.

As SystemZVectorConstantInfo() does not really care about the bitwidth of its initial APInt, it seems simple enough to just not recreate it with a limited bitwidth (corresponding to MemVT or Splat VT), but rather just pass the APInt of the ConstantSDNode directly instead.


>From eab4adb0ae5f7478d8d853287a86ea6cf58f77d3 Mon Sep 17 00:00:00 2001
From: Jonas Paulsson <paulson1 at linux.ibm.com>
Date: Thu, 7 Nov 2024 16:11:29 -0600
Subject: [PATCH] Fix bitwidth problem in FindReplicatedImm().

---
 .../Target/SystemZ/SystemZISelLowering.cpp    |  8 +++---
 llvm/test/CodeGen/SystemZ/dag-combine-07.ll   | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/CodeGen/SystemZ/dag-combine-07.ll

diff --git a/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp b/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
index 3999b54de81b65..ed170bae6aec30 100644
--- a/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
+++ b/llvm/lib/Target/SystemZ/SystemZISelLowering.cpp
@@ -7218,12 +7218,12 @@ SDValue SystemZTargetLowering::combineSTORE(
 
     // Find a replicated immediate and return it if found in Word and its
     // type in WordVT.
-    auto FindReplicatedImm = [&](ConstantSDNode *C, unsigned TotBytes) {
+    auto FindReplicatedImm = [&](ConstantSDNode *C) {
       // Some constants are better handled with a scalar store.
       if (C->getAPIntValue().getBitWidth() > 64 || C->isAllOnes() ||
           isInt<16>(C->getSExtValue()) || MemVT.getStoreSize() <= 2)
         return;
-      SystemZVectorConstantInfo VCI(APInt(TotBytes * 8, C->getZExtValue()));
+      SystemZVectorConstantInfo VCI(C->getAPIntValue());
       if (VCI.isVectorConstantLegal(Subtarget) &&
           VCI.Opcode == SystemZISD::REPLICATE) {
         Word = DAG.getConstant(VCI.OpVals[0], SDLoc(SN), MVT::i32);
@@ -7261,12 +7261,12 @@ SDValue SystemZTargetLowering::combineSTORE(
         DAG.isSplatValue(Op1, true/*AllowUndefs*/)) {
       SDValue SplatVal = Op1->getOperand(0);
       if (auto *C = dyn_cast<ConstantSDNode>(SplatVal))
-        FindReplicatedImm(C, SplatVal.getValueType().getStoreSize());
+        FindReplicatedImm(C);
       else
         FindReplicatedReg(SplatVal);
     } else {
       if (auto *C = dyn_cast<ConstantSDNode>(Op1))
-        FindReplicatedImm(C, MemVT.getStoreSize());
+        FindReplicatedImm(C);
       else
         FindReplicatedReg(Op1);
     }
diff --git a/llvm/test/CodeGen/SystemZ/dag-combine-07.ll b/llvm/test/CodeGen/SystemZ/dag-combine-07.ll
new file mode 100644
index 00000000000000..11359009c2899b
--- /dev/null
+++ b/llvm/test/CodeGen/SystemZ/dag-combine-07.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=s390x-linux-gnu -mcpu=z16 | FileCheck %s
+;
+; Test that SystemZTargetLowering::combineSTORE() does not crash on a
+; truncated immediate.
+
+ at G1 = external global i64, align 8
+ at G2 = external global i64, align 8
+
+define void @func_5(ptr %Dst) {
+; CHECK-LABEL: func_5:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lgrl %r1, G2 at GOT
+; CHECK-NEXT:    llihl %r0, 50
+; CHECK-NEXT:    oill %r0, 2
+; CHECK-NEXT:    stg %r0, 0(%r1)
+; CHECK-NEXT:    lgrl %r1, G1 at GOT
+; CHECK-NEXT:    stg %r0, 0(%r1)
+; CHECK-NEXT:    st %r0, 0(%r2)
+; CHECK-NEXT:    br %r14
+  store i64 214748364802, ptr @G2, align 8
+  store i64 214748364802, ptr @G1, align 8
+  %1 = load i32, ptr getelementptr inbounds (i8, ptr @G1, i64 4), align 4
+  store i32 %1, ptr %Dst, align 4
+  ret void
+}