[InstCombine] enhance vector demanded elements to look at a vector select condition operand

[spatel]

Summary

I noticed that we were not back-propagating undef lanes to shuffle masks when we have a shuffle that reduces the vector width. This is part of investigating/solving PR38691:
https://bugs.llvm.org/show_bug.cgi?id=38691

[lebedev.ri]

I don't see any obvious problems here, but i'm unfamiliar with this area at all..
Do we not have any one-use concerns in such transforms?

[RKSimon]

In D51433#1219656, @lebedev.ri wrote:

I don't see any obvious problems here, but i'm unfamiliar with this area at all..
Do we not have any one-use concerns in such transforms?

Its handled at the top of the function

[RKSimon]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1267 ↗

(On Diff #163102)

Won't using UndefElts interfere with the Left/Right cases below?

1275 ↗

(On Diff #163102)

Separate NFC change

[spatel]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1267 ↗

(On Diff #163102)

Yes, good catch. Not sure how to expose that in a test, but I'll fix the code in the next rev.

[spatel]

Patch updated:
Corrected the recursive call on the condition to use its own 'UndefElts' constant (which is unused currently as explained in the new code comments).
I cleaned up the surrounding code with rL341545.

[craig.topper]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1263 ↗

(On Diff #164213)

What would you think about casting to SelectInst and using getCondition, getTrueValue, getFalseValue, setCondition, etc. throughout this code?

[spatel]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1263 ↗

(On Diff #164213)

Sure - let me clean that up and rebase.

[spatel]

Patch updated:
Rebased after cleanup and rL341708.
I made a subtle change here to not use a local Cond variable because that could become stale if we setCond() in the new block.

[craig.topper]

LGTM

llvm/trunk/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (NewUndefElts) {
                                           Shuffle->getMaskValue(i)));
       }
       I->setOperand(2, ConstantVector::get(Elts));
       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);
-    Value *Cond = Sel->getCondition();
+    if (Sel->getCondition()->getType()->isVectorTy()) {
+      // 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.
+      if (Value *V = SimplifyDemandedVectorElts(Sel->getCondition(),
+                                                DemandedElts, UndefElts,
+                                                Depth + 1)) {
+        Sel->setCondition(V);
+        MadeChange = true;
+      }
+    }
 
+    // Next, see if we can transform the arms of the select.
     APInt DemandedLHS(DemandedElts), DemandedRHS(DemandedElts);
-    if (auto *CV = dyn_cast<ConstantVector>(Cond)) {
+    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.
         if (CElt->isNullValue())
           DemandedLHS.clearBit(i);
         else
           DemandedRHS.clearBit(i);
       }
     }
 
     if (Value *V = SimplifyDemandedVectorElts(Sel->getTrueValue(), DemandedLHS,
                                               UndefElts2, Depth + 1)) {
       Sel->setTrueValue(V);
       MadeChange = true;
     }
 
     if (Value *V = SimplifyDemandedVectorElts(Sel->getFalseValue(), DemandedRHS,
                                               UndefElts3, Depth + 1)) {
       Sel->setFalseValue(V);
       MadeChange = true;
     }
 
     // Output elements are undefined if the element from each arm is undefined.
+    // TODO: This can be improved. See comment in select condition handling.
     UndefElts = UndefElts2 & UndefElts3;
     break;
   }
   case Instruction::BitCast: {
     // Vector->vector casts only.
     VectorType *VTy = dyn_cast<VectorType>(I->getOperand(0)->getType());
     if (!VTy) break;
     unsigned InVWidth = VTy->getNumElements();

llvm/trunk/test/Transforms/InstCombine/shuffle-select-narrow.ll

 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %widecmp = shufflevector <2 x i1> %cmp, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
   %widesel = select <4 x i1> %widecmp, <4 x i8> %x, <4 x i8> %y
   %r = shufflevector <4 x i8> %widesel, <4 x i8> undef, <2 x i32> <i32 0, i32 1>
   ret <2 x i8> %r
 }
 
-; TODO: The 1st shuffle is not extending with undefs, but demanded elements should correct that.
+; The 1st shuffle is not extending with undefs, but demanded elements corrects that.
 
 define <2 x i8> @narrow_shuffle_of_select_overspecified_extend(<2 x i1> %cmp, <4 x i8> %x, <4 x i8> %y) {
 ; CHECK-LABEL: @narrow_shuffle_of_select_overspecified_extend(
-; CHECK-NEXT:    [[WIDECMP:%.*]] = shufflevector <2 x i1> [[CMP:%.*]], <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
-; CHECK-NEXT:    [[WIDESEL:%.*]] = select <4 x i1> [[WIDECMP]], <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i8> [[WIDESEL]], <4 x i8> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i8> [[Y:%.*]], <4 x i8> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP:%.*]], <2 x i8> [[TMP1]], <2 x i8> [[TMP2]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %widecmp = shufflevector <2 x i1> %cmp, <2 x i1> undef, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
   %widesel = select <4 x i1> %widecmp, <4 x i8> %x, <4 x i8> %y
   %r = shufflevector <4 x i8> %widesel, <4 x i8> undef, <2 x i32> <i32 0, i32 1>
   ret <2 x i8> %r
 }
 
;
   %r = shufflevector <4 x i8> %widesel, <4 x i8> undef, <3 x i32> <i32 0, i32 1, i32 2>
   ret <3 x i8> %r
 }
 
 ; Negative test - mismatched types would require extra shuffling.
 
 define <3 x i8> @narrow_shuffle_of_select_mismatch_types2(<4 x i1> %cmp, <6 x i8> %x, <6 x i8> %y) {
 ; CHECK-LABEL: @narrow_shuffle_of_select_mismatch_types2(
-; CHECK-NEXT:    [[WIDECMP:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 undef, i32 undef>
+; CHECK-NEXT:    [[WIDECMP:%.*]] = shufflevector <4 x i1> [[CMP:%.*]], <4 x i1> undef, <6 x i32> <i32 0, i32 1, i32 2, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    [[WIDESEL:%.*]] = select <6 x i1> [[WIDECMP]], <6 x i8> [[X:%.*]], <6 x i8> [[Y:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = shufflevector <6 x i8> [[WIDESEL]], <6 x i8> undef, <3 x i32> <i32 0, i32 1, i32 2>
 ; CHECK-NEXT:    ret <3 x i8> [[R]]
 ;
   %widecmp = shufflevector <4 x i1> %cmp, <4 x i1> undef, <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 undef, i32 undef>
   %widesel = select <6 x i1> %widecmp, <6 x i8> %x, <6 x i8> %y
   %r = shufflevector <6 x i8> %widesel, <6 x i8> undef, <3 x i32> <i32 0, i32 1, i32 2>
   ret <3 x i8> %r