[SimplifyCFG] Find an arrayless index for the covered lookup table

[dianqk]

This PR restores the optimization I found in #67885.

If we are generating a covered lookup table, try to find an index that doesn't create an array of values.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Quentin Dian (DianQK)

Changes

This PR restores the optimization I found in #67885.

I'm marking this as a draft PR because I copied too much code.

---

Full diff: https://github.com/llvm/llvm-project/pull/73269.diff

4 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+128)
• (modified) llvm/test/Transforms/SimplifyCFG/X86/CoveredLookupTable.ll (+4-5)
• (modified) llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll (+2-5)
• (modified) llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll (+4-5)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 4b4d0595cbf776..190e6c5b2085bf 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6081,6 +6081,11 @@ class SwitchLookupTable {
   static bool WouldFitInRegister(const DataLayout &DL, uint64_t TableSize,
                                  Type *ElementType);
 
+  static bool isArrayKind(
+      Module &M, uint64_t TableSize, ConstantInt *Offset,
+      const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
+      Constant *DefaultValue, const DataLayout &DL);
+
 private:
   // Depending on the contents of the table, it can be represented in
   // different ways.
@@ -6122,6 +6127,95 @@ class SwitchLookupTable {
 
 } // end anonymous namespace
 
+bool SwitchLookupTable::isArrayKind(
+    Module &M, uint64_t TableSize, ConstantInt *Offset,
+    const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
+    Constant *DefaultValue, const DataLayout &DL) {
+
+  assert(Values.size() && "Can't build lookup table without values!");
+  assert(TableSize >= Values.size() && "Can't fit values in table!");
+  // For SingleValueKind, this is the single value.
+  Constant *SingleValue = nullptr;
+
+  // If all values in the table are equal, this is that value.
+  SingleValue = Values.begin()->second;
+
+  Type *ValueType = Values.begin()->second->getType();
+
+  // Build up the table contents.
+  SmallVector<Constant *, 64> TableContents(TableSize);
+  for (size_t I = 0, E = Values.size(); I != E; ++I) {
+    ConstantInt *CaseVal = Values[I].first;
+    Constant *CaseRes = Values[I].second;
+    assert(CaseRes->getType() == ValueType);
+
+    uint64_t Idx = (CaseVal->getValue() - Offset->getValue()).getLimitedValue();
+    TableContents[Idx] = CaseRes;
+
+    if (CaseRes != SingleValue)
+      SingleValue = nullptr;
+  }
+
+  // Fill in any holes in the table with the default result.
+  if (Values.size() < TableSize) {
+    assert(DefaultValue &&
+           "Need a default value to fill the lookup table holes.");
+    assert(DefaultValue->getType() == ValueType);
+    for (uint64_t I = 0; I < TableSize; ++I) {
+      if (!TableContents[I])
+        TableContents[I] = DefaultValue;
+    }
+
+    if (DefaultValue != SingleValue)
+      SingleValue = nullptr;
+  }
+
+  // If each element in the table contains the same value, we only need to store
+  // that single value.
+  if (SingleValue) {
+    return false;
+  }
+
+  // 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;
+    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.
+        LinearMappingPossible = false;
+        break;
+      }
+      const APInt &Val = ConstVal->getValue();
+      if (I != 0) {
+        APInt Dist = Val - PrevVal;
+        if (I == 1) {
+          DistToPrev = Dist;
+        } else if (Dist != DistToPrev) {
+          LinearMappingPossible = false;
+          break;
+        }
+      }
+      PrevVal = Val;
+    }
+    if (LinearMappingPossible) {
+      return false;
+    }
+  }
+
+  // If the type is integer and the table fits in a register, build a bitmap.
+  if (WouldFitInRegister(DL, TableSize, ValueType)) {
+    return false;
+  }
+  return true;
+}
+
 SwitchLookupTable::SwitchLookupTable(
     Module &M, uint64_t TableSize, ConstantInt *Offset,
     const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
@@ -6638,6 +6732,40 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
   Module &Mod = *CommonDest->getParent()->getParent();
   BasicBlock *LookupBB = BasicBlock::Create(
       Mod.getContext(), "switch.lookup", CommonDest->getParent(), CommonDest);
+  // If we are generating a covered lookup table, try to find an index that does
+  // create an array of values.
+  // TODO: We could more expensive check, and choose the index with the best sum
+  // of all kinds.
+  if (MaxTableSize == TableSize && TableSize * PHIs.size() <= 128) {
+    for (uint64_t Offset = 0; Offset < TableSize; Offset++) {
+      ConstantInt *TableIndexOffset =
+          ConstantInt::get(MaxCaseVal->getType(), Offset);
+      bool NoArrayKind = true;
+      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];
+        if (SwitchLookupTable::isArrayKind(Mod, TableSize, TableIndexOffset,
+                                           ResultList, DV, DL)) {
+          NoArrayKind = false;
+          break;
+        }
+      }
+      if (NoArrayKind) {
+        if (Offset == 0)
+          UseSwitchConditionAsTableIndex = true;
+        MinCaseVal = TableIndexOffset;
+        APInt One(TableIndexOffset->getValue().getBitWidth(), 1);
+        bool Overflow = false;
+        MaxCaseVal = cast<ConstantInt>(ConstantInt::get(
+            MaxCaseVal->getType(),
+            TableIndexOffset->getValue().usub_ov(One, Overflow)));
+        break;
+      }
+    }
+  }
 
   // Compute the table index value.
   Builder.SetInsertPoint(SI);
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/CoveredLookupTable.ll b/llvm/test/Transforms/SimplifyCFG/X86/CoveredLookupTable.ll
index 56fa249d23bba2..a333a0dc70752d 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/CoveredLookupTable.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/CoveredLookupTable.ll
@@ -9,11 +9,10 @@ target triple = "x86_64-apple-darwin12.0.0"
 define i3 @coveredswitch_test(i3 %input) {
 ; CHECK-LABEL: @coveredswitch_test(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[SWITCH_TABLEIDX:%.*]] = sub i3 [[INPUT:%.*]], -4
-; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i3 [[SWITCH_TABLEIDX]] to i24
-; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i24 [[SWITCH_CAST]], 3
-; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i24 7507338, [[SWITCH_SHIFTAMT]]
-; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i24 [[SWITCH_DOWNSHIFT]] to i3
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = zext i3 [[INPUT:%.*]] to i21
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i21 [[SWITCH_CAST]], 3
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i21 -481496, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i21 [[SWITCH_DOWNSHIFT]] to i3
 ; CHECK-NEXT:    ret i3 [[SWITCH_MASKED]]
 ;
 entry:
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll b/llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll
index 37001f4fba2aa8..75cee53d230374 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch-table-bug.ll
@@ -9,11 +9,8 @@ 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_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]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i2 [[TMP0:%.*]] to i64
+; CHECK-NEXT:    ret i64 [[SWITCH_IDX_CAST]]
 ;
 entry:
   switch i2 %0, label %1 [
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 3873f0c0ae0bbd..1fbabce97e8578 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll
@@ -1696,11 +1696,10 @@ 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]]
+; CHECK-NEXT:    [[SWITCH_IDX_CAST:%.*]] = zext i2 [[TRUNC]] 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]]
 ;
 start:
   %trunc = trunc i8 %n to i2

[dtcxzyw]

@dianqk Could you please fix the compilation error first?

FAILED: lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o 
ccache /usr/bin/c++ -DGTEST_HAS_RTTI=0 -D_DEBUG -D_GLIBCXX_ASSERTIONS -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-build/lib/Transforms/Utils -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-build/include -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/include -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -fno-lifetime-dse -w -fdiagnostics-color -ffunction-sections -fdata-sections -O3 -DNDEBUG -std=c++17 -fPIC  -fno-exceptions -funwind-tables -fno-rtti -UNDEBUG -MD -MT lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o -MF lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o.d -o lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o -c /home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp: In function ‘bool SwitchToLookupTable(llvm::SwitchInst*, llvm::IRBuilder<>&, llvm::DomTreeUpdater*, const llvm::DataLayout&, const llvm::TargetTransformInfo&)’:
/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp:6766:27: error: invalid conversion from ‘llvm::Constant*’ to ‘llvm::ConstantInt*’ [-fpermissive]
 6766 |           ConstantInt::get(MaxCaseVal->getType(), Offset);
      |           ~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      |                           |
      |                           llvm::Constant*

[dianqk]

@dianqk Could you please fix the compilation error first?

FAILED: lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o 
ccache /usr/bin/c++ -DGTEST_HAS_RTTI=0 -D_DEBUG -D_GLIBCXX_ASSERTIONS -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-build/lib/Transforms/Utils -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-build/include -I/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/include -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -fno-lifetime-dse -w -fdiagnostics-color -ffunction-sections -fdata-sections -O3 -DNDEBUG -std=c++17 -fPIC  -fno-exceptions -funwind-tables -fno-rtti -UNDEBUG -MD -MT lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o -MF lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o.d -o lib/Transforms/Utils/CMakeFiles/LLVMTransformUtils.dir/SimplifyCFG.cpp.o -c /home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp: In function ‘bool SwitchToLookupTable(llvm::SwitchInst*, llvm::IRBuilder<>&, llvm::DomTreeUpdater*, const llvm::DataLayout&, const llvm::TargetTransformInfo&)’:
/home/dtcxzyw/llvm-opt/actions-runner/_work/llvm-opt-benchmark/llvm-opt-benchmark/llvm/llvm-project/llvm/lib/Transforms/Utils/SimplifyCFG.cpp:6766:27: error: invalid conversion from ‘llvm::Constant*’ to ‘llvm::ConstantInt*’ [-fpermissive]
 6766 |           ConstantInt::get(MaxCaseVal->getType(), Offset);
      |           ~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      |                           |
      |                           llvm::Constant*

No problem. Updated. But please note this PR I am not ready to review.

@khei4 Is there any reason you marked it as ready for review? :)

[dianqk]

Ping.

[khei4]

@khei4 Is there any reason you marked it as ready for review? :)

I'm very sorry about this. That was just a miss-touch...

[dianqk]

@khei4 Is there any reason you marked it as ready for review? :)

I'm very sorry about this. That was just a miss-touch...

Sometimes I worry about accidentally clicking the "Squash and merge".

[dianqk]

Rebase and ping.

[dianqk]

Ping