[InstCombine] fold select X, (ext X), C

[spatel]

Summary

If we're going to canonicalize IR towards select of constants, try harder to create those.
Also, don't lose the metadata.

[spatel]

Ping.

[majnemer]

LGTM

[bjope]

lib/Transforms/InstCombine/InstCombineSelect.cpp

966 ↗

(On Diff #73134)

Why aren't you using the Builder here? Is it to avoid an implicit fold by the CreateSelect?
To me it looks like you are using a different style compared to the optimization above involving the trunc.

And I guess that if you use the Builder, then you won't need to explicitly copyMetadata? Or is that missing from the optimization involving the trunc. There was not a single explicit copyMetadata in this file until now.

[majnemer]

lib/Transforms/InstCombine/InstCombineSelect.cpp

966 ↗

(On Diff #73134)

The code added is in the style of the rest of InstCombine. The builder is used for intermediate instructions but not the final instruction which is RAUW. This has to do with the particular behavior InstCombine has augmented the builder with.

I'm pretty sure the copyMetadata is around to make sure we don't discard the !prof. Other select combines probably drop this information.

[spatel]

lib/Transforms/InstCombine/InstCombineSelect.cpp

966 ↗

(On Diff #73134)

Correct about the !prof - although hopefully we're not dropping that anymore. In the trunc transform above, it is copied with:

Builder->CreateSelect(Cond, X, TruncCVal, "narrow", &Sel)  <-- last param means copy metadata from Sel.

[bjope]

lib/Transforms/InstCombine/InstCombineSelect.cpp

966 ↗

(On Diff #73134)

Ok. Nothing wrong with the style then. Sorry about that comment.

And yes, the Builder->CreateSelect is copying metadata if the original select is provided. So in many situations the metadata will survive. That is kind of why I wondered why there was an explicit copy of metadata instead of using the Builder in this partiular case.

(I won't learn if I don't ask questions. Thanks for you answers!)

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

Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel) {
 
   Constant *C;
   if (!match(Sel.getTrueValue(), m_Constant(C)) &&
       !match(Sel.getFalseValue(), m_Constant(C)))
     return nullptr;
 
   // If the constant is the same after truncation to the smaller type and
   // extension to the original type, we can narrow the select.
+  Value *Cond = Sel.getCondition();
   Type *SelType = Sel.getType();
   Constant *TruncC = ConstantExpr::getTrunc(C, SmallType);
   Constant *ExtC = ConstantExpr::getCast(ExtOpcode, TruncC, SelType);
   if (ExtC == C) {
-    Value *Cond = Sel.getCondition();
     Value *TruncCVal = cast<Value>(TruncC);
     if (ExtInst == Sel.getFalseValue())
       std::swap(X, TruncCVal);
 
     // select Cond, (ext X), C --> ext(select Cond, X, C')
     // select Cond, C, (ext X) --> ext(select Cond, C', X)
     Value *NewSel = Builder->CreateSelect(Cond, X, TruncCVal, "narrow", &Sel);
     return CastInst::Create(Instruction::CastOps(ExtOpcode), NewSel, SelType);
   }
 
+  // If one arm of the select is the extend of the condition, replace that arm
+  // with the extension of the appropriate known bool value.
+  if (Cond == X) {
+    SelectInst *NewSel;
+    if (ExtInst == Sel.getTrueValue()) {
+      // select X, (sext X), C --> select X, -1, C
+      // select X, (zext X), C --> select X,  1, C
+      Constant *One = ConstantInt::getTrue(SmallType);
+      Constant *AllOnesOrOne = ConstantExpr::getCast(ExtOpcode, One, SelType);
+      NewSel = SelectInst::Create(Cond, AllOnesOrOne, C);
+    } else {
+      // select X, C, (sext X) --> select X, C, 0
+      // select X, C, (zext X) --> select X, C, 0
+      Constant *Zero = ConstantInt::getNullValue(SelType);
+      NewSel = SelectInst::Create(Cond, C, Zero);
+    }
+    NewSel->copyMetadata(Sel);
+    return NewSel;
+  }
+
   return nullptr;
 }
 
 /// Try to transform a vector select with a constant condition vector into a
 /// shuffle for easier combining with other shuffles and insert/extract.
 static Instruction *canonicalizeSelectToShuffle(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Constant *CondC;

llvm/trunk/test/Transforms/InstCombine/select-bitext.ll

 ;
   %ccax = zext <2 x i1> %cca to <2 x i32>
   %r = select i1 %ccb, <2 x i32> %ccax, <2 x i32> <i32 0, i32 0>
   ret <2 x i32> %r
 }
 
 define i32 @sext_true_val_must_be_all_ones(i1 %x) {
 ; CHECK-LABEL: @sext_true_val_must_be_all_ones(
-; CHECK-NEXT:    [[EXT:%.*]] = sext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 [[EXT]], i32 42, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 -1, i32 42, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = sext i1 %x to i32
   %sel = select i1 %x, i32 %ext, i32 42, !prof !0
   ret i32 %sel
 }
 
 define <2 x i32> @sext_true_val_must_be_all_ones_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @sext_true_val_must_be_all_ones_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = sext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> [[EXT]], <2 x i32> <i32 42, i32 12>, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 -1, i32 -1>, <2 x i32> <i32 42, i32 12>, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = sext <2 x i1> %x to <2 x i32>
   %sel = select <2 x i1> %x, <2 x i32> %ext, <2 x i32> <i32 42, i32 12>, !prof !0
   ret <2 x i32> %sel
 }
 
 define i32 @zext_true_val_must_be_one(i1 %x) {
 ; CHECK-LABEL: @zext_true_val_must_be_one(
-; CHECK-NEXT:    [[EXT:%.*]] = zext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 [[EXT]], i32 42, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 1, i32 42, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = zext i1 %x to i32
   %sel = select i1 %x, i32 %ext, i32 42, !prof !0
   ret i32 %sel
 }
 
 define <2 x i32> @zext_true_val_must_be_one_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @zext_true_val_must_be_one_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = zext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> [[EXT]], <2 x i32> <i32 42, i32 12>, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 1, i32 1>, <2 x i32> <i32 42, i32 12>, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = zext <2 x i1> %x to <2 x i32>
   %sel = select <2 x i1> %x, <2 x i32> %ext, <2 x i32> <i32 42, i32 12>, !prof !0
   ret <2 x i32> %sel
 }
 
 define i32 @sext_false_val_must_be_zero(i1 %x) {
 ; CHECK-LABEL: @sext_false_val_must_be_zero(
-; CHECK-NEXT:    [[EXT:%.*]] = sext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 0, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = sext i1 %x to i32
   %sel = select i1 %x, i32 42, i32 %ext, !prof !0
   ret i32 %sel
 }
 
 define <2 x i32> @sext_false_val_must_be_zero_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @sext_false_val_must_be_zero_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = sext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> zeroinitializer, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = sext <2 x i1> %x to <2 x i32>
   %sel = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> %ext, !prof !0
   ret <2 x i32> %sel
 }
 
 define i32 @zext_false_val_must_be_zero(i1 %x) {
 ; CHECK-LABEL: @zext_false_val_must_be_zero(
-; CHECK-NEXT:    [[EXT:%.*]] = zext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 0, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = zext i1 %x to i32
   %sel = select i1 %x, i32 42, i32 %ext, !prof !0
   ret i32 %sel
 }
 
 define <2 x i32> @zext_false_val_must_be_zero_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @zext_false_val_must_be_zero_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = zext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> zeroinitializer, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = zext <2 x i1> %x to <2 x i32>
   %sel = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> %ext, !prof !0
   ret <2 x i32> %sel
 }
 
 !0 = !{!"branch_weights", i32 3, i32 5}