[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

6126

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.

6144

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.

6512

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

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

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

6512

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

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

6517

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

6517

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

public:
   /// of Values, using DefaultValue to fill any holes in the table.
   SwitchLookupTable(
       Module &M, uint64_t TableSize, ConstantInt *Offset,
       const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
       Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName);
 
   /// Build instructions with Builder to retrieve the value at
   /// the position given by Index in the lookup table.
-  Value *BuildLookup(Value *Index, IRBuilder<> &Builder);
+  Value *BuildLookup(Value *Index, IRBuilder<> &Builder, bool DefaultIsReachable);
 
   /// Return true if a table with TableSize elements of
   /// type ElementType would fit in a target-legal register.
   static bool WouldFitInRegister(const DataLayout &DL, uint64_t TableSize,
                                  Type *ElementType);
 
 private:
   // Depending on the contents of the table, it can be represented in
Array = new GlobalVariable(M, ArrayTy, /*isConstant=*/true,
                              "switch.table." + FuncName);
   Array->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
   // Set the alignment to that of an array items. We will be only loading one
   // value out of it.
   Array->setAlignment(DL.getPrefTypeAlign(ValueType));
   Kind = ArrayKind;
 }
 
-Value *SwitchLookupTable::BuildLookup(Value *Index, IRBuilder<> &Builder) {
+Value *SwitchLookupTable::BuildLookup(Value *Index, IRBuilder<> &Builder,
+                                      bool DefaultIsReachable) {
   switch (Kind) {
   case SingleValueKind:
     return SingleValue;
   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");
+      Result = Builder.CreateMul(Result, LinearMultiplier, "switch.idx.mult",
+                                 /*HasNUW =*/false,
+                                 /*HasNSW =*/!DefaultIsReachable);
+
     if (!LinearOffset->isZero())
-      Result = Builder.CreateAdd(Result, LinearOffset, "switch.offset");
+      Result = Builder.CreateAdd(Result, LinearOffset, "switch.offset",
+                                 /*HasNUW =*/false,
+                                 /*HasNSW =*/!DefaultIsReachable);
+
     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.
     ShiftAmt = Builder.CreateMul(
         ShiftAmt, ConstantInt::get(MapTy, BitMapElementTy->getBitWidth()),
-        "switch.shiftamt");
+        "switch.shiftamt",
+        /*HasNUW =*/false,
+        /*HasNSW =*/!DefaultIsReachable);
 
     // Shift down.
     Value *DownShifted =
         Builder.CreateLShr(BitMap, ShiftAmt, "switch.downshift");
     // Mask off.
     return Builder.CreateTrunc(DownShifted, BitMapElementTy, "switch.masked");
   }
   case ArrayKind: {
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");
+    TableIndex = Builder.CreateSub(SI->getCondition(), TableIndexOffset,
+                                   "switch.tableidx", /*HasNUW =*/false,
+                                   /*HasNSW =*/!DefaultIsReachable);
   }
 
-  // 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.
   } else {
for (PHINode *PHI : PHIs) {
     const ResultListTy &ResultList = ResultLists[PHI];
 
     // If using a bitmask, use any value to fill the lookup table holes.
     Constant *DV = NeedMask ? ResultLists[PHI][0].second : DefaultResults[PHI];
     StringRef FuncName = Fn->getName();
     SwitchLookupTable Table(Mod, TableSize, TableIndexOffset, ResultList, DV,
                             DL, FuncName);
 
-    Value *Result = Table.BuildLookup(TableIndex, Builder);
+    Value *Result = Table.BuildLookup(TableIndex, Builder, DefaultIsReachable);
 
     // Do a small peephole optimization: re-use the switch table compare if
     // possible.
     if (!TableHasHoles && HasDefaultResults && RangeCheckBranch) {
       BasicBlock *PhiBlock = PHI->getParent();
       // Search for compare instructions which use the phi.
       for (auto *User : PHI->users()) {
         reuseTableCompare(User, PhiBlock, RangeCheckBranch, DV, ResultList);

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -switch-to-lookup < %s -mtriple=x86_64-apple-darwin12.0.0 | FileCheck %s
 ; RUN: opt -S -passes='simplifycfg<switch-to-lookup>' < %s -mtriple=x86_64-apple-darwin12.0.0 | FileCheck %s
 
 ; rdar://17735071
 target datalayout = "e-p:64:64:64-S128-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f16:16:16-f32:32:32-f64:64:64-f128:128:128-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 target triple = "x86_64-apple-darwin12.0.0"
 
 define i64 @_TFO6reduce1E5toRawfS0_FT_Si(i2) {
 ; CHECK-LABEL: @_TFO6reduce1E5toRawfS0_FT_Si(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i2 [[TMP0:%.*]], -2
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i2 [[TMP0:%.*]], -2
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX_ZEXT:%.*]] = zext i2 [[SWITCH_TABLEIDX]] to i3
 ; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [4 x i64], ptr @switch.table._TFO6reduce1E5toRawfS0_FT_Si, i32 0, i3 [[SWITCH_TABLEIDX_ZEXT]]
 ; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i64, ptr [[SWITCH_GEP]], align 8
 ; CHECK-NEXT:    ret i64 [[SWITCH_LOAD]]
 ;
 entry:
   switch i2 %0, label %1 [
   i2 0, label %2

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

return:
   %storemerge = phi i1 [ true, %sw.bb ], [ false, %sw.default ]
   ret i1 %storemerge
 }
 
 define i32 @signed_overflow1(i8 %n) {
 ; CHECK-LABEL: @signed_overflow1(
 ; CHECK-NEXT:  start:
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i2
-; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i2 [[TRUNC]], -2
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i2 [[TRUNC]], -2
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX_ZEXT:%.*]] = zext i2 [[SWITCH_TABLEIDX]] to i3
 ; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [4 x i32], ptr @switch.table.signed_overflow1, i32 0, i3 [[SWITCH_TABLEIDX_ZEXT]]
 ; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
 ; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
 ;
 start:
   %trunc = trunc i8 %n to i2
   switch i2 %trunc, label %bb1 [
bb6:                                              ; preds = %start, %bb3, %bb4, %bb5
   %.sroa.0.0 = phi i32 [ 4444, %bb5 ], [ 3333, %bb4 ], [ 2222, %bb3 ]
   ret i32 %.sroa.0.0
 }
 
 define i32 @signed_overflow_negative(i8 %n) {
 ; CHECK-LABEL: @signed_overflow_negative(
 ; CHECK-NEXT:  start:
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[N:%.*]] to i2
-; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i2 [[TRUNC]], -2
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i2 [[TRUNC]], -2
 ; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i2 [[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:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[SWITCH_IDX_CAST]], 1111
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], 1111
 ; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 start:
   %trunc = trunc i8 %n to i2
   switch i2 %trunc, label %bb1 [
   i2 0, label %bb6
   i2 1, label %bb3
   i2 -2, label %bb4
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
 }