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

[Sanjay Patel via Phabricator via llvm-commits]

spatel created this revision.
spatel added reviewers: xbolva00, nikic, lebedev.ri.
Herald added subscribers: hiraditya, mcrosier.
Herald added a project: LLVM.

As mentioned in D72643 <https://reviews.llvm.org/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.


https://reviews.llvm.org/D72958

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


Index: llvm/test/Transforms/InstSimplify/select.ll
===================================================================
--- llvm/test/Transforms/InstSimplify/select.ll
+++ llvm/test/Transforms/InstSimplify/select.ll
@@ -59,8 +59,7 @@
 
 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 <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_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
Index: llvm/lib/Analysis/InstructionSimplify.cpp
===================================================================
--- llvm/lib/Analysis/InstructionSimplify.cpp
+++ llvm/lib/Analysis/InstructionSimplify.cpp
@@ -4014,6 +4014,39 @@
   if (isa<UndefValue>(FalseVal))   // select ?, X, undef -> X
     return TrueVal;
 
+  // Deal with partial undef vector constants: select ?, VecC, VecC' --> VecC''
+  Constant *TrueC, *FalseC;
+  if (match(TrueVal, m_Constant(TrueC)) &&
+      match(FalseVal, m_Constant(FalseC)) &&
+      TrueC->isElementWiseEqual(FalseC)) {
+    assert(TrueC->getType()->isVectorTy() && FalseC->getType()->isVectorTy() &&
+           "Expected vector constants");
+
+    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)) {
+        assert(!isa<UndefValue>(FEltC) && "Expected non-undef constant");
+        NewC.push_back(FEltC);
+      } else {
+        assert(!isa<UndefValue>(TEltC) && "Expected non-undef constant");
+        NewC.push_back(TEltC);
+      }
+    }
+    if (NewC.size() == NumElts)
+      return ConstantVector::get(NewC);
+  }
+
   if (Value *V =
           simplifySelectWithICmpCond(Cond, TrueVal, FalseVal, Q, MaxRecurse))
     return V;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D72958.238880.patch
Type: text/x-patch
Size: 3625 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200117/0f89da30/attachment.bin>