[InstCombine] if the condition of a select is known, eliminate the select

[spatel]

Summary

This is a limited solution for PR31512:
https://llvm.org/bugs/show_bug.cgi?id=31512

The motivation is that we will need to increase usage of llvm.assume and/or metadata to solve PR28430:
https://llvm.org/bugs/show_bug.cgi?id=28430

...and this kind of simplification is needed to take advantage of that extra information.

[efriedma]

This seems like a dangerous direction, in the sense that it could make instsimplify very expensive... I mean, if we didn't care about compile-time cost, we could call computeKnownBits on all the operands for every call into instsimplify, and hope something simplifies. Could we put this sort of transform somewhere it isn't called so frequently?

lib/Analysis/InstructionSimplify.cpp

3712 ↗

(On Diff #83169)

There's probably a more straightforward way to check whether the condition is a vector type.

[majnemer]

In D28337#640118, @efriedma wrote:

This seems like a dangerous direction, in the sense that it could make instsimplify very expensive... I mean, if we didn't care about compile-time cost, we could call computeKnownBits on all the operands for every call into instsimplify, and hope something simplifies. Could we put this sort of transform somewhere it isn't called so frequently?

I agree, I'm not sure this is the right place for this.

[spatel]

In D28337#640127, @majnemer wrote:

In D28337#640118, @efriedma wrote:

This seems like a dangerous direction, in the sense that it could make instsimplify very expensive... I mean, if we didn't care about compile-time cost, we could call computeKnownBits on all the operands for every call into instsimplify, and hope something simplifies. Could we put this sort of transform somewhere it isn't called so frequently?

I agree, I'm not sure this is the right place for this.

Ok, thanks for the guidance. I see a couple of ways to limit the cost and still get the motivating case that has an 'llvm.assume':

1. Guard this call to computeKnownBits with a check of the select condition's users (make sure there's an assume in the list).
2. Trigger the fold from an earlier instruction via isKnownNonNullAt(). I think this would involve adding an llvm.assume check that's similar to the change in D20044. This probably also needs the extra isKnownNonNullAt() call from D28204.

Any others? Is one of the above preferable to the other?
Any

[spatel]

I hope I'm not burning everyone out with 'assume' questions, but I should mention that the test in D28485 works because we have a naked call to computeKnownBits() in InstCombiner::visitReturn(). I'm guessing that there are usually less 'ret' insts than 'select' insts, but should that fold be guarded in the same way as we decide here?

// There might be assume intrinsics dominating this return that completely
// determine the value. If so, constant fold it.

[efriedma]

Another option is to handle this in EarlyCSE; we already do a very similar optimization there for branches.

Another option is to do this in instcombine.

[spatel]

Patch updated to limit compile-time cost:

1. Move the fold to InstCombine
2. Guard the fold with a check of the assumption cache

Note that this patch now depends on D28485 to work correctly.

[efriedma]

I think we need a more comprehensive approach to analyzing value ranges/known bits/etc., but this is fine for now, I think.

[spatel]

In D28337#643135, @efriedma wrote:

I think we need a more comprehensive approach to analyzing value ranges/known bits/etc., but this is fine for now, I think.

Thanks! And agreed - if there's consensus about how to fix this (nonnull and llvm.assume are my motivations in this series of patches), I'll try to improve it.

[spatel]

I made the 2nd test a FIXME (removing the dependency on D28485) and committed this at r291915 so we could see if there is any compile-time impact from this on the bots or other testers.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (PBI && PBI->isConditional() &&
         PBI->getCondition(), SI.getCondition(), DL, CondIsFalse);
       if (Implication) {
         Value *V = *Implication ? TrueVal : FalseVal;
         return replaceInstUsesWith(SI, V);
       }
     }
   }
 
+  // If we can compute the condition, there's no need for a select.
+  // Like the above fold, we are attempting to reduce compile-time cost by
+  // putting this fold here with limitations rather than in InstSimplify.
+  // The motivation for this call into value tracking is to take advantage of
+  // the assumption cache, so make sure that is populated.
+  if (!CondVal->getType()->isVectorTy() && !AC.assumptions().empty()) {
+    APInt KnownOne(1, 0), KnownZero(1, 0);
+    computeKnownBits(CondVal, KnownZero, KnownOne, 0, &SI);
+    if (KnownOne == 1)
+      return replaceInstUsesWith(SI, TrueVal);
+    if (KnownZero == 1)
+      return replaceInstUsesWith(SI, FalseVal);
+  }
+
   if (Instruction *BitCastSel = foldSelectCmpBitcasts(SI, *Builder))
     return BitCastSel;
 
   return nullptr;
 }

test/Transforms/InstCombine/select.ll

 ; CHECK-LABEL: @cannot_canonicalize_to_shuffle2(
 ; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> <i1 true, i1 undef, i1 false, i1 icmp sle (i16 ptrtoint (i32* @g to i16), i16 4)>, <4 x i32> %a, <4 x i32> %b
 ; CHECK-NEXT:    ret <4 x i32> [[SEL]]
 ;
   %sel = select <4 x i1> <i1 true, i1 undef, i1 false, i1 icmp sle (i16 ptrtoint (i32* @g to i16), i16 4)>, <4 x i32> %a, <4 x i32> %b
   ret <4 x i32> %sel
 }
 
+declare void @llvm.assume(i1)
+
+define i8 @assume_cond_true(i1 %cond, i8 %x, i8 %y) {
+; CHECK-LABEL: @assume_cond_true(
+; CHECK-NEXT:    call void @llvm.assume(i1 %cond)
+; CHECK-NEXT:    ret i8 %x
+;
+  call void @llvm.assume(i1 %cond)
+  %sel = select i1 %cond, i8 %x, i8 %y
+  ret i8 %sel
+}
+
+define i8 @assume_cond_false(i1 %cond, i8 %x, i8 %y) {
+; CHECK-LABEL: @assume_cond_false(
+; CHECK-NEXT:    [[NOTCOND:%.*]] = xor i1 %cond, true
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NOTCOND]])
+; CHECK-NEXT:    ret i8 %y
+;
+  %notcond = icmp eq i1 %cond, false
+  call void @llvm.assume(i1 %notcond)
+  %sel = select i1 %cond, i8 %x, i8 %y
+  ret i8 %sel
+}
+

test/Transforms/InstSimplify/select.ll

 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32* %x, i32* %y
 ; CHECK-NEXT:    ret i32* [[SEL]]
 ;
   %cmp = icmp slt i32* %x, null
   %sel = select i1 %cmp, i32* %x, i32* %y
   ret i32* %sel
 }
 
-; FIXME: If the condition is known, we don't need to select.
+; If the condition is known, we don't need to select, but we're not
+; doing this fold here to avoid compile-time cost.
 
 declare void @llvm.assume(i1)
 
 define i8 @assume_sel_cond(i1 %cond, i8 %x, i8 %y) {
 ; CHECK-LABEL: @assume_sel_cond(
 ; CHECK-NEXT:    call void @llvm.assume(i1 %cond)
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 %cond, i8 %x, i8 %y
 ; CHECK-NEXT:    ret i8 [[SEL]]