[llvm] [SROA] Prevent poison propagation in insertInteger() (PR #109240)

[Junda Liu via llvm-commits]

https://github.com/LLJJDD created https://github.com/llvm/llvm-project/pull/109240

When inserting an new integer to a larger old integer, freeze the old one first
in case it contains poison, otherwise the insertion still results in poison.

>From 241dc39ed196d3175a46d241a00b3b23c1b37825 Mon Sep 17 00:00:00 2001
From: Junda Liu <Junda.Liu at amd.com>
Date: Thu, 19 Sep 2024 12:02:32 +0800
Subject: [PATCH 1/2] [SROA] Precommit test to demonstrate problem

---
 .../Transforms/SROA/poison-propagation.ll     | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 llvm/test/Transforms/SROA/poison-propagation.ll

diff --git a/llvm/test/Transforms/SROA/poison-propagation.ll b/llvm/test/Transforms/SROA/poison-propagation.ll
new file mode 100644
index 00000000000000..94778235466746
--- /dev/null
+++ b/llvm/test/Transforms/SROA/poison-propagation.ll
@@ -0,0 +1,28 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=sroa %s | FileCheck %s
+
+; Negative test: https://alive2.llvm.org/ce/z/xvzHav
+define double @test([2 x i32] %val) {
+; CHECK-LABEL: define double @test(
+; CHECK-SAME: [2 x i32] [[VAL:%.*]]) {
+; CHECK-NEXT:    [[VAL_FCA_0_EXTRACT:%.*]] = extractvalue [2 x i32] [[VAL]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast double poison to i64
+; CHECK-NEXT:    [[PTR_0_INSERT_EXT:%.*]] = zext i32 [[VAL_FCA_0_EXTRACT]] to i64
+; CHECK-NEXT:    [[PTR_0_INSERT_MASK:%.*]] = and i64 [[TMP1]], -4294967296
+; CHECK-NEXT:    [[PTR_0_INSERT_INSERT:%.*]] = or i64 [[PTR_0_INSERT_MASK]], [[PTR_0_INSERT_EXT]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i64 [[PTR_0_INSERT_INSERT]] to double
+; CHECK-NEXT:    [[VAL_FCA_1_EXTRACT:%.*]] = extractvalue [2 x i32] [[VAL]], 1
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast double [[TMP2]] to i64
+; CHECK-NEXT:    [[PTR_4_INSERT_EXT:%.*]] = zext i32 [[VAL_FCA_1_EXTRACT]] to i64
+; CHECK-NEXT:    [[PTR_4_INSERT_SHIFT:%.*]] = shl i64 [[PTR_4_INSERT_EXT]], 32
+; CHECK-NEXT:    [[PTR_4_INSERT_MASK:%.*]] = and i64 [[TMP3]], 4294967295
+; CHECK-NEXT:    [[PTR_4_INSERT_INSERT:%.*]] = or i64 [[PTR_4_INSERT_MASK]], [[PTR_4_INSERT_SHIFT]]
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast i64 [[PTR_4_INSERT_INSERT]] to double
+; CHECK-NEXT:    ret double [[TMP4]]
+;
+  %ptr = alloca double
+  store double poison, ptr %ptr
+  store [2 x i32] %val, ptr %ptr
+  %load = load double, ptr %ptr
+  ret double %load
+}

>From 12f7b930c9dbebede7770efc3d3db559a9e7cb0a Mon Sep 17 00:00:00 2001
From: Junda Liu <Junda.Liu at amd.com>
Date: Thu, 19 Sep 2024 12:16:28 +0800
Subject: [PATCH 2/2] [SROA] Prevent poison propagation in insertInteger()

When inserting an new integer to a larger old integer, freeze the old one first
in case it contains poison, otherwise the insertion still results in poison.
---
 llvm/lib/Transforms/Scalar/SROA.cpp             | 3 +++
 llvm/test/Transforms/SROA/poison-propagation.ll | 6 ++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index d0186da1bc5e22..ace60604a75b7b 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2519,6 +2519,9 @@ static Value *insertInteger(const DataLayout &DL, IRBuilderTy &IRB, Value *Old,
 
   if (ShAmt || Ty->getBitWidth() < IntTy->getBitWidth()) {
     APInt Mask = ~Ty->getMask().zext(IntTy->getBitWidth()).shl(ShAmt);
+    // Freeze old value first in case it contains poison, otherwise the
+    // combination still results in poison.
+    Old = IRB.CreateFreeze(Old);
     Old = IRB.CreateAnd(Old, Mask, Name + ".mask");
     LLVM_DEBUG(dbgs() << "      masked: " << *Old << "\n");
     V = IRB.CreateOr(Old, V, Name + ".insert");
diff --git a/llvm/test/Transforms/SROA/poison-propagation.ll b/llvm/test/Transforms/SROA/poison-propagation.ll
index 94778235466746..e57547df7b547b 100644
--- a/llvm/test/Transforms/SROA/poison-propagation.ll
+++ b/llvm/test/Transforms/SROA/poison-propagation.ll
@@ -8,14 +8,16 @@ define double @test([2 x i32] %val) {
 ; CHECK-NEXT:    [[VAL_FCA_0_EXTRACT:%.*]] = extractvalue [2 x i32] [[VAL]], 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast double poison to i64
 ; CHECK-NEXT:    [[PTR_0_INSERT_EXT:%.*]] = zext i32 [[VAL_FCA_0_EXTRACT]] to i64
-; CHECK-NEXT:    [[PTR_0_INSERT_MASK:%.*]] = and i64 [[TMP1]], -4294967296
+; CHECK-NEXT:    [[TMP6:%.*]] = freeze i64 [[TMP1]]
+; CHECK-NEXT:    [[PTR_0_INSERT_MASK:%.*]] = and i64 [[TMP6]], -4294967296
 ; CHECK-NEXT:    [[PTR_0_INSERT_INSERT:%.*]] = or i64 [[PTR_0_INSERT_MASK]], [[PTR_0_INSERT_EXT]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i64 [[PTR_0_INSERT_INSERT]] to double
 ; CHECK-NEXT:    [[VAL_FCA_1_EXTRACT:%.*]] = extractvalue [2 x i32] [[VAL]], 1
 ; CHECK-NEXT:    [[TMP3:%.*]] = bitcast double [[TMP2]] to i64
 ; CHECK-NEXT:    [[PTR_4_INSERT_EXT:%.*]] = zext i32 [[VAL_FCA_1_EXTRACT]] to i64
 ; CHECK-NEXT:    [[PTR_4_INSERT_SHIFT:%.*]] = shl i64 [[PTR_4_INSERT_EXT]], 32
-; CHECK-NEXT:    [[PTR_4_INSERT_MASK:%.*]] = and i64 [[TMP3]], 4294967295
+; CHECK-NEXT:    [[TMP5:%.*]] = freeze i64 [[TMP3]]
+; CHECK-NEXT:    [[PTR_4_INSERT_MASK:%.*]] = and i64 [[TMP5]], 4294967295
 ; CHECK-NEXT:    [[PTR_4_INSERT_INSERT:%.*]] = or i64 [[PTR_4_INSERT_MASK]], [[PTR_4_INSERT_SHIFT]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = bitcast i64 [[PTR_4_INSERT_INSERT]] to double
 ; CHECK-NEXT:    ret double [[TMP4]]