From 2438988d3a7082abd4e4b9362dae1cc23751f664 Mon Sep 17 00:00:00 2001
From: Jorge Botto <jorge.botto.16@ucl.ac.uk>
Date: Thu, 8 Aug 2024 18:29:05 +0100
Subject: [PATCH] Fixing miscompilation due to undef picking the wrong branch

---
 .../InstCombine/InstCombineSimplifyDemanded.cpp       | 11 +++--------
 llvm/test/Transforms/InstCombine/pr98435.ll           |  3 ++-
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index c494fec84c1e6e..153d8c238ed4b6 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -1735,17 +1735,12 @@ Value *InstCombinerImpl::SimplifyDemandedVectorElts(Value *V,
     APInt DemandedLHS(DemandedElts), DemandedRHS(DemandedElts);
     if (auto *CV = dyn_cast<ConstantVector>(Sel->getCondition())) {
       for (unsigned i = 0; i < VWidth; i++) {
-        // isNullValue() always returns false when called on a ConstantExpr.
-        // Skip constant expressions to avoid propagating incorrect information.
         Constant *CElt = CV->getAggregateElement(i);
-        if (isa<ConstantExpr>(CElt))
-          continue;
-        // TODO: If a select condition element is undef, we can demand from
-        // either side. If one side is known undef, choosing that side would
-        // propagate undef.
+
+        // isNullValue() always returns false when called on a ConstantExpr.
         if (CElt->isNullValue())
           DemandedLHS.clearBit(i);
-        else
+        else if (CElt->isOneValue())
           DemandedRHS.clearBit(i);
       }
     }
diff --git a/llvm/test/Transforms/InstCombine/pr98435.ll b/llvm/test/Transforms/InstCombine/pr98435.ll
index 97ac4b987f4fc0..b400801d342fe7 100644
--- a/llvm/test/Transforms/InstCombine/pr98435.ll
+++ b/llvm/test/Transforms/InstCombine/pr98435.ll
@@ -4,7 +4,8 @@
 define <2 x i1> @pr98435(<2 x i1> %val) {
 ; CHECK-LABEL: define <2 x i1> @pr98435(
 ; CHECK-SAME: <2 x i1> [[VAL:%.*]]) {
-; CHECK-NEXT:    ret <2 x i1> <i1 poison, i1 true>
+; CHECK-NEXT:    [[VAL1:%.*]] = select <2 x i1> <i1 undef, i1 true>, <2 x i1> <i1 true, i1 true>, <2 x i1> [[VAL]]
+; CHECK-NEXT:    ret <2 x i1> [[VAL1]]
 ;
   %val1 = select <2 x i1> <i1 undef, i1 true>, <2 x i1> <i1 true, i1 true>, <2 x i1> %val
   ret <2 x i1> %val1