[SimplifyCFG] add nsw on BuildLookuptable LinearMap calculation

[khei4]

Summary

1-c step of https://reviews.llvm.org/D146903

We can attach NSW on the addition and multiplication of LinearMap from the switch if the values are monotonic.

alive proofs: https://alive2.llvm.org/ce/z/F8mZUf

Note: SwitchToLookUpTable is inherently target-independent except for BitMap patterns, which embed the values into one register. It might be better to separate independent tests from target-specific test cases also.

[khei4]

rebase

[khei4]

remove comment

[khei4]

fix style

[khei4]

remove unnecessary initializer

TODO: track compile time effect

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6122

I don't think this is sufficient. Let's say we have multiplier 2 and MaxIndex is -1 (that is, full range). Then 2*-1 does not overflow, but there are still plenty of values in the range that do overflow (e.g. INT_MAX).

A possible way to properly handle this would be using ConstantRange. I think another possibility is to used ssub_ov for the Val - PrevVal calculation above and set NSW based on the overflow flag there.

In any case, we need additional tests to cover this.

[khei4]

TODO: use ConstantRange,
TODO: add test case which have select case value ranges [2^n-1, 2^n].

TODO: add signed overflow results

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6122

Let's say we have multiplier 2 and MaxIndex is -1 (that is, full range). Then 2*-1 does not overflow, but there are still plenty of values in the range that do overflow (e.g. INT_MAX).

Yeah, its case matters. Overflowed MaxIndex should be handled. I'll use ConstantRange!

Also, I'll add a linear map test that uses full range.

[khei4]

I would skip the LinearMap part! Until the overflow benefits are revealed!G

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6122

@nikic It seems like I missed Max/MinIndex is signed sense.
I think the tests we need are unsigned results with signed-overflowed/underflowed values

[khei4]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6122

It seems like I missed Max/MinIndex in a signed sense.

I mean we don't need to worry about signed-over/underflow checks with falsy unsigned sense Max/MinIndex.

And I guess this is index full range tests.
(from above URL)

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_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 [
  i2 0, label %bb6
  i2 1, label %bb3
  i2 -2, label %bb4
  i2 -1, label %bb5
  ]

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 [ 4444, %bb5 ], [ 3333, %bb4 ], [ 2222, %bb3 ], [ 1111, %start ]
  ret i32 %.sroa.0.0
}

[khei4]

Change the plan to use PrevDist to check monotonicity.
Rebase for pre-commit test https://reviews.llvm.org/D153238.

[nikic]

LGTM

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6064

I would drop the Inc var and directly write

Wrapped |= Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);

here.

[khei4]

inline Inc.

[uabelho]

Hi @khei4 and @nikic

I think I've found a case where this patch causes a miscompile.
If we do

opt bbi-86196.ll -o - -S -passes="simplifycfg<switch-to-lookup>"

on

define i8 @foo(i3 %0) {                              ; Assume 1
entry:
  switch i3 %0, label %cond.end [
    i3 -1, label %cond.false
    i3 -2, label %cond.false
    i3 1, label %cond.false                          ; -> %cond.false
    i3 0, label %cond.false
  ]

cond.false:                                       ; preds = %entry, %entry, %entry, %entry
  %mul = shl nsw i3 %0, 1                            ; 1 << 1 = 2
  br label %cond.end

cond.end:                                         ; preds = %entry, %cond.false
  %cond = phi i3 [ %mul, %cond.false ], [ 2, %entry ]; 2
  %conv = sext i3 %cond to i8
  ret i8 %conv
}

we get

define i16 @foo(i3 %0) {                            ; Assume 1
entry:
  %switch.tableidx = sub i3 %0, -2                  ; 1 - -2 = 3
  %1 = icmp ult i3 %switch.tableidx, -4             ; 3 (011) u< -4 (100) = 1
  %switch.idx.mult = mul nsw i3 %switch.tableidx, 2 ; 3 * 2 = 6 ovf -2 poison?!
  %switch.offset = add nsw i3 %switch.idx.mult, -4  ; -2 + -4 = -6 ovf 2 poison?!
  %cond = select i1 %1, i3 %switch.offset, i3 2     ; %1 is 1 -> 2 poison?!
  %conv = sext i3 %cond to i16                      ; 2 poison
  ret i16 %conv                                     ; 2 poison
}

which I think is wrong.
I've added some comments in the input and output showing what I think happens for the input value 1.
There are "nsw" on the mul and add instructions, but both caluculations overflow the i3 so I think we get poison there after the transformation?

bbi-86196.ll1 KBDownload

[khei4]

@uabelho Thank you for catching this! As you see, attaching NSW to this example you gave seems unreasonable!

I thought switch results values' monotonicity was enough to say wrapping must not happen, but it wasn't correct.
(-2, -1, 1, 2) |→ (-4, -2, 0, 2) is monotonic, but actual calculation contains ov (-2, -1, 1, 2) |→ (0, 1, 2, 3) (*2) |→ (0, 2, 4(ov), 6(ov)) |→ (-4, -2, 0, 2).

I think we can handle this case if we calculate the maximum value, with monotonicity checking.

Candidate patch: https://github.com/llvm/llvm-project/pull/65882

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, APInt &MaxIndex);
 
   /// 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
for (uint64_t I = 0; I < TableSize; ++I) {
       if (I != 0) {
         APInt Dist = Val - PrevVal;
         if (I == 1) {
           DistToPrev = Dist;
         } else if (Dist != DistToPrev) {
           LinearMappingPossible = false;
           break;
         }
       }
       PrevVal = Val;
     }
     if (LinearMappingPossible) {
       LinearOffset = cast<ConstantInt>(TableContents[0]);
       LinearMultiplier = ConstantInt::get(M.getContext(), DistToPrev);
       Kind = LinearMapKind;
       ++NumLinearMaps;
       return;
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, APInt &MaxIndex) {
   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");
+    // Check whether signed wrap needed.
+    bool MulSW = true;
+    bool AddSW = true;
+    // TODO: cast BitWidth by cleaner way
+    MaxIndex = APInt(LinearMultiplier->getBitWidth(), MaxIndex.getZExtValue());
+    LinearOffset->getValue().sadd_ov(
+        LinearMultiplier->getValue().smul_ov(MaxIndex, MulSW), AddSW);
     if (!LinearMultiplier->isOne())
-      Result = Builder.CreateMul(Result, LinearMultiplier, "switch.idx.mult");
+      Result = Builder.CreateMul(Result, LinearMultiplier, "switch.idx.mult",
+                                 /*HasNUW = */ false, /*HasNSW = */ !MulSW);
     if (!LinearOffset->isZero())
-      Result = Builder.CreateAdd(Result, LinearOffset, "switch.offset");
+      Result = Builder.CreateAdd(Result, LinearOffset, "switch.offset",
+                                 /*HasNUW = */ false,
+                                 /*HasNSW = */ !AddSW);
     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
static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   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;
+  APInt MaxIndex(MaxCaseVal->getBitWidth(), 0);
   if (UseSwitchConditionAsTableIndex) {
     TableIndexOffset = ConstantInt::get(MaxCaseVal->getType(), 0);
     TableIndex = SI->getCondition();
+    MaxIndex = MaxCaseVal->getValue();
   } else {
     TableIndexOffset = MinCaseVal;
     // If the default is unreachable, all case values are s>= MinCaseVal. Then
     // we can try to attach nsw.
     bool CanBeWrapped = true;
     if (!DefaultIsReachable)
+      MaxIndex =
           MaxCaseVal->getValue().ssub_ov(MinCaseVal->getValue(), CanBeWrapped);
-
     TableIndex = Builder.CreateSub(SI->getCondition(), TableIndexOffset,
                                    "switch.tableidx", /*HasNUW =*/false,
-                                   /*HasNSW =*/!CanBeWrapped);
+                                   /*HasNSW =*/!(CanBeWrapped || DefaultIsReachable));
   }
 
   BranchInst *RangeCheckBranch = nullptr;
 
   if (!DefaultIsReachable || GeneratingCoveredLookupTable) {
     Builder.CreateBr(LookupBB);
     if (DTU)
       Updates.push_back({DominatorTree::Insert, BB, LookupBB});
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, MaxIndex);
 
     // 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_to_lookup_table.ll

 
 ; We can use linear mapping.
 define i8 @linearmap1(i32 %c) {
 ; CHECK-LABEL: @linearmap1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i32 [[C:%.*]], 10
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[SWITCH_TABLEIDX]], 4
 ; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[SWITCH_TABLEIDX]] to i8
-; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i8 [[SWITCH_IDX_CAST]], -5
-; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i8 [[SWITCH_IDX_MULT]], 18
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i8 [[SWITCH_IDX_CAST]], -5
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i8 [[SWITCH_IDX_MULT]], 18
 ; CHECK-NEXT:    [[X:%.*]] = select i1 [[TMP0]], i8 [[SWITCH_OFFSET]], i8 3
 ; CHECK-NEXT:    ret i8 [[X]]
 ;
 entry:
   switch i32 %c, label %sw.default [
   i32 10, label %return
   i32 11, label %sw.bb1
   i32 12, label %sw.bb2

 
 ; Linear mapping in a different configuration.
 define i32 @linearmap2(i8 %c) {
 ; CHECK-LABEL: @linearmap2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i8 [[C:%.*]], -13
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i8 [[SWITCH_TABLEIDX]], 4
 ; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i8 [[SWITCH_TABLEIDX]] to i32
-; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[SWITCH_IDX_CAST]], 18
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_CAST]], 18
 ; CHECK-NEXT:    [[X:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 3
 ; CHECK-NEXT:    ret i32 [[X]]
 ;
 entry:
   switch i8 %c, label %sw.default [
   i8 -10, label %return
   i8 -11, label %sw.bb1
   i8 -12, label %sw.bb2

 
 ; Linear mapping with overflows.
 define i8 @linearmap3(i32 %c) {
 ; CHECK-LABEL: @linearmap3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i32 [[C:%.*]], 10
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[SWITCH_TABLEIDX]], 4
 ; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[SWITCH_TABLEIDX]] to i8
-; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i8 [[SWITCH_IDX_CAST]], 100
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i8 [[SWITCH_IDX_CAST]], 100
 ; CHECK-NEXT:    [[X:%.*]] = select i1 [[TMP0]], i8 [[SWITCH_IDX_MULT]], i8 3
 ; CHECK-NEXT:    ret i8 [[X]]
 ;
 entry:
   switch i32 %c, label %sw.default [
   i32 10, label %return
   i32 11, label %sw.bb1
   i32 12, label %sw.bb2

 }
 
 ; Reuse the inverted table range compare.
 define i32 @reuse_cmp1(i32 %x) {
 ; CHECK-LABEL: @reuse_cmp1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X:%.*]], 4
 ; CHECK-NEXT:    [[INVERTED_CMP:%.*]] = xor i1 [[TMP0]], true
-; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[X]], 10
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[X]], 10
 ; CHECK-NEXT:    [[R_0:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 0
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[R_0]], 0
 ; CHECK-NEXT:    [[RETVAL_0:%.*]] = select i1 [[INVERTED_CMP]], i32 100, i32 [[R_0]]
 ; CHECK-NEXT:    ret i32 [[RETVAL_0]]
 ;
 entry:
   switch i32 %x, label %sw.default [
   i32 0, label %sw.bb

 }
 
 ; Cannot reuse the table range compare, because the default value is the same
 ; as one of the case values.
 define i32 @no_reuse_cmp(i32 %x) {
 ; CHECK-LABEL: @no_reuse_cmp(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X:%.*]], 4
-; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[X]], 10
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[X]], 10
 ; CHECK-NEXT:    [[R_0:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 12
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[R_0]], 0
 ; CHECK-NEXT:    [[RETVAL_0:%.*]] = select i1 [[CMP]], i32 [[R_0]], i32 100
 ; CHECK-NEXT:    ret i32 [[RETVAL_0]]
 ;
 entry:
   switch i32 %x, label %sw.default [
   i32 0, label %sw.bb

 
 ; Cannot reuse the table range compare, because the phi at the switch merge
 ; point is not dominated by the switch.
 define i32 @no_reuse_cmp2(i32 %x, i32 %y) {
 ; CHECK-LABEL: @no_reuse_cmp2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[EC:%.*]] = icmp ne i32 [[Y:%.*]], 0
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X:%.*]], 4
-; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i32 [[X]], 10
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[X]], 10
 ; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 0
 ; CHECK-NEXT:    [[R_0:%.*]] = select i1 [[EC]], i32 [[SPEC_SELECT]], i32 100
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[R_0]], 0
 ; CHECK-NEXT:    [[DOTR_0:%.*]] = select i1 [[CMP]], i32 100, i32 [[R_0]]
 ; CHECK-NEXT:    ret i32 [[DOTR_0]]
 ;
 entry:
   %ec = icmp ne i32 %y, 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_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