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

[spatel]

Summary

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.

[nikic]

llvm/lib/Analysis/InstructionSimplify.cpp

4044

I'm having a bit of a hard time convincing myself that this is always correct. In particular, the relationship between isElementWiseEqual(), which is based on ConstantExpr::getICmp() and this code, which is based on simple identity comparison, is not entirely clear. If the vector includes complex constant expressions, and ConstantExpr::getICmp() manages to fold that code while the identity comparison fails, this would trigger an assertion failure here. Or is there an invariant that this cannot happen?

[spatel]

llvm/lib/Analysis/InstructionSimplify.cpp

4044

I've tried to concoct an example of this:

@g = global i32 0

define <3 x i32> @sel_constant_expression(<3 x i1> %cond) {
  %r = select <3 x i1> %cond, <3 x i32> <i32 ptrtoint (i32* @g to i32), i32 0, i32 1>, <3 x i32> <i32 bitcast (<2 x i16> bitcast (i32 ptrtoint (i32* @g to i32) to <2 x i16>) to i32), i32 0, i32 undef>
  ret <3 x i32> %r
}

...but the 2nd ConstantExpr get constant folded into the reduced form of the 1st (the bitcasts go away). So we can either leave this as-is, or we could bail out on ConstantExpr to be safe(r)? I can't imagine that this will ever make a difference in any real-world code.

[nikic]

llvm/lib/Analysis/InstructionSimplify.cpp

4044

I wasn't able to come up with an example either.

I do wonder though if this code shouldn't just forgo isElementWiseEqual() entirely and just cover all four cases in this if/else chain instead? I.e. equality, undef op1, undef op2, else abort the transform. Going through isElementWiseEqual() doesn't really seem to make this code simpler.

[spatel]

Patch updated:
Reduce code - don't bother with isElementWiseEqual(). If elements don't match, break and carry on with other simplifications.

llvm/lib/Analysis/InstructionSimplify.cpp

4044

Good point. We have to walk through the vector elements anyway, so we can save time and space by doing it all in one place.

[nikic]

LGTM

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifySelectInst(Value *Cond, Value *TrueVal, Value *FalseVal,
   if (TrueVal == FalseVal)
     return TrueVal;
 
   if (isa<UndefValue>(TrueVal))   // select ?, undef, X -> X
     return 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;
 
   if (Value *V = simplifySelectWithFCmp(Cond, TrueVal, FalseVal, Q))
     return V;
 
   if (Value *V = foldSelectWithBinaryOp(Cond, TrueVal, FalseVal))

llvm/test/Transforms/InstSimplify/select.ll

 ; CHECK-NEXT:    ret <2 x i32> [[X:%.*]]
 ;
   %V = select <2 x i1> %cond, <2 x i32> %x, <2 x i32> %x
   ret <2 x i32> %V
 }
 
 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
 }
 
 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
 }
 
 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
 }
 
 define <2 x i8> @vsel_tvec(<2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @vsel_tvec(
 ; CHECK-NEXT:    ret <2 x i8> [[X:%.*]]