[InstCombine] enhance vector demanded elements for select condition operand

[spatel]

Summary

This is part of reducing IR in PR42755:
https://bugs.llvm.org/show_bug.cgi?id=42755

If we can determine that both arms of a given element of a vector select are the same value, then the select's condition operand is not demanded at that element.

Alive2 examples based on the regression tests:
http://volta.cs.utah.edu:8080/z/2PCSs2
http://volta.cs.utah.edu:8080/z/s4hWra
http://volta.cs.utah.edu:8080/z/Jgexop

[efriedma]

I'm concerned that in some cases, the condition might be "demanded" due to poison rules, even if we don't actually care about the value. See, for example, http://volta.cs.utah.edu:8080/z/YDSX7k

[spatel]

In D80658#2058339, @efriedma wrote:

I'm concerned that in some cases, the condition might be "demanded" due to poison rules, even if we don't actually care about the value. See, for example, http://volta.cs.utah.edu:8080/z/YDSX7k

Yes, I overlooked poison in the condition operand itself. Given that, I don't see a way to salvage this approach vs. a pattern-specific fold for the motivating case (reducing a shuffle of the condition operand).
I'll leave this up for a bit in case anyone else sees a way forward.

[spatel]

Another test added with:
rGb874dc4ddabe

Abandoning.

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

 
 #include "InstCombineInternal.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/IntrinsicsX86.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/KnownBits.h"
+#include <numeric>
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "instcombine"
 
 namespace {
 
if (NewUndefElts) {
       MadeChange = true;
     }
     break;
   }
   case Instruction::Select: {
     // If this is a vector select, try to transform the select condition based
     // on the current demanded elements.
     SelectInst *Sel = cast<SelectInst>(I);
-    if (Sel->getCondition()->getType()->isVectorTy()) {
+    if (isa<FixedVectorType>(Sel->getCondition()->getType())) {
+      // Peek through shuffled true/false operands of the select.
+      // If an operand is not shuffled, simulate with an identity shuffle.
+      Value *LHS = Sel->getTrueValue(), *RHS = Sel->getFalseValue();
+      ArrayRef<int> TMask, FMask;
+      SmallVector<int, 16> IdentityMask(VWidth);
+      std::iota(IdentityMask.begin(), IdentityMask.end(), 0);
+      Value *T0, *T1;
+      if (!match(LHS, m_Shuffle(m_Value(T0), m_Value(T1), m_Mask(TMask)))) {
+        TMask = IdentityMask;
+        T0 = LHS;
+        T1 = nullptr;
+      }
+      Value *F0, *F1;
+      if (!match(RHS, m_Shuffle(m_Value(F0), m_Value(F1), m_Mask(FMask)))) {
+        FMask = IdentityMask;
+        F0 = RHS;
+        F1 = nullptr;
+      }
+
+      // If an element is the same in the true/false operands, the corresponding
+      // element of the select condition is unnecessary.
+      APInt DemandedCond(DemandedElts);
+      int NumElts = VWidth;
+      for (int i = 0; i < NumElts; ++i) {
+        // If both operands are undefined at this element, the condition element
+        // does not matter. Poison can't propagate through this element of the
+        // select.
+        if (TMask[i] == UndefMaskElem && FMask[i] == UndefMaskElem) {
+          DemandedCond.clearBit(i);
+          continue;
+        }
+
+        // If only one operand is undefined in the element, we can't change
+        // the condition because the select may block or propagate poison.
+        if (TMask[i] == UndefMaskElem || FMask[i] == UndefMaskElem)
+          continue;
+
+        assert(TMask[i] >= 0 && FMask[i] >= 0 && TMask[i] < NumElts * 2 &&
+               FMask[i] < NumElts * 2 && "Unexpected shuffle mask");
+        Value *TSrc = TMask[i] < NumElts ? T0 : T1;
+        Value *FSrc = FMask[i] < NumElts ? F0 : F1;
+
+        // Are these the same element of the same vector?
+        assert(TSrc && FSrc && "Unexpected shuffle operand");
+        if (TSrc == FSrc && (TMask[i] % NumElts) == (FMask[i] % NumElts))
+          DemandedCond.clearBit(i);
+      }
+
       // TODO: We are not doing anything with UndefElts based on this call.
       // It is overwritten below based on the other select operands. If an
       // element of the select condition is known undef, then we are free to
       // choose the output value from either arm of the select. If we know that
       // one of those values is undef, then the output can be undef.
-      simplifyAndSetOp(I, 0, DemandedElts, UndefElts);
+      simplifyAndSetOp(I, 0, DemandedCond, UndefElts);
     }
 
     // Next, see if we can transform the arms of the select.
     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.

llvm/test/Transforms/InstCombine/vec_demanded_elts.ll

bb:
   ret i32* %tmp2
 }
 
 ; The non-zero elements of the result are always 'min', so the splat is unnecessary.
 
 define <4 x i8> @select_cond_with_eq_true_false_elts(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_eq_true_false_elts(
 ; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[SPLAT]], <4 x i8> [[TVAL]], <4 x i8> [[Y]]
+; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[CMP:%.*]], <4 x i8> [[TVAL]], <4 x i8> [[Y]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
   %splat = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> zeroinitializer
   %r = select <4 x i1> %splat, <4 x i8> %tval, <4 x i8> %y
   ret <4 x i8> %r
 }
 
 ; First element of the result is always x[0], so first element of select condition is unnecessary.
 
 define <4 x i8> @select_cond_with_eq_true_false_elts2(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_eq_true_false_elts2(
 ; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 7>
-; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
+; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 undef, i32 1, i32 0, i32 1>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[COND]], <4 x i8> [[TVAL]], <4 x i8> [[X]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 7>
   %cond = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %r = select <4 x i1> %cond, <4 x i8> %tval, <4 x i8> %x
   ret <4 x i8> %r
 }
 
 ; Second element of the result is always x[3], so second element of select condition is unnecessary.
 ; Fourth element of the result is always undef, so fourth element of select condition is unnecessary.
 
 define <4 x float> @select_cond_with_eq_true_false_elts3(<4 x float> %x, <4 x float> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_eq_true_false_elts3(
 ; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x float> [[X:%.*]], <4 x float> [[Y:%.*]], <4 x i32> <i32 1, i32 3, i32 5, i32 undef>
 ; CHECK-NEXT:    [[FVAL:%.*]] = shufflevector <4 x float> [[Y]], <4 x float> [[X]], <4 x i32> <i32 0, i32 7, i32 6, i32 undef>
-; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 undef, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 undef, i32 undef, i32 2, i32 undef>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[COND]], <4 x float> [[TVAL]], <4 x float> [[FVAL]]
 ; CHECK-NEXT:    ret <4 x float> [[R]]
 ;
   %tval = shufflevector <4 x float> %x, <4 x float> %y, <4 x i32> <i32 1, i32 3, i32 5, i32 undef>
   %fval = shufflevector <4 x float> %y, <4 x float> %x, <4 x i32> <i32 0, i32 7, i32 6, i32 undef>
   %cond = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> <i32 undef, i32 1, i32 2, i32 3>
   %r = select <4 x i1> %cond, <4 x float> %tval, <4 x float> %fval
   ret <4 x float> %r
 }
 
+; Negative test - only 1 side of the select is undef at element 0; can't change the condition.
+
 define <4 x i8> @select_cond_with_undef_true_false_elts(<4 x i8> %x, <4 x i8> %y, <4 x i1> %cmp) {
 ; CHECK-LABEL: @select_cond_with_undef_true_false_elts(
 ; CHECK-NEXT:    [[TVAL:%.*]] = shufflevector <4 x i8> [[Y:%.*]], <4 x i8> undef, <4 x i32> <i32 undef, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    [[COND:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
 ; CHECK-NEXT:    [[R:%.*]] = select <4 x i1> [[COND]], <4 x i8> [[TVAL]], <4 x i8> [[X:%.*]]
 ; CHECK-NEXT:    ret <4 x i8> [[R]]
 ;
   %tval = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 undef, i32 5, i32 6, i32 7>
   %cond = shufflevector <4 x i1> %cmp, <4 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %r = select <4 x i1> %cond, <4 x i8> %tval, <4 x i8> %x
   ret <4 x i8> %r
 }