[llvm] da9c93f - [InstSimplify] fold select of vector constants that include undef elements

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2020-01-20T08:48:32-05:00
New Revision: da9c93f330e0e4bb12972ee7c67229c36943a0c1

URL: https://github.com/llvm/llvm-project/commit/da9c93f330e0e4bb12972ee7c67229c36943a0c1
DIFF: https://github.com/llvm/llvm-project/commit/da9c93f330e0e4bb12972ee7c67229c36943a0c1.diff

LOG: [InstSimplify] fold select of vector constants that include undef elements

As mentioned in D72643, we'd like to be able to assert that any select
of equivalent constants has been removed before we're deep into InstCombine.

But there's a loophole in that assertion for vectors with undef elements
that don't match exactly.

This patch should close that gap. If we have undefs, we can't safely
propagate those unless both constants elements for that lane are undef.

Differential Revision: https://reviews.llvm.org/D72958

Added: 
    

Modified: 
    llvm/lib/Analysis/InstructionSimplify.cpp
    llvm/test/Transforms/InstSimplify/select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index d7510c899101..be24e6fc6da6 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -4014,6 +4014,34 @@ static Value *SimplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
   if (isa<UndefValue>(FalseVal))   // select ?, X, undef -> X
     return TrueVal;
 
+  // Deal with partial undef vector constants: select ?, VecC, VecC' --> VecC''
+  Constant *TrueC, *FalseC;
+  if (TrueVal->getType()->isVectorTy() && match(TrueVal, m_Constant(TrueC)) &&
+      match(FalseVal, m_Constant(FalseC))) {
+    unsigned NumElts = TrueC->getType()->getVectorNumElements();
+    SmallVector<Constant *, 16> NewC;
+    for (unsigned i = 0; i != NumElts; ++i) {
+      // Bail out on incomplete vector constants.
+      Constant *TEltC = TrueC->getAggregateElement(i);
+      Constant *FEltC = FalseC->getAggregateElement(i);
+      if (!TEltC || !FEltC)
+        break;
+
+      // If the elements match (undef or not), that value is the result. If only
+      // one element is undef, choose the defined element as the safe result.
+      if (TEltC == FEltC)
+        NewC.push_back(TEltC);
+      else if (isa<UndefValue>(TEltC))
+        NewC.push_back(FEltC);
+      else if (isa<UndefValue>(FEltC))
+        NewC.push_back(TEltC);
+      else
+        break;
+    }
+    if (NewC.size() == NumElts)
+      return ConstantVector::get(NewC);
+  }
+
   if (Value *V =
           simplifySelectWithICmpCond(Cond, TrueVal, FalseVal, Q, MaxRecurse))
     return V;

diff  --git a/llvm/test/Transforms/InstSimplify/select.ll b/llvm/test/Transforms/InstSimplify/select.ll
index 47a57f80a863..b3f215a44754 100644
--- a/llvm/test/Transforms/InstSimplify/select.ll
+++ b/llvm/test/Transforms/InstSimplify/select.ll
@@ -59,8 +59,7 @@ define <2 x i32> @equal_arms_vec(<2 x i1> %cond, <2 x i32> %x) {
 
 define <2 x i32> @equal_arms_vec_undef(<2 x i1> %cond) {
 ; CHECK-LABEL: @equal_arms_vec_undef(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> <i32 42, i32 undef>, <2 x i32> <i32 undef, i32 42>
-; CHECK-NEXT:    ret <2 x i32> [[V]]
+; CHECK-NEXT:    ret <2 x i32> <i32 42, i32 42>
 ;
   %V = select <2 x i1> %cond, <2 x i32> <i32 42, i32 undef>, <2 x i32> <i32 undef, i32 42>
   ret <2 x i32> %V
@@ -68,8 +67,7 @@ define <2 x i32> @equal_arms_vec_undef(<2 x i1> %cond) {
 
 define <3 x float> @equal_arms_vec_less_undef(<3 x i1> %cond) {
 ; CHECK-LABEL: @equal_arms_vec_less_undef(
-; CHECK-NEXT:    [[V:%.*]] = select <3 x i1> [[COND:%.*]], <3 x float> <float 4.200000e+01, float undef, float 4.300000e+01>, <3 x float> <float 4.200000e+01, float 4.200000e+01, float 4.300000e+01>
-; CHECK-NEXT:    ret <3 x float> [[V]]
+; CHECK-NEXT:    ret <3 x float> <float 4.200000e+01, float 4.200000e+01, float 4.300000e+01>
 ;
   %V = select <3 x i1> %cond, <3 x float> <float 42.0, float undef, float 43.0>, <3 x float> <float 42.0, float 42.0, float 43.0>
   ret <3 x float> %V
@@ -77,8 +75,7 @@ define <3 x float> @equal_arms_vec_less_undef(<3 x i1> %cond) {
 
 define <3 x float> @equal_arms_vec_more_undef(<3 x i1> %cond) {
 ; CHECK-LABEL: @equal_arms_vec_more_undef(
-; CHECK-NEXT:    [[V:%.*]] = select <3 x i1> [[COND:%.*]], <3 x float> <float 4.200000e+01, float undef, float undef>, <3 x float> <float undef, float undef, float 4.300000e+01>
-; CHECK-NEXT:    ret <3 x float> [[V]]
+; CHECK-NEXT:    ret <3 x float> <float 4.200000e+01, float undef, float 4.300000e+01>
 ;
   %V = select <3 x i1> %cond, <3 x float> <float 42.0, float undef, float undef>, <3 x float> <float undef, float undef, float 43.0>
   ret <3 x float> %V