[SimplifyCFG] add nsw on SwitchToLookupTable Index calculation

[khei4]

Summary

This is the first step to resolve the following issue, no-signed wrap flags propagation on SwitchToLookupTable.
https://github.com/llvm/llvm-project/issues/59671

motivated transformation: https://alive2.llvm.org/ce/z/GZs_dj

1. add nsw/nuw on SwitchToLookupTable in SimplifyCFG
   
   a. add nsw on SwitchToLookupTable index calculation on MinCaseVal subtraction <- this revision will be
   
   https://alive2.llvm.org/ce/z/WSFRK5
   
   b. add nuw/nsw on BuildLookuptable BitMap shiftwidth calculation https://reviews.llvm.org/D150838
   
   https://alive2.llvm.org/ce/z/5Cmc72
   
   c. add nsw on BuildLookuptable LinearMap calculation https://reviews.llvm.org/D150943
   
   https://alive2.llvm.org/ce/z/cKBQfs

2. preserve nsw on InstCombine for the cases. https://reviews.llvm.org/D152088

transformation: https://alive2.llvm.org/ce/z/ZWmZLv

[khei4]

update test (without confirmation)

[khei4]

change plan to use Unreachabiliity of DefaultDest. (not yet change on test except X86)

[khei4]

reuse DefaultIsReachable and fix tests

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6122

Looking at how LinearMultiplier and LinearOffset are generated in SwitchLookupTable, I don't think that always setting NSW is correct. Nothing in the calculation suggests that it can't overflow. I believe we would have to track this separate when these values are computed. I also think that DefaultIsReachable doesn't matter for this part: At this point the range of the input values has already been checked.

6139

I believe in this case both NUW and NSW can always be set. We know that this multiply must be < BitWidth to be a valid shift amount, so it can't overflow in either signed or unsigned sense.

6504–6505

Add something like // If the default is unreachable, all case values are s>= MinCaseVal, so the subtraction cannot wrap in a signed sense.

This part of the change looks correct to me. It might be worthwhile to split it from the changes to BuildLookup.

[khei4]

@nikic Thank you for the review!

This part of the change looks correct to me. It might be worthwhile to split it from the changes to BuildLookup.

As you suggested, I'll first focus on flags only by DeafultsReachable outside of BuildLookup on this revision!

[khei4]

temporary remove changes in BuildLookup

[khei4]

update non-X86 platform test, narrow revision scope

[nikic]

llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll

12 ↗

(On Diff #523248)

This is a miscompile: E.g. if the input was 0 then we'll do 0 - (-2) which should be 2 but overflows to -2.

The bit I missed here is that we're effectively mapping a signed range to an unsigned one, doing something like -2 -> 0, -1 -> 1, 0 -> 2, 1 -> 3, but of course the last two become -2 and -1 when interpreted as signed.

I think we need an additional check that TableSize is small enough (should not be larger than half the integer space, or something like that).

[nikic]

llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll

12 ↗

(On Diff #523248)

Probably the simplest way to check this is to do something like MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue()) and check that it doesn't overflow. That directly corresponds to the operation we're doing...

[khei4]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6504–6505

Add something like // If the default is unreachable, all case values are s>= MinCaseVal, so the subtraction cannot wrap in a signed sense.

Thanks! I'll basically took this comment as it is

llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll

12 ↗

(On Diff #523248)

It seems right! Thank you for a good catch!

Probably the simplest way to check this is to do something like MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue()) and check that it doesn't overflow.

Yeah, it seems like repetitions though, seem simpler than
1 << (TableIndexOffset->getBitWidth() - 1) < TableSize - 1. :)

[khei4]

fix to check signed overflow directly. (may need to be update tests

[nikic]

Looks like this no longer affects any of the existing tests. Could you please add a test where this adds the nsw flag? Preferably two tests, one directly before overflow happens, and one directly after.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6507

CanBeWrapped -> CanWrap?

[khei4]

Thank you for the review!

Looks like this no longer affects any of the existing tests. Could you please add a test where this adds the nsw flag? Preferably two tests, one directly before overflow happens, and one directly after.

Sure!

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6507

Thanks! MayWrap might be correct :)

[khei4]

add tests.

[khei4]

change flag name and remove obsolete comment.

[nikic]

LGTM

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

1812

signd -> signed

[khei4]

fix typo

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

case LinearMapKind: {
     // Derive the result value from the input value.
     Value *Result = Builder.CreateIntCast(Index, LinearMultiplier->getType(),
                                           false, "switch.idx.cast");
     if (!LinearMultiplier->isOne())
       Result = Builder.CreateMul(Result, LinearMultiplier, "switch.idx.mult");
     if (!LinearOffset->isZero())
       Result = Builder.CreateAdd(Result, LinearOffset, "switch.offset");
     return Result;
   }
   case BitMapKind: {
     // Type of the bitmap (e.g. i59).
     IntegerType *MapTy = BitMap->getType();
 
     // Cast Index to the same type as the bitmap.
     // Note: The Index is <= the number of elements in the table, so
     // truncating it to the width of the bitmask is safe.
     Value *ShiftAmt = Builder.CreateZExtOrTrunc(Index, MapTy, "switch.cast");
 
     // Multiply the shift amount by the element width. NUW/NSW can always be
     // set, because WouldFitInRegister guarantees Index * ShiftAmt is in
     // BitMap's bit width.
     ShiftAmt = Builder.CreateMul(
         ShiftAmt, ConstantInt::get(MapTy, BitMapElementTy->getBitWidth()),
         "switch.shiftamt",/*HasNUW =*/true,/*HasNSW =*/true);
 
     // Shift down.
     Value *DownShifted =
         Builder.CreateLShr(BitMap, ShiftAmt, "switch.downshift");
     // Mask off.
     return Builder.CreateTrunc(DownShifted, BitMapElementTy, "switch.masked");
   }
   case ArrayKind: {
     // Make sure the table index will not overflow when treated as signed.
     IntegerType *IT = cast<IntegerType>(Index->getType());
if (!DL.fitsInLegalInteger(TableSize))
       return false;
   }
 
   if (!ShouldBuildLookupTable(SI, TableSize, TTI, DL, ResultTypes))
     return false;
 
   std::vector<DominatorTree::UpdateType> Updates;
 
+  // Compute the maximum table size representable by the integer type we are
+  // switching upon.
+  unsigned CaseSize = MinCaseVal->getType()->getPrimitiveSizeInBits();
+  uint64_t MaxTableSize = CaseSize > 63 ? UINT64_MAX : 1ULL << CaseSize;
+  assert(MaxTableSize >= TableSize &&
+         "It is impossible for a switch to have more entries than the max "
+         "representable value of its input integer type's size.");
+
+  // If the default destination is unreachable, or if the lookup table covers
+  // all values of the conditional variable, branch directly to the lookup table
+  // BB. Otherwise, check that the condition is within the case range.
+  const bool DefaultIsReachable =
+      !isa<UnreachableInst>(SI->getDefaultDest()->getFirstNonPHIOrDbg());
+  const bool GeneratingCoveredLookupTable = (MaxTableSize == TableSize);
+
   // Create the BB that does the lookups.
   Module &Mod = *CommonDest->getParent()->getParent();
   BasicBlock *LookupBB = BasicBlock::Create(
       Mod.getContext(), "switch.lookup", CommonDest->getParent(), CommonDest);
 
   // Compute the table index value.
   Builder.SetInsertPoint(SI);
   Value *TableIndex;
   ConstantInt *TableIndexOffset;
   if (UseSwitchConditionAsTableIndex) {
     TableIndexOffset = ConstantInt::get(MaxCaseVal->getType(), 0);
     TableIndex = SI->getCondition();
   } else {
     TableIndexOffset = MinCaseVal;
-    TableIndex =
-        Builder.CreateSub(SI->getCondition(), TableIndexOffset, "switch.tableidx");
+    // If the default is unreachable, all case values are s>= MinCaseVal. Then
+    // we can try to attach nsw.
+    bool MayWrap = true;
+    if (!DefaultIsReachable)
+      MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue(), MayWrap);
+
+    TableIndex = Builder.CreateSub(SI->getCondition(), TableIndexOffset,
+                                   "switch.tableidx", /*HasNUW =*/false,
+                                   /*HasNSW =*/!MayWrap);
   }
 
-  // Compute the maximum table size representable by the integer type we are
-  // switching upon.
-  unsigned CaseSize = MinCaseVal->getType()->getPrimitiveSizeInBits();
-  uint64_t MaxTableSize = CaseSize > 63 ? UINT64_MAX : 1ULL << CaseSize;
-  assert(MaxTableSize >= TableSize &&
-         "It is impossible for a switch to have more entries than the max "
-         "representable value of its input integer type's size.");
-
-  // If the default destination is unreachable, or if the lookup table covers
-  // all values of the conditional variable, branch directly to the lookup table
-  // BB. Otherwise, check that the condition is within the case range.
-  const bool DefaultIsReachable =
-      !isa<UnreachableInst>(SI->getDefaultDest()->getFirstNonPHIOrDbg());
-  const bool GeneratingCoveredLookupTable = (MaxTableSize == TableSize);
   BranchInst *RangeCheckBranch = nullptr;
 
   if (!DefaultIsReachable || GeneratingCoveredLookupTable) {
     Builder.CreateBr(LookupBB);
     if (DTU)
       Updates.push_back({DominatorTree::Insert, BB, LookupBB});
     // Note: We call removeProdecessor later since we need to be able to get the
     // PHI value for the default case in case we're using a bit mask.

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

 
 bb5:                                              ; preds = %start
   br label %bb6
 
 bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
   %.sroa.0.0 = phi i32 [ 2222, %bb5 ], [ 1111, %bb4 ], [ 4444, %bb3 ], [ 3333, %start ]
   ret i32 %.sroa.0.0
 }
+
+; can attach nsw because of default's unreachability and case value type size
+; is big enough to hold max signed value.
+define i32 @nsw_on_index_sub(i8 %n) {
+; CHECK-LABEL: @nsw_on_index_sub(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i3
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i3 [[TRUNC]], -2
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i32
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i32 [[SWITCH_IDX_CAST]], 1111
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[SWITCH_IDX_MULT]], 1111
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
+;
+start:
+  %trunc = trunc i8 %n to i3
+  ; we can hold max(cases) - min(cases) = 0b001 - 0b110 = 0b011
+  switch i3 %trunc, label %bb1 [
+  i3 0, label %bb6 ; 0b000
+  i3 1, label %bb3 ; 0b001
+  i3 -2, label %bb4; 0b110
+  i3 -1, label %bb5; 0b111
+  ]
+
+bb1:                                              ; preds = %start
+  unreachable
+
+bb3:                                              ; preds = %start
+  br label %bb6
+
+bb4:                                              ; preds = %start
+  br label %bb6
+
+bb5:                                              ; preds = %start
+  br label %bb6
+
+bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
+  %.sroa.0.0 = phi i32 [ 2222, %bb5 ], [ 1111, %bb4 ], [ 4444, %bb3 ], [ 3333, %start ]
+  ret i32 %.sroa.0.0
+}
+
+; cannot attach nsw on tableidx subtraction
+; because maximum index is greater than INT3_MAX = 3
+define i32 @wrapped_on_index_sub(i8 %n) {
+; CHECK-LABEL: @wrapped_on_index_sub(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i3
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TRUNC]], -2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i3 [[SWITCH_TABLEIDX]], -3
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i32
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i32 [[SWITCH_IDX_CAST]], 1111
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[SWITCH_IDX_MULT]], 1111
+; CHECK-NEXT:    [[DOTSROA_0_0:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 0
+; CHECK-NEXT:    ret i32 [[DOTSROA_0_0]]
+;
+start:
+  %trunc = trunc i8 %n to i3
+  switch i3 %trunc, label %bb1 [
+  i3 0, label %bb6
+  i3 1, label %bb3
+  i3 2, label %bb7
+  i3 -2, label %bb4
+  i3 -1, label %bb5
+  ]
+
+bb1:                                              ; preds = %start
+  br label %bb6
+
+bb3:                                              ; preds = %start
+  br label %bb6
+
+bb4:                                              ; preds = %start
+  br label %bb6
+
+bb5:                                              ; preds = %start
+  br label %bb6
+
+bb7:                                              ; preds = %start
+  br label %bb6
+
+bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
+  %.sroa.0.0 = phi i32 [ 0, %bb1 ], [ 2222, %bb5 ], [ 1111, %bb4 ], [ 5555, %bb7 ], [ 4444, %bb3 ], [ 3333, %start ]
+  ret i32 %.sroa.0.0
+}
+
+
+; cannot attach nsw on tableidx subtraction
+; because default reachability shows overflow could happen.
+define i32 @wrapped_on_index_sub_default(i8 %n) {
+; CHECK-LABEL: @wrapped_on_index_sub_default(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i3
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TRUNC]], -2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i3 [[SWITCH_TABLEIDX]], -4
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i32
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i32 [[SWITCH_IDX_CAST]], 1111
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[SWITCH_IDX_MULT]], 1111
+; CHECK-NEXT:    [[DOTSROA_0_0:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 0
+; CHECK-NEXT:    ret i32 [[DOTSROA_0_0]]
+;
+start:
+  %trunc = trunc i8 %n to i3
+  switch i3 %trunc, label %bb1 [
+  i3 0, label %bb6
+  i3 1, label %bb3
+  i3 -2, label %bb4
+  i3 -1, label %bb5
+  ]
+
+bb1:                                              ; preds = %start
+  br label %bb6
+
+bb3:                                              ; preds = %start
+  br label %bb6
+
+bb4:                                              ; preds = %start
+  br label %bb6
+
+bb5:                                              ; preds = %start
+  br label %bb6
+
+bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
+  %.sroa.0.0 = phi i32 [ 0, %bb1 ], [ 2222, %bb5 ], [ 1111, %bb4 ], [ 4444, %bb3 ], [ 3333, %start ]
+  ret i32 %.sroa.0.0
+}