[InstCombine] Fix to switch narrowing

[Gerolf]

Summary

This fixes an issue with switch narrowing that gets exposed
by r274233 ( https://reviews.llvm.org/D12965 - [InstCombine] shrink
switch conditions better (PR24766)). The problem is that case constants are
narrowed (to, say, width "k"), then the condition like cond = expr + const is
simplified by adding const to the case constants, and finally the new case
constants are narrowed again. However, when the width for the second round
of narrowing increases to, say, "k+n", then cond is a "k" bit expression
which won't match properly any "k+n" wide case constants.
The fix is to do the cond optimization first and only then narrow the case
constants appropriately.

[majnemer]

Adding a switch expert to take a look :)

[hans]

lib/Transforms/InstCombine/InstructionCombining.cpp

2276

What happened to this part?

test/Transforms/InstCombine/narrow-switch.ll

108

Why does this need to be i61 now, and should the test be renamed to reflect that?

[Gerolf]

lib/Transforms/InstCombine/InstructionCombining.cpp

2276

Gone as part of the fix: the condition expression a+const is evaluated at full width before the case constants are inspected.

test/Transforms/InstCombine/narrow-switch.ll

108

The original code evaluates the case constants.
+ It recognizes they can be truncated to 61b.
+ It truncates the case values to 61b
+ It truncates the cond value to 61b
+ It checks the condition a+const in 61b and adjusts the case value - const to 61b.
+ It returns SI, so switch narrowing is invoked again and now finds the case values to be 59b. It is this double adjustment of the case values that potentially can lead to a bug. Obviously the test cases in narrow-switch.ll did not catch it.

The fix evaluates a+const in 64b and changes the case value to value -const (still 64b). Then it finds the case values to be 61b , and adjusts with of the values and cond accordingly.

[hans]

I'm sorry, but I'm having trouble fully understanding exactly what the bug is that this is fixing.

It seems that this patch is regressing the @trunc64to59 test since we're now failing to truncate the condition to i59.

It does seem reasonable to optimize the switch (x + C) before doing truncation, so this is probably the right thing to do, I'd just like to understand it better.

And if we're going to do truncation after we've changed the switch condition, don't we need to move the computeKnownBits part as well? With your patch, we're first computing known bits of the condition, then changing the condition, and then later using those known bits. That seems wrong.

lib/Transforms/InstCombine/InstructionCombining.cpp

2276

Oh, I see now.

[Gerolf]

Hans, you were right. The computeKnowBits etc parts should have been moved also. Sanjay committed the fix proper in r289442. I only kept the regression test in narrow-switch.ll (r293018).

-Gerolf

[sanjoy]

lib/Transforms/InstCombine/InstructionCombining.cpp

2247

Range-for?

lib/Transforms/InstCombine/InstructionCombining.cpp

 Instruction *InstCombiner::visitSwitchInst(SwitchInst &SI) {
   Value *Cond = SI.getCondition();
   unsigned BitWidth = cast<IntegerType>(Cond->getType())->getBitWidth();
   APInt KnownZero(BitWidth, 0), KnownOne(BitWidth, 0);
   computeKnownBits(Cond, KnownZero, KnownOne, 0, &SI);
   unsigned LeadingKnownZeros = KnownZero.countLeadingOnes();
   unsigned LeadingKnownOnes = KnownOne.countLeadingOnes();
 
+  Value *Op0 = nullptr;
+  ConstantInt *AddRHS = nullptr;
+  if (match(Cond, m_Add(m_Value(Op0), m_ConstantInt(AddRHS)))) {
+    // Change 'switch (X+4) case 1:' into 'switch (X) case -3'.
+    for (SwitchInst::CaseIt i = SI.case_begin(), e = SI.case_end(); i != e;
+         ++i) {
+      ConstantInt *CaseVal = i.getCaseValue();
+      Constant *LHS = CaseVal;
+      Constant *NewCaseVal = ConstantExpr::getSub(LHS, AddRHS);
+      assert(isa<ConstantInt>(NewCaseVal) &&
+             "Result of expression should be constant");
+      i.setValue(cast<ConstantInt>(NewCaseVal));
+    }
+    SI.setCondition(Op0);
+    if (auto *CondI = dyn_cast<Instruction>(Cond))
+      Worklist.Add(CondI);
+  }
+
   // Compute the number of leading bits we can ignore.
   // TODO: A better way to determine this would use ComputeNumSignBits().
   for (auto &C : SI.cases()) {
     LeadingKnownZeros = std::min(
         LeadingKnownZeros, C.getCaseValue()->getValue().countLeadingZeros());
     LeadingKnownOnes = std::min(
         LeadingKnownOnes, C.getCaseValue()->getValue().countLeadingOnes());
   }
if (NewWidth > 0 && NewWidth < BitWidth) {
     Value *NewCond = Builder->CreateTrunc(Cond, Ty, "trunc");
     SI.setCondition(NewCond);
 
     for (auto &C : SI.cases())
       static_cast<SwitchInst::CaseIt *>(&C)->setValue(ConstantInt::get(
           SI.getContext(), C.getCaseValue()->getValue().trunc(NewWidth)));
   }
 
-  Value *Op0 = nullptr;
-  ConstantInt *AddRHS = nullptr;
-  if (match(Cond, m_Add(m_Value(Op0), m_ConstantInt(AddRHS)))) {
-    // Change 'switch (X+4) case 1:' into 'switch (X) case -3'.
-    for (SwitchInst::CaseIt i = SI.case_begin(), e = SI.case_end(); i != e;
-         ++i) {
-      ConstantInt *CaseVal = i.getCaseValue();
-      Constant *LHS = CaseVal;
-      if (TruncCond) {
-        LHS = LeadingKnownZeros
-                  ? ConstantExpr::getZExt(CaseVal, Cond->getType())
-                  : ConstantExpr::getSExt(CaseVal, Cond->getType());
-      }
-      Constant *NewCaseVal = ConstantExpr::getSub(LHS, AddRHS);
-      assert(isa<ConstantInt>(NewCaseVal) &&
-             "Result of expression should be constant");
-      i.setValue(cast<ConstantInt>(NewCaseVal));
-    }
-    SI.setCondition(Op0);
-    if (auto *CondI = dyn_cast<Instruction>(Cond))
-      Worklist.Add(CondI);
-    return &SI;
-  }
-
   return TruncCond ? &SI : nullptr;
 }
 
 Instruction *InstCombiner::visitExtractValueInst(ExtractValueInst &EV) {
   Value *Agg = EV.getAggregateOperand();
 
   if (!EV.hasIndices())
     return replaceInstUsesWith(EV, Agg);

test/Transforms/InstCombine/narrow-switch.ll

 }
 
 ; Make sure to avoid assertion crashes and use the type before
 ; truncation to generate the sub constant expressions that leads
 ; to the recomputed condition.
 
 define void @trunc64to59(i64 %a) {
 ; ALL-LABEL: @trunc64to59(
-; ALL:         switch i59
-; ALL-NEXT:    i59 0, label %sw.bb1
-; ALL-NEXT:    i59 18717182647723699, label %sw.bb2
+; ALL:         switch i61
+; ALL-NEXT:    i61 0, label %sw.bb1
+; ALL-NEXT:    i61 18717182647723699, label %sw.bb2
 ; ALL-NEXT:    ]
 ;
 entry:
   %tmp0 = and i64 %a, 15
   %tmp1 = mul i64 %tmp0, -6425668444178048401
   %tmp2 = add i64 %tmp1, 5170979678563097242
   %tmp3 = mul i64 %tmp2, 1627972535142754813
   switch i64 %tmp3, label %sw.default [
   i64 847514119312061490, label %sw.bb1
   i64 866231301959785189, label %sw.bb2
   ]
 
 sw.bb1:
   br label %sw.default
 
 sw.bb2:
   br label %sw.default
 
 sw.default:
   ret void
 }
+
+; Make sure the arithmetic evaluation of the switch
+; condition is evaluated on the original type
+define i32 @trunc32to16(i32 %a0) #0 {
+; ALL-LABEL: @trunc32to16(
+; ALL:         switch i16
+; ALL-NEXT:    i16 63, label %sw.bb
+; ALL-NEXT:    i16 1, label %sw.bb1
+; ALL-NEXT:    i16 100, label %sw.bb2
+; ALL-NEXT:    ]
+;
+entry:
+  %retval = alloca i32, align 4
+  %xor = xor i32 %a0, 1034460917
+  %shr = lshr i32 %xor, 16
+  %add = add i32 %shr, -917677090
+  switch i32 %add, label %sw.epilog [
+    i32 -917677027, label %sw.bb
+    i32 -917677089, label %sw.bb1
+    i32 -917676990, label %sw.bb2
+  ]
+
+sw.bb:                                            ; preds = %entry
+  store i32 90, i32* %retval, align 4
+  br label %return
+
+sw.bb1:                                           ; preds = %entry
+  store i32 91, i32* %retval, align 4
+  br label %return
+
+sw.bb2:                                           ; preds = %entry
+  store i32 92, i32* %retval, align 4
+  br label %return
+
+sw.epilog:                                        ; preds = %entry
+  store i32 113, i32* %retval, align 4
+  br label %return
+
+return:                                           ; preds = %sw.epilog, %sw.bb2,
+  %rval = load i32, i32* %retval, align 4
+  ret i32 %rval
+}
+