[SimplifyCFG] Find the smallest table considering overflow in `switchToLookupTable`

[DianQK]

Summary

Closes https://github.com/llvm/llvm-project/issues/64231.

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6440

I don't like that the whole fold gets invoked twice with this flag flipped. Can we just pick the Min/Max that result in the smallest table size (i.e. Max-Min) here?

Ideally we would not even base this on signed/unsigned at all, because the smallest range might be both signed *and* unsigned wrapping. Consider something like -128, -1, 0, 127. This would pick [-128, 127] as the signed range, [0, -1] as the unsigned range, but a minimal range would be [127, 0] or [-1, -128]. This would be based around the largest gap between the case value. Though just signed/unsigned is probably the main useful case.

[DianQK]

Try to find a minimal range.
Some test cases got better and some got worse.

[DianQK]

I need to check these test case changes carefully.

I ran llvm-lit with alive. The result is:

Failed Tests (2):
  LLVM :: Transforms/SimplifyCFG/fold-branch-to-common-dest.ll
  LLVM :: Transforms/SimplifyCFG/speculate-call.ll

I haven't focused on whether this is an issue with alive or in itself, but it doesn't seem to be related to this patch.

I believe @signed_overflow1, @signed_overflow2, and @signed_overflow_negative have all changed for the better. @signed_overflow1 and @signed_overflow_negative exchange results. But since we prefer to start at 0, this reduces the number of sub/add instructions by one.
@covered_switch_with_bit_tests also becomes better, smaller table and fewer commands.
@f_i8_128 is our new test. @_TFO6reduce1E5toRawfS0_FT_Si becomes better.

@f_min_max_2, @f_min_max, @test, and @coveredswitch_test look to become worse. But their table size did get smaller.
I rechecked @f_min_max_2, @f_min_max and @test. They become from table with holes to table without holes. It's a "perfect" table?
I've also found that they luckily have a max and min value. Then we created a cover table.
If we want to deal with this situation, maybe we should add a condition to extend to a cover table.

@coveredswitch_test also luckily have a max and min value.

[DianQK]

Rebase.

[DianQK]

Maybe ready for a new review.

These test cases that went worse just happened to have the max and min values of bit width, and then created a cover table with holes.
I think I should create another patch that adds a threshold condition to grow into a cover table to handle this.

[khei4]

I'm still not sure about the correctness of this. It seems like only setting the BeginCaseVal and EndCaseVal to the appropriate values is enough to handle overflowing switch cases. Test transformations themselves look correct!

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6443

I think it's ok to use auto on this.

[DianQK]

In D156612#4647819, @khei4 wrote:

I'm still not sure about the correctness of this. It seems like only setting the BeginCaseVal and EndCaseVal to the appropriate values is enough to handle overflowing switch cases. Test transformations themselves look correct!

Could you explain more about the overflow?
If you're referring to RangeOverflow in the code, this is a special handling to avoid redundant calculation.
Or do you mean that we can always find the smallest table whose End minus Begin does not overflow? Yes. But we need to introduce an overflow offset calculation.
For example, [122, -128]([122, 128]) in the @f_i8_128 case.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6443

Thanks, if Nikita will continue to review it here, I'll change it directly. Otherwise I'll change it when I move to GitHub.

[DianQK]

In D156612#4647819, @khei4 wrote:

I'm still not sure about the correctness of this. It seems like only setting the BeginCaseVal and EndCaseVal to the appropriate values is enough to handle overflowing switch cases. Test transformations themselves look correct!

I seem to understand what you mean. Yes, this patch is looking for the appropriate BeginCaseVal and EndCaseVal.

[DianQK]

Migrated to https://github.com/llvm/llvm-project/pull/67885.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

if (DefaultConst == FalseConst) {
     CmpInst->replaceAllUsesWith(InvertedTableCmp);
     ++NumTableCmpReuses;
   }
 }
 
 /// If the switch is only used to initialize one or more phi nodes in a common
 /// successor block with different constant values, replace the switch with
 /// lookup tables.
-static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
+static bool switchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
                                 DomTreeUpdater *DTU, const DataLayout &DL,
-                                const TargetTransformInfo &TTI) {
+                                const TargetTransformInfo &TTI,
+                                bool UnsignedComparison) {
   assert(SI->getNumCases() > 1 && "Degenerate switch?");
 
   BasicBlock *BB = SI->getParent();
   Function *Fn = BB->getParent();
   // Only build lookup table when we have a target that supports it or the
   // attribute is not set.
   if (!TTI.shouldBuildLookupTables() ||
       (Fn->getFnAttribute("no-jump-tables").getValueAsBool()))
static bool switchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   SmallDenseMap<PHINode *, ResultListTy> ResultLists;
 
   SmallDenseMap<PHINode *, Constant *> DefaultResults;
   SmallDenseMap<PHINode *, Type *> ResultTypes;
   SmallVector<PHINode *, 4> PHIs;
 
   for (SwitchInst::CaseIt E = SI->case_end(); CI != E; ++CI) {
     ConstantInt *CaseVal = CI->getCaseValue();
+    if (UnsignedComparison) {
+      if (CaseVal->getValue().ult(MinCaseVal->getValue()))
+        MinCaseVal = CaseVal;
+      if (CaseVal->getValue().ugt(MaxCaseVal->getValue()))
+        MaxCaseVal = CaseVal;
+    } else {
       if (CaseVal->getValue().slt(MinCaseVal->getValue()))
         MinCaseVal = CaseVal;
       if (CaseVal->getValue().sgt(MaxCaseVal->getValue()))
         MaxCaseVal = CaseVal;
+    }
 
     // Resulting value at phi nodes for this case value.
     using ResultsTy = SmallVector<std::pair<PHINode *, Constant *>, 4>;
     ResultsTy Results;
     if (!getCaseResults(SI, CaseVal, CI->getCaseSuccessor(), &CommonDest,
                         Results, DL, TTI))
       return false;
 
     // Append the result from this case to the list for each phi.
     for (const auto &I : Results) {
       PHINode *PHI = I.first;
static bool switchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   if (UseSwitchConditionAsTableIndex) {
     TableIndexOffset = ConstantInt::get(MaxCaseVal->getType(), 0);
     TableIndex = SI->getCondition();
   } else {
     TableIndexOffset = MinCaseVal;
     // If the default is unreachable, all case values are s>= MinCaseVal. Then
     // we can try to attach nsw.
     bool MayWrap = true;
-    if (!DefaultIsReachable) {
+    if (!DefaultIsReachable && !UnsignedComparison) {
       APInt Res = MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue(), MayWrap);
       (void)Res;
     }
 
     TableIndex = Builder.CreateSub(SI->getCondition(), TableIndexOffset,
                                    "switch.tableidx", /*HasNUW =*/false,
                                    /*HasNSW =*/!MayWrap);
   }
if (Options.ForwardSwitchCondToPhi && ForwardSwitchConditionToPHI(SI))
     return requestResimplify();
 
   // The conversion from switch to lookup tables results in difficult-to-analyze
   // code and makes pruning branches much harder. This is a problem if the
   // switch expression itself can still be restricted as a result of inlining or
   // CVP. Therefore, only apply this transformation during late stages of the
   // optimisation pipeline.
   if (Options.ConvertSwitchToLookupTable &&
-      SwitchToLookupTable(SI, Builder, DTU, DL, TTI))
+      switchToLookupTable(SI, Builder, DTU, DL, TTI, false))
+    return requestResimplify();
+  if (Options.ConvertSwitchToLookupTable &&
+      switchToLookupTable(SI, Builder, DTU, DL, TTI, true))
     return requestResimplify();
 
   if (ReduceSwitchRange(SI, Builder, DL, TTI))
     return requestResimplify();
 
   return false;
 }
 

llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll

 sw.bb8: br label %return
 sw.default: br label %return
 return:
   %retval.0 = phi i8 [ 15, %sw.default ], [ 84, %sw.bb8 ], [ 33, %sw.bb7 ], [ 1, %sw.bb6 ], [ 62, %sw.bb5 ], [ 27, %sw.bb4 ], [ -1, %sw.bb3 ], [ 0, %sw.bb2 ], [ 123, %sw.bb1 ], [ 55, %entry ]
   ret i8 %retval.0
 
 }
 
+; The maximum value can be `-128`(`128`) when considering only indexed ones.
+
+define i32 @f_i8_128(i8 %c) {
+; CHECK-LABEL: @f_i8_128(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i8 [[C:%.*]], 122
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i8 [[SWITCH_TABLEIDX]], 7
+; CHECK-NEXT:    br i1 [[TMP0]], label [[SWITCH_LOOKUP:%.*]], label [[RETURN:%.*]]
+; CHECK:       switch.lookup:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.f_i8_128, i32 0, i8 [[SWITCH_TABLEIDX]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       return:
+; CHECK-NEXT:    [[RETVAL_0:%.*]] = phi i32 [ [[SWITCH_LOAD]], [[SWITCH_LOOKUP]] ], [ 15, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret i32 [[RETVAL_0]]
+;
+entry:
+  switch i8 %c, label %sw.default [
+  i8 122, label %return
+  i8 123, label %sw.bb1
+  i8 124, label %sw.bb2
+  i8 125, label %sw.bb3
+  i8 126, label %sw.bb4
+  i8 127, label %sw.bb5
+  i8 -128, label %sw.bb6
+  ]
+
+sw.bb1: br label %return
+sw.bb2: br label %return
+sw.bb3: br label %return
+sw.bb4: br label %return
+sw.bb5: br label %return
+sw.bb6: br label %return
+sw.default: br label %return
+return:
+  %retval.0 = phi i32 [ 15, %sw.default ], [ 1, %sw.bb6 ], [ 62, %sw.bb5 ], [ 27, %sw.bb4 ], [ -1, %sw.bb3 ], [ 0, %sw.bb2 ], [ 123, %sw.bb1 ], [ 55, %entry ]
+  ret i32 %retval.0
+}
+
 ; A switch used to initialize two variables, an i8 and a float.
 
 declare void @dummy(i8 signext, float)
 define void @h(i32 %x) {
 ; CHECK-LABEL: @h(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X:%.*]], 4
 ; CHECK-NEXT:    br i1 [[TMP0]], label [[SWITCH_LOOKUP:%.*]], label [[SW_EPILOG:%.*]]
bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
   %.sroa.0.0 = phi i32 [ 4444, %bb5 ], [ 3333, %bb4 ], [ 2222, %bb3 ], [ 1111, %start ]
   ret i32 %.sroa.0.0
 }
 
 define i32 @signed_overflow2(i8 %n) {
 ; CHECK-LABEL: @signed_overflow2(
 ; CHECK-NEXT:  start:
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i2
-; CHECK-NEXT:    switch i2 [[TRUNC]], label [[BB1:%.*]] [
-; CHECK-NEXT:    i2 1, label [[BB6:%.*]]
-; CHECK-NEXT:    i2 -2, label [[BB4:%.*]]
-; CHECK-NEXT:    i2 -1, label [[BB5:%.*]]
-; CHECK-NEXT:    ]
-; CHECK:       bb1:
-; CHECK-NEXT:    unreachable
-; CHECK:       bb4:
-; CHECK-NEXT:    br label [[BB6]]
-; CHECK:       bb5:
-; CHECK-NEXT:    br label [[BB6]]
-; CHECK:       bb6:
-; CHECK-NEXT:    [[DOTSROA_0_0:%.*]] = phi i32 [ 4444, [[BB5]] ], [ 3333, [[BB4]] ], [ 2222, [[START:%.*]] ]
-; CHECK-NEXT:    ret i32 [[DOTSROA_0_0]]
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i2 [[TRUNC]], 1
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i2 [[SWITCH_TABLEIDX]] to i32
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[SWITCH_IDX_CAST]], 1111
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], 2222
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 start:
   %trunc = trunc i8 %n to i2
   switch i2 %trunc, label %bb1 [
   i2 1, label %bb3
   i2 -2, label %bb4
   i2 -1, label %bb5
   ]