[llvm] Switch lookup tables undef linear mapping (PR #95539)

via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 14 06:08:54 PDT 2024


https://github.com/DaMatrix created https://github.com/llvm/llvm-project/pull/95539

None

>From 68b795d837ade923ea9c12265925ae8302b8806a Mon Sep 17 00:00:00 2001
From: DaPorkchop_ <daporkchop at daporkchop.net>
Date: Mon, 10 Jun 2024 16:30:42 +0200
Subject: [PATCH 1/4] [SimplifyCFG] Lookup table test cases

---
 .../SimplifyCFG/X86/switch_to_lookup_table.ll | 134 ++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index 845c5008e3837..e4a73b0c0eebc 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -39,6 +39,9 @@ target triple = "x86_64-unknown-linux-gnu"
 ; CHECK: @switch.table.covered_switch_with_bit_tests = private unnamed_addr constant [8 x i32] [i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 1, i32 1], align 4
 ; CHECK: @switch.table.signed_overflow1 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 1111, i32 2222], align 4
 ; CHECK: @switch.table.signed_overflow2 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 2222, i32 2222], align 4
+; CHECK: @switch.table.constant_hole_unreachable_default_firstpoison = private unnamed_addr constant [5 x i32] [i32 poison, i32 poison, i32 1, i32 1, i32 1], align 4
+; CHECK: @switch.table.constant_hole_unreachable_default_lastpoison = private unnamed_addr constant [5 x i32] [i32 1, i32 1, i32 1, i32 1, i32 poison], align 4
+; CHECK: @switch.table.linearmap_hole_unreachable_default = private unnamed_addr constant [5 x i32] [i32 1, i32 1, i32 5, i32 7, i32 9], align 4
 ;.
 define i32 @f(i32 %c) {
 ; CHECK-LABEL: @f(
@@ -2151,3 +2154,134 @@ return:                                           ; preds = %sw.default, %entry,
   %retval.0 = phi { i8, i8 } [ undef, %entry ], [ undef, %entry ], [ undef, %entry ], [ %1, %sw.default ]
   ret { i8, i8 } %retval.0
 }
+
+; The switch has a hole which falls through to an unreachable default case, but it can still be optimized into a constant load because
+; the poison value used for the hole is ignored.
+define i32 @constant_hole_unreachable_default(i32 %x) {
+; CHECK-LABEL: @constant_hole_unreachable_default(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb0
+  i32 3, label %bb0
+  i32 4, label %bb0
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+
+return:
+  %res = phi i32 [ 1, %bb0 ]
+  ret i32 %res
+}
+
+; The switch has a hole which falls through to an unreachable default case and the first case explicitly returns poison, but it can still
+; be optimized into a constant load because the poison values are ignored.
+define i32 @constant_hole_unreachable_default_firstpoison(i32 %x) {
+; CHECK-LABEL: @constant_hole_unreachable_default_firstpoison(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.constant_hole_unreachable_default_firstpoison, i32 0, i32 [[X:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb.poison
+  i32 2, label %bb0
+  i32 3, label %bb0
+  i32 4, label %bb0
+  ]
+
+sw.default: unreachable
+bb.poison: br label %return
+bb0: br label %return
+
+return:
+  %res = phi i32 [ poison, %bb.poison ], [ 1, %bb0 ]
+  ret i32 %res
+}
+
+; The switch has a hole which falls through to an unreachable default case and the first case explicitly returns poison, but it can still
+; be optimized into a constant load because the poison values are ignored.
+define i32 @constant_hole_unreachable_default_lastpoison(i32 %x) {
+; CHECK-LABEL: @constant_hole_unreachable_default_lastpoison(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.constant_hole_unreachable_default_lastpoison, i32 0, i32 [[X:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb0
+  i32 3, label %bb0
+  i32 4, label %bb.poison
+  ]
+
+sw.default: unreachable
+bb.poison: br label %return
+bb0: br label %return
+
+return:
+  %res = phi i32 [ poison, %bb.poison ], [ 1, %bb0 ]
+  ret i32 %res
+}
+
+; The switch has a hole which falls through to an unreachable default case, which prevents it from being optimized into a linear mapping 2*x+1.
+; TODO: We should add support for this, at least in certain cases.
+define i32 @linearmap_hole_unreachable_default(i32 %x) {
+; CHECK-LABEL: @linearmap_hole_unreachable_default(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.linearmap_hole_unreachable_default, i32 0, i32 [[X:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 3, label %bb3
+  i32 4, label %bb4
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb3: br label %return
+bb4: br label %return
+
+return:
+  %res = phi i32 [ 1, %bb0 ], [ 5, %bb2 ], [ 7, %bb3 ], [ 9, %bb4 ]
+  ret i32 %res
+}
+
+; The switch has a hole which falls through to an unreachable default case, but it can still be optimized into a bitmask extraction because
+; the poison value used for the hole is simply replaced with zero.
+define i1 @bitset_hole_unreachable_default(i32 %x) {
+; CHECK-LABEL: @bitset_hole_unreachable_default(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i32 [[X:%.*]] to i5
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i5 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i5 8, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i5 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb0
+  i32 3, label %bb1
+  i32 4, label %bb0
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb1: br label %return
+
+return:
+  %res = phi i1 [ 0, %bb0 ], [ 1, %bb1 ]
+  ret i1 %res
+}

>From d72eeb223aad029e4d5442a4a1553ae0845a617a Mon Sep 17 00:00:00 2001
From: DaPorkchop_ <daporkchop at daporkchop.net>
Date: Mon, 10 Jun 2024 16:31:04 +0200
Subject: [PATCH 2/4] [SimplifyCFG] Replace unreachable switch lookup table
 holes with poison

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 32 ++++++++++++++-----
 .../SimplifyCFG/X86/switch_to_lookup_table.ll | 18 ++++-------
 .../X86/switch_to_lookup_table_big.ll         |  6 ++--
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 292739b6c5fda..50ef68ec02cc6 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6266,7 +6266,8 @@ SwitchLookupTable::SwitchLookupTable(
   assert(Values.size() && "Can't build lookup table without values!");
   assert(TableSize >= Values.size() && "Can't fit values in table!");
 
-  // If all values in the table are equal, this is that value.
+  // If all values in the table are equal, ignoring any values that are poison,
+  // this is that value.
   SingleValue = Values.begin()->second;
 
   Type *ValueType = Values.begin()->second->getType();
@@ -6281,8 +6282,15 @@ SwitchLookupTable::SwitchLookupTable(
     uint64_t Idx = (CaseVal->getValue() - Offset->getValue()).getLimitedValue();
     TableContents[Idx] = CaseRes;
 
-    if (CaseRes != SingleValue)
-      SingleValue = nullptr;
+    if (CaseRes != SingleValue) {
+      if (SingleValue && isa<PoisonValue>(SingleValue)) {
+        // All of the switch cases until now have returned poison; ignore them
+        // and use this value as the single constant value.
+        SingleValue = CaseRes;
+      } else if (!isa<PoisonValue>(CaseRes)) {
+        SingleValue = nullptr;
+      }
+    }
   }
 
   // Fill in any holes in the table with the default result.
@@ -6295,7 +6303,10 @@ SwitchLookupTable::SwitchLookupTable(
         TableContents[I] = DefaultValue;
     }
 
-    if (DefaultValue != SingleValue)
+    // If the default value is poison, all the holes are poison.
+    bool DefaultValueIsPoison = isa<PoisonValue>(DefaultValue);
+
+    if (DefaultValue != SingleValue && !DefaultValueIsPoison)
       SingleValue = nullptr;
   }
 
@@ -6322,6 +6333,9 @@ SwitchLookupTable::SwitchLookupTable(
       if (!ConstVal) {
         // This is an undef. We could deal with it, but undefs in lookup tables
         // are very seldom. It's probably not worth the additional complexity.
+        // TODO: In switches with holes and an unreachable default branch, this
+        //       will actually occur every time, as the holes will be filled
+        //       with poison.
         LinearMappingPossible = false;
         break;
       }
@@ -6752,8 +6766,8 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
 
   // If the table has holes but the default destination doesn't produce any
   // constant results, the lookup table entries corresponding to the holes will
-  // contain undefined values.
-  bool AllHolesAreUndefined = TableHasHoles && !HasDefaultResults;
+  // contain poison.
+  bool AllHolesArePoison = TableHasHoles && !HasDefaultResults;
 
   // If the default destination doesn't produce a constant result but is still
   // reachable, and the lookup table has holes, we need to use a mask to
@@ -6761,7 +6775,7 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   // to the default case.
   // The mask is unnecessary if the table has holes but the default destination
   // is unreachable, as in that case the holes must also be unreachable.
-  bool NeedMask = AllHolesAreUndefined && DefaultIsReachable;
+  bool NeedMask = AllHolesArePoison && DefaultIsReachable;
   if (NeedMask) {
     // As an extra penalty for the validity test we require more cases.
     if (SI->getNumCases() < 4) // FIXME: Find best threshold value (benchmark).
@@ -6906,9 +6920,11 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   for (PHINode *PHI : PHIs) {
     const ResultListTy &ResultList = ResultLists[PHI];
 
+    Type *ResultType = ResultList.begin()->second->getType();
+
     // Use any value to fill the lookup table holes.
     Constant *DV =
-        AllHolesAreUndefined ? ResultLists[PHI][0].second : DefaultResults[PHI];
+        AllHolesArePoison ? PoisonValue::get(ResultType) : DefaultResults[PHI];
     StringRef FuncName = Fn->getName();
     SwitchLookupTable Table(Mod, TableSize, TableIndexOffset, ResultList, DV,
                             DL, FuncName);
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index e4a73b0c0eebc..345476481f570 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -34,14 +34,12 @@ target triple = "x86_64-unknown-linux-gnu"
 ; CHECK: @switch.table.unreachable_case = private unnamed_addr constant [9 x i32] [i32 0, i32 0, i32 0, i32 2, i32 -1, i32 1, i32 1, i32 1, i32 1], align 4
 ; CHECK: @switch.table.unreachable_default = private unnamed_addr constant [4 x i32] [i32 42, i32 52, i32 1, i32 2], align 4
 ; CHECK: @switch.table.nodefaultnoholes = private unnamed_addr constant [4 x i32] [i32 55, i32 123, i32 0, i32 -1], align 4
-; CHECK: @switch.table.nodefaultwithholes = private unnamed_addr constant [6 x i32] [i32 55, i32 123, i32 0, i32 -1, i32 55, i32 -1], align 4
+; CHECK: @switch.table.nodefaultwithholes = private unnamed_addr constant [6 x i32] [i32 55, i32 123, i32 0, i32 -1, i32 poison, i32 -1], align 4
 ; CHECK: @switch.table.threecases = private unnamed_addr constant [3 x i32] [i32 10, i32 7, i32 5], align 4
-; CHECK: @switch.table.covered_switch_with_bit_tests = private unnamed_addr constant [8 x i32] [i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 1, i32 1], align 4
+; CHECK: @switch.table.covered_switch_with_bit_tests = private unnamed_addr constant [8 x i32] [i32 2, i32 2, i32 poison, i32 poison, i32 poison, i32 poison, i32 1, i32 1], align 4
 ; CHECK: @switch.table.signed_overflow1 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 1111, i32 2222], align 4
-; CHECK: @switch.table.signed_overflow2 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 2222, i32 2222], align 4
-; CHECK: @switch.table.constant_hole_unreachable_default_firstpoison = private unnamed_addr constant [5 x i32] [i32 poison, i32 poison, i32 1, i32 1, i32 1], align 4
-; CHECK: @switch.table.constant_hole_unreachable_default_lastpoison = private unnamed_addr constant [5 x i32] [i32 1, i32 1, i32 1, i32 1, i32 poison], align 4
-; CHECK: @switch.table.linearmap_hole_unreachable_default = private unnamed_addr constant [5 x i32] [i32 1, i32 1, i32 5, i32 7, i32 9], align 4
+; CHECK: @switch.table.signed_overflow2 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 poison, i32 2222], align 4
+; CHECK: @switch.table.linearmap_hole_unreachable_default = private unnamed_addr constant [5 x i32] [i32 1, i32 poison, i32 5, i32 7, i32 9], align 4
 ;.
 define i32 @f(i32 %c) {
 ; CHECK-LABEL: @f(
@@ -2183,9 +2181,7 @@ return:
 define i32 @constant_hole_unreachable_default_firstpoison(i32 %x) {
 ; CHECK-LABEL: @constant_hole_unreachable_default_firstpoison(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.constant_hole_unreachable_default_firstpoison, i32 0, i32 [[X:%.*]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    ret i32 1
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2209,9 +2205,7 @@ return:
 define i32 @constant_hole_unreachable_default_lastpoison(i32 %x) {
 ; CHECK-LABEL: @constant_hole_unreachable_default_lastpoison(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.constant_hole_unreachable_default_lastpoison, i32 0, i32 [[X:%.*]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    ret i32 1
 ;
 entry:
   switch i32 %x, label %sw.default [
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table_big.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table_big.ll
index 7988e3057a2c2..4ebf09ae3b127 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table_big.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table_big.ll
@@ -7,11 +7,11 @@ target triple = "i386-pc-linux-gnu"
 ;.
 ; CHECK: @switch.table.reachable_default_dense_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1], align 4
 ; CHECK: @switch.table.unreachable_default_dense_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1], align 4
-; CHECK: @switch.table.reachable_default_holes_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 0, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 0, i32 2, i32 1, i32 0, i32 7, i32 0, i32 5, i32 4, i32 3, i32 2, i32 0, i32 0, i32 7, i32 6, i32 5, i32 0, i32 3, i32 2, i32 1], align 4
-; CHECK: @switch.table.unreachable_default_holes_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 0, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 0, i32 2, i32 1, i32 0, i32 7, i32 0, i32 5, i32 4, i32 3, i32 2, i32 0, i32 0, i32 7, i32 6, i32 5, i32 0, i32 3, i32 2, i32 1], align 4
+; CHECK: @switch.table.reachable_default_holes_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 poison, i32 4, i32 3, i32 2, i32 1, i32 poison, i32 7, i32 6, i32 5, i32 4, i32 poison, i32 2, i32 1, i32 0, i32 7, i32 poison, i32 5, i32 4, i32 3, i32 2, i32 poison, i32 0, i32 7, i32 6, i32 5, i32 poison, i32 3, i32 2, i32 1], align 4
+; CHECK: @switch.table.unreachable_default_holes_0to31 = private unnamed_addr constant [32 x i32] [i32 0, i32 7, i32 6, i32 poison, i32 4, i32 3, i32 2, i32 1, i32 poison, i32 7, i32 6, i32 5, i32 4, i32 poison, i32 2, i32 1, i32 0, i32 7, i32 poison, i32 5, i32 4, i32 3, i32 2, i32 poison, i32 0, i32 7, i32 6, i32 5, i32 poison, i32 3, i32 2, i32 1], align 4
 ; CHECK: @switch.table.reachable_default_dense_0to32 = private unnamed_addr constant [33 x i32] [i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0], align 4
 ; CHECK: @switch.table.unreachable_default_dense_0to32 = private unnamed_addr constant [33 x i32] [i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0], align 4
-; CHECK: @switch.table.unreachable_default_holes_0to32 = private unnamed_addr constant [33 x i32] [i32 0, i32 7, i32 6, i32 0, i32 4, i32 3, i32 2, i32 1, i32 0, i32 7, i32 6, i32 5, i32 4, i32 0, i32 2, i32 1, i32 0, i32 7, i32 0, i32 5, i32 4, i32 3, i32 2, i32 0, i32 0, i32 7, i32 6, i32 5, i32 0, i32 3, i32 2, i32 1, i32 0], align 4
+; CHECK: @switch.table.unreachable_default_holes_0to32 = private unnamed_addr constant [33 x i32] [i32 0, i32 7, i32 6, i32 poison, i32 4, i32 3, i32 2, i32 1, i32 poison, i32 7, i32 6, i32 5, i32 4, i32 poison, i32 2, i32 1, i32 0, i32 7, i32 poison, i32 5, i32 4, i32 3, i32 2, i32 poison, i32 0, i32 7, i32 6, i32 5, i32 poison, i32 3, i32 2, i32 1, i32 0], align 4
 ;.
 define i32 @reachable_default_dense_0to31(i32 %x, i32 %y) {
 ; CHECK-LABEL: @reachable_default_dense_0to31(

>From c26e21f26854d99903dc8f1352da6e9a7925394b Mon Sep 17 00:00:00 2001
From: DaPorkchop_ <daporkchop at daporkchop.net>
Date: Tue, 11 Jun 2024 18:51:24 +0200
Subject: [PATCH 3/4] [SimplifyCFG] Lookup table test cases

---
 .../SimplifyCFG/X86/switch_to_lookup_table.ll | 279 +++++++++++++++++-
 1 file changed, 273 insertions(+), 6 deletions(-)

diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index 345476481f570..f2a50cc66d06f 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -40,6 +40,10 @@ target triple = "x86_64-unknown-linux-gnu"
 ; CHECK: @switch.table.signed_overflow1 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 1111, i32 2222], align 4
 ; CHECK: @switch.table.signed_overflow2 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 poison, i32 2222], align 4
 ; CHECK: @switch.table.linearmap_hole_unreachable_default = private unnamed_addr constant [5 x i32] [i32 1, i32 poison, i32 5, i32 7, i32 9], align 4
+; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive = private unnamed_addr constant [7 x i32] [i32 1, i32 poison, i32 5, i32 poison, i32 9, i32 poison, i32 13], align 4
+; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_neg = private unnamed_addr constant [8 x i32] [i32 -1, i32 poison, i32 poison, i32 5, i32 poison, i32 9, i32 poison, i32 13], align 4
+; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative1 = private unnamed_addr constant [7 x i32] [i32 1111, i32 poison, i32 2222, i32 poison, i32 3333, i32 poison, i32 4444], align 4
+; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2 = private unnamed_addr constant [7 x i32] [i32 1111, i32 poison, i32 3333, i32 poison, i32 5555, i32 poison, i32 7777], align 4
 ;.
 define i32 @f(i32 %c) {
 ; CHECK-LABEL: @f(
@@ -2224,8 +2228,7 @@ return:
   ret i32 %res
 }
 
-; The switch has a hole which falls through to an unreachable default case, which prevents it from being optimized into a linear mapping 2*x+1.
-; TODO: We should add support for this, at least in certain cases.
+; The switch has a hole which falls through to an unreachable default case, but can still be optimized into a linear mapping 2*x+1.
 define i32 @linearmap_hole_unreachable_default(i32 %x) {
 ; CHECK-LABEL: @linearmap_hole_unreachable_default(
 ; CHECK-NEXT:  entry:
@@ -2252,10 +2255,246 @@ return:
   ret i32 %res
 }
 
-; The switch has a hole which falls through to an unreachable default case, but it can still be optimized into a bitmask extraction because
-; the poison value used for the hole is simply replaced with zero.
-define i1 @bitset_hole_unreachable_default(i32 %x) {
-; CHECK-LABEL: @bitset_hole_unreachable_default(
+; We can optimize this into a linear mapping even though every case is separated from the others by an unreachable hole.
+define i32 @linearmap_unreachable_holes_noconsecutive(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive, i32 0, i32 [[X:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i32 [ 1, %bb0 ], [ 5, %bb2 ], [ 9, %bb4 ], [ 13, %bb6 ]
+  ret i32 %res
+}
+
+; The switch is very sparse; every case is separated from the othery by a hole which falls through to an unreachable default case.
+; Despite that, it can still be optimized into a linear mapping 2*x+1.
+define i32 @linearmap_unreachable_holes_noconsecutive_neg(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_neg(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i32 [[X:%.*]], -1
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [8 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive_neg, i32 0, i32 [[SWITCH_TABLEIDX]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 -1, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i32 [ -1, %bb0 ], [ 5, %bb2 ], [ 9, %bb4 ], [ 13, %bb6 ]
+  ret i32 %res
+}
+
+; This corresponds to @signed_overflow_negative above.
+; We can't optimize this into a linear mapping because it would need a multiplier of 1111/2 (which isn't an integer)
+define i32 @linearmap_unreachable_holes_noconsecutive_signed_overflow_negative1(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_overflow_negative1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[X:%.*]] to i3
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TRUNC]], -4
+; CHECK-NEXT:    [[SWITCH_TABLEIDX_ZEXT:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i4
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative1, i32 0, i4 [[SWITCH_TABLEIDX_ZEXT]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  %trunc = trunc i32 %x to i3
+  switch i3 %trunc, label %sw.default [
+  i3 0, label %bb0
+  i3 2, label %bb2
+  i3 -4, label %bbm4
+  i3 -2, label %bbm2
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bbm4: br label %return
+bbm2: br label %return
+
+return:
+  %res = phi i32 [ 3333, %bb0 ], [ 4444, %bb2 ], [ 1111, %bbm4 ], [ 2222, %bbm2 ]
+  ret i32 %res
+}
+
+; This corresponds to @signed_overflow_negative above.
+; We can optimize this into a linear mapping.
+define i32 @linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[X:%.*]] to i3
+; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TRUNC]], -4
+; CHECK-NEXT:    [[SWITCH_TABLEIDX_ZEXT:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i4
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2, i32 0, i4 [[SWITCH_TABLEIDX_ZEXT]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+entry:
+  %trunc = trunc i32 %x to i3
+  switch i3 %trunc, label %sw.default [
+  i3 0, label %bb0
+  i3 2, label %bb2
+  i3 -4, label %bbm4
+  i3 -2, label %bbm2
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bbm4: br label %return
+bbm2: br label %return
+
+return:
+  %res = phi i32 [ 5555, %bb0 ], [ 7777, %bb2 ], [ 1111, %bbm4 ], [ 3333, %bbm2 ]
+  ret i32 %res
+}
+
+; This corresponds to @linearmap_inc_nsw above.
+define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_nsw(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_nsw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 34058863028011039, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i8 [ 31, %bb0 ], [ 61, %bb2 ], [ 91, %bb4 ], [ 121, %bb6 ]
+  ret i8 %res
+}
+
+; This corresponds to @linearmap_inc_wrapped above.
+define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_wrapped(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_wrapped(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 -36028384697909216, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i8 [ 32, %bb0 ], [ 64, %bb2 ], [ 96, %bb4 ], [ -128, %bb6 ]
+  ret i8 %res
+}
+
+; This corresponds to @linearmap_dec_nsw above.
+define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_nsw(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_nsw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 -36028109811613472, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i8 [ -32, %bb0 ], [ -64, %bb2 ], [ -96, %bb4 ], [ -128, %bb6 ]
+  ret i8 %res
+}
+
+; This corresponds to @linearmap_dec_wrapped above.
+define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_wrapped(i32 %x) {
+; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_wrapped(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 34059137914306783, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
+; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb2
+  i32 4, label %bb4
+  i32 6, label %bb6
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb2: br label %return
+bb4: br label %return
+bb6: br label %return
+
+return:
+  %res = phi i8 [ -33, %bb0 ], [ -67, %bb2 ], [ -101, %bb4 ], [ 121, %bb6 ]
+  ret i8 %res
+}
+
+; The switch has a hole which falls through to an unreachable default case, but it can still be optimized into a
+; linear mapping (with an offset of 0 and a multiplier of 1)
+define i1 @bitset_hole_unreachable_default1(i32 %x) {
+; CHECK-LABEL: @bitset_hole_unreachable_default1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i32 [[X:%.*]] to i5
 ; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i5 [[SWITCH_CAST]], 1
@@ -2279,3 +2518,31 @@ return:
   %res = phi i1 [ 0, %bb0 ], [ 1, %bb1 ]
   ret i1 %res
 }
+
+; The switch has a hole which falls through to an unreachable default case, but it can still be optimized into a
+; linear mapping (with an offset of 0 and a multiplier of 1)
+define i1 @bitset_hole_unreachable_default2(i32 %x) {
+; CHECK-LABEL: @bitset_hole_unreachable_default2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i32 [[X:%.*]] to i6
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i6 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i6 -32, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i6 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  switch i32 %x, label %sw.default [
+  i32 0, label %bb0
+  i32 2, label %bb0
+  i32 4, label %bb0
+  i32 5, label %bb1
+  ]
+
+sw.default: unreachable
+bb0: br label %return
+bb1: br label %return
+
+return:
+  %res = phi i1 [ 0, %bb0 ], [ 1, %bb1 ]
+  ret i1 %res
+}

>From 61e6e22fb492dead4a8fff2408923e90ead485bc Mon Sep 17 00:00:00 2001
From: DaPorkchop_ <daporkchop at daporkchop.net>
Date: Tue, 11 Jun 2024 18:55:57 +0200
Subject: [PATCH 4/4] [SimplifyCFG] Transform lookup tables with holes into
 linear mappings

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 147 ++++++++++++++----
 .../SimplifyCFG/X86/switch_to_lookup_table.ll |  80 ++++------
 2 files changed, 151 insertions(+), 76 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 50ef68ec02cc6..26f9be403d2d8 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6317,48 +6317,137 @@ SwitchLookupTable::SwitchLookupTable(
     return;
   }
 
+  assert(
+      2 <= std::count_if(TableContents.begin(), TableContents.end(),
+                         [](const auto *C) { return !isa<PoisonValue>(C); }) &&
+      "Should be a SingleValue table.");
+
   // Check if we can derive the value with a linear transformation from the
   // table index.
   if (isa<IntegerType>(ValueType)) {
-    bool LinearMappingPossible = true;
-    APInt PrevVal;
-    APInt DistToPrev;
+    // If we can find two consecutive lookup table entries with a defined value,
+    // we can unambiguously determine the multiplier by computing the difference
+    bool FoundMultiplier = false;
+    APInt OffsetVal;
+    APInt MultiplierVal;
+    for (uint64_t I = 1; I < TableSize; ++I) {
+      ConstantInt *ConstValPrev = dyn_cast<ConstantInt>(TableContents[I - 1]);
+      ConstantInt *ConstValNext = dyn_cast<ConstantInt>(TableContents[I]);
+      if (ConstValPrev && ConstValNext) {
+        FoundMultiplier = true;
+
+        const APInt &PrevVal = ConstValPrev->getValue();
+        const APInt &NextVal = ConstValNext->getValue();
+        MultiplierVal = NextVal - PrevVal;
+        OffsetVal = NextVal - MultiplierVal * I;
+        break;
+      }
+    }
+
+    // If there are no two consecutive lookup table entries with a defined
+    // value, find any two defined entries and try to compute a multiplier
+    // based on their values using the multiplicative inverse.
+    if (!FoundMultiplier) {
+      auto NotPoisonPredicate = [](const auto *C) {
+        return !isa<PoisonValue>(C);
+      };
+
+      // Find the first two table entries with a defined value.
+      // This choice is not optimal - we'd get best results if we preferred to
+      // choose a pair of entries with different values, separated by an odd
+      // number of table indices. However, that would add a lot of complexity,
+      // and it would only improve detection of linear mappings with overflows,
+      // which are not as common.
+      auto First = std::find_if(TableContents.begin(), TableContents.end(),
+                                NotPoisonPredicate);
+      auto Second =
+          std::find_if(First + 1, TableContents.end(), NotPoisonPredicate);
+      assert(First != TableContents.end() && Second != TableContents.end() &&
+             "Should have at least two defined values");
+      assert(isa<ConstantInt>(*First) && isa<ConstantInt>(*Second) &&
+             "Should both be ConstantInt");
+
+      uint64_t FirstIndex = First - TableContents.begin();
+      uint64_t SecondIndex = Second - TableContents.begin();
+      const APInt &FirstVal = cast<ConstantInt>(*First)->getValue();
+      const APInt &SecondVal = cast<ConstantInt>(*Second)->getValue();
+
+      // The number of table indices separating the two values, truncated
+      APInt IndexDist(FirstVal.getBitWidth(), SecondIndex - FirstIndex);
+
+      // The difference between the two values
+      APInt ValDist = SecondVal - FirstVal;
+
+      // We need to find an integer M such that M * IndexDist = ValDist.
+      // If the multiplicative inverse of IndexDist exists, we can compute M by:
+      //   M = ValDist * (1 / IndexDist)
+      // However, this only works if IndexDist is odd.
+
+      // If IndexDist is even (specifically, if it's divisible by 2^N), we can
+      // make it odd if ValDist is also divisible by 2^N.
+      //   M * IndexDist         = ValDist
+      //   M * (IndexDist / 2^N) = ValDist / 2^N
+      unsigned N = IndexDist.countTrailingZeros();
+      if (N <= ValDist.countTrailingZeros()) {
+        IndexDist.ashrInPlace(N);
+        ValDist.ashrInPlace(N);
+      }
+
+      // If IndexDist is odd (i.e. bit 0 is set), we can proceed:
+      if (IndexDist[0]) {
+        FoundMultiplier = true;
+        MultiplierVal = IndexDist.multiplicativeInverse() * ValDist;
+        OffsetVal = FirstVal - MultiplierVal * FirstIndex;
+
+        assert(OffsetVal + MultiplierVal * FirstIndex == FirstVal &&
+               OffsetVal + MultiplierVal * SecondIndex == SecondVal);
+      }
+    }
+
+    bool LinearMappingPossible = false;
     // When linear map is monotonic and signed overflow doesn't happen on
     // maximum index, we can attach nsw on Add and Mul.
     bool NonMonotonic = false;
-    assert(TableSize >= 2 && "Should be a SingleValue table.");
-    // Check if there is the same distance between two consecutive values.
-    for (uint64_t I = 0; I < TableSize; ++I) {
-      ConstantInt *ConstVal = dyn_cast<ConstantInt>(TableContents[I]);
-      if (!ConstVal) {
-        // This is an undef. We could deal with it, but undefs in lookup tables
-        // are very seldom. It's probably not worth the additional complexity.
-        // TODO: In switches with holes and an unreachable default branch, this
-        //       will actually occur every time, as the holes will be filled
-        //       with poison.
-        LinearMappingPossible = false;
-        break;
-      }
-      const APInt &Val = ConstVal->getValue();
-      if (I != 0) {
-        APInt Dist = Val - PrevVal;
-        if (I == 1) {
-          DistToPrev = Dist;
-        } else if (Dist != DistToPrev) {
+    // Check if the offset and multiplier actually give the correct result in
+    // all cases.
+    if (FoundMultiplier) {
+      LinearMappingPossible = true;
+      for (uint64_t I = 0; I < TableSize; ++I) {
+        ConstantInt *ConstVal = dyn_cast<ConstantInt>(TableContents[I]);
+        if (!ConstVal) {
+          if (isa<PoisonValue>(TableContents[I])) {
+            // This is poison. We don't care what the linear mapping evaluates
+            // to, so we'll allow it.
+            continue;
+          }
+
+          // TODO: Can we treat undef the same as poison?
+          assert(isa<UndefValue>(TableContents[I]) && "Should be undef");
+          // This is an undef. We could deal with it, but undefs in lookup
+          // tables are very seldom. It's probably not worth the additional
+          // complexity.
           LinearMappingPossible = false;
           break;
         }
-        NonMonotonic |=
-            Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);
+        const APInt &Val = ConstVal->getValue();
+        APInt ComputedVal = OffsetVal + MultiplierVal * I;
+        if (Val != ComputedVal) {
+          LinearMappingPossible = false;
+          break;
+        }
+        if (I != 0) {
+          APInt PrevVal = ComputedVal - MultiplierVal;
+          NonMonotonic |= MultiplierVal.isStrictlyPositive() ? Val.sle(PrevVal)
+                                                             : Val.sgt(PrevVal);
+        }
       }
-      PrevVal = Val;
     }
     if (LinearMappingPossible) {
-      LinearOffset = cast<ConstantInt>(TableContents[0]);
-      LinearMultiplier = ConstantInt::get(M.getContext(), DistToPrev);
+      LinearOffset = ConstantInt::get(M.getContext(), OffsetVal);
+      LinearMultiplier = ConstantInt::get(M.getContext(), MultiplierVal);
       bool MayWrap = false;
-      APInt M = LinearMultiplier->getValue();
-      (void)M.smul_ov(APInt(M.getBitWidth(), TableSize - 1), MayWrap);
+      (void)MultiplierVal.smul_ov(
+          APInt(MultiplierVal.getBitWidth(), TableSize - 1), MayWrap);
       LinearMapValWrapped = NonMonotonic || MayWrap;
       Kind = LinearMapKind;
       ++NumLinearMaps;
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
index f2a50cc66d06f..884046b2efa7a 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -39,11 +39,7 @@ target triple = "x86_64-unknown-linux-gnu"
 ; CHECK: @switch.table.covered_switch_with_bit_tests = private unnamed_addr constant [8 x i32] [i32 2, i32 2, i32 poison, i32 poison, i32 poison, i32 poison, i32 1, i32 1], align 4
 ; CHECK: @switch.table.signed_overflow1 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 1111, i32 2222], align 4
 ; CHECK: @switch.table.signed_overflow2 = private unnamed_addr constant [4 x i32] [i32 3333, i32 4444, i32 poison, i32 2222], align 4
-; CHECK: @switch.table.linearmap_hole_unreachable_default = private unnamed_addr constant [5 x i32] [i32 1, i32 poison, i32 5, i32 7, i32 9], align 4
-; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive = private unnamed_addr constant [7 x i32] [i32 1, i32 poison, i32 5, i32 poison, i32 9, i32 poison, i32 13], align 4
-; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_neg = private unnamed_addr constant [8 x i32] [i32 -1, i32 poison, i32 poison, i32 5, i32 poison, i32 9, i32 poison, i32 13], align 4
 ; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative1 = private unnamed_addr constant [7 x i32] [i32 1111, i32 poison, i32 2222, i32 poison, i32 3333, i32 poison, i32 4444], align 4
-; CHECK: @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2 = private unnamed_addr constant [7 x i32] [i32 1111, i32 poison, i32 3333, i32 poison, i32 5555, i32 poison, i32 7777], align 4
 ;.
 define i32 @f(i32 %c) {
 ; CHECK-LABEL: @f(
@@ -2232,9 +2228,9 @@ return:
 define i32 @linearmap_hole_unreachable_default(i32 %x) {
 ; CHECK-LABEL: @linearmap_hole_unreachable_default(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [5 x i32], ptr @switch.table.linearmap_hole_unreachable_default, i32 0, i32 [[X:%.*]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], 1
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2259,9 +2255,9 @@ return:
 define i32 @linearmap_unreachable_holes_noconsecutive(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive, i32 0, i32 [[X:%.*]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[X:%.*]], 2
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], 1
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2288,9 +2284,9 @@ define i32 @linearmap_unreachable_holes_noconsecutive_neg(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_neg(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub nsw i32 [[X:%.*]], -1
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [8 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive_neg, i32 0, i32 [[SWITCH_TABLEIDX]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[SWITCH_TABLEIDX]], 2
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], -1
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2350,10 +2346,10 @@ define i32 @linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[X:%.*]] to i3
 ; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[TRUNC]], -4
-; CHECK-NEXT:    [[SWITCH_TABLEIDX_ZEXT:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i4
-; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.linearmap_unreachable_holes_noconsecutive_signed_overflow_negative2, i32 0, i4 [[SWITCH_TABLEIDX_ZEXT]]
-; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
-; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i3 [[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]], 1111
+; CHECK-NEXT:    ret i32 [[SWITCH_OFFSET]]
 ;
 entry:
   %trunc = trunc i32 %x to i3
@@ -2379,11 +2375,10 @@ return:
 define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_nsw(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_nsw(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 34058863028011039, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
-; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i8 [[SWITCH_IDX_CAST]], 15
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i8 [[SWITCH_IDX_MULT]], 31
+; CHECK-NEXT:    ret i8 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2408,11 +2403,10 @@ return:
 define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_wrapped(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_inc_wrapped(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 -36028384697909216, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
-; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i8 [[SWITCH_IDX_CAST]], 16
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i8 [[SWITCH_IDX_MULT]], 32
+; CHECK-NEXT:    ret i8 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2437,11 +2431,10 @@ return:
 define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_nsw(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_nsw(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 -36028109811613472, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
-; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i8 [[SWITCH_IDX_CAST]], -16
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i8 [[SWITCH_IDX_MULT]], -32
+; CHECK-NEXT:    ret i8 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2466,11 +2459,10 @@ return:
 define i8 @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_wrapped(i32 %x) {
 ; CHECK-LABEL: @linearmap_unreachable_holes_noconsecutive_signed_linearmap_dec_wrapped(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i32 [[X:%.*]] to i56
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i56 [[SWITCH_CAST]], 8
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i56 34059137914306783, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i56 [[SWITCH_DOWNSHIFT]] to i8
-; CHECK-NEXT:    ret i8 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul i8 [[SWITCH_IDX_CAST]], -17
+; CHECK-NEXT:    [[SWITCH_OFFSET:%.*]] = add i8 [[SWITCH_IDX_MULT]], -33
+; CHECK-NEXT:    ret i8 [[SWITCH_OFFSET]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2496,11 +2488,8 @@ return:
 define i1 @bitset_hole_unreachable_default1(i32 %x) {
 ; CHECK-LABEL: @bitset_hole_unreachable_default1(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i32 [[X:%.*]] to i5
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i5 [[SWITCH_CAST]], 1
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i5 8, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i5 [[SWITCH_DOWNSHIFT]] to i1
-; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_IDX_CAST]]
 ;
 entry:
   switch i32 %x, label %sw.default [
@@ -2524,11 +2513,8 @@ return:
 define i1 @bitset_hole_unreachable_default2(i32 %x) {
 ; CHECK-LABEL: @bitset_hole_unreachable_default2(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i32 [[X:%.*]] to i6
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i6 [[SWITCH_CAST]], 1
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i6 -32, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i6 [[SWITCH_DOWNSHIFT]] to i1
-; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = trunc i32 [[X:%.*]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_IDX_CAST]]
 ;
 entry:
   switch i32 %x, label %sw.default [



More information about the llvm-commits mailing list