[JumpThreading] Don't restrict cast-traversal to i1

[loladiro]

Summary

In D17663, JumpThreading learned to look trough simple cast instructions,
but only if the source of those cast instructions was a phi/cmp i1
(in an effort to limit compile time effects). I think this condition
is too restrictive. For switches with limited value range, InstCombine
will readily introduce an extra trunc instruction to a smaller
integer type (e.g. from i8 to i2), leaving us in the somewhat perverse
situation that jump-threading would work before running instcombine,
but not after. Since instcombine produces this pattern, I think we
need to consider it canonical and support it in JumpThreading.
In general, for limiting recursion, I think the existing restriction
to phi and cmp nodes should be sufficient to avoid looking through
unprofitable chains of instructions.

[haicheng]

Hi Keno,

Sorry for the late response. What you said totally makes sense to me. Do you have any performance and/or compilation data to share?

Haicheng

[loladiro]

This happens fairly frequently for us in julia, because we use i8 to represent tags for tagged unions and then switch on them to discover which of the union components is active.
Without this patch I saw ~3x performance regressions in the worst benchmark (because it prevented other optimizations from seeing optimization opportunities based on the
active union component). I did not see compilation performance impact in my testing.

[haicheng]

LGTM

lib/Transforms/Scalar/JumpThreading.cpp

for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
           Result.push_back(std::make_pair(KC, PN->getIncomingBlock(i)));
       }
     }
 
     return !Result.empty();
   }
 
   // Handle Cast instructions.  Only see through Cast when the source operand is
-  // PHI or Cmp and the source type is i1 to save the compilation time.
+  // PHI or Cmp to save the compilation time.
   if (CastInst *CI = dyn_cast<CastInst>(I)) {
     Value *Source = CI->getOperand(0);
-    if (!Source->getType()->isIntegerTy(1))
-      return false;
     if (!isa<PHINode>(Source) && !isa<CmpInst>(Source))
       return false;
     ComputeValueKnownInPredecessors(Source, BB, Result, Preference, CxtI);
     if (Result.empty())
       return false;
 
     // Convert the known values.
     for (auto &R : Result)

test/Transforms/JumpThreading/basic.ll

l4:
   ret void
 
 l5:
   call void @k()
   ret void
 ; CHECK: }
 }
 
+define i1 @trunc_switch(i1 %arg) {
+; CHECK-LABEL: @trunc_switch
+top:
+; CHECK: br i1 %arg, label %exitA, label %exitB
+  br i1 %arg, label %common, label %B
+
+B:
+  br label %common
+
+common:
+  %phi = phi i8 [ 2, %B ], [ 1, %top ]
+  %trunc = trunc i8 %phi to i2
+; CHECK-NOT: switch
+  switch i2 %trunc, label %unreach [
+    i2 1, label %exitA
+    i2 -2, label %exitB
+  ]
+
+unreach:
+  unreachable
+
+exitA:
+  ret i1 true
+
+exitB:
+  ret i1 false
+}
+
 ; CHECK-LABEL: define void @h_con(i32 %p) {
 define void @h_con(i32 %p) {
   %x = icmp ult i32 %p, 5
   br i1 %x, label %l1, label %l2
 
 l1:
   call void @j()
   br label %l3