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

DianQK · 2023-11-24T00:18:08Z

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 · 2023-12-28T15:24:22Z

@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

PR Link: llvm/llvm-project#73269

dtcxzyw · 2023-12-28T15:55:17Z

@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 · 2023-12-28T22:50:57Z

@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 · 2024-01-12T00:55:07Z

Ping.

PR Link: llvm/llvm-project#73269

khei4 · 2024-01-22T08:56:30Z

@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 · 2024-01-22T09:19:52Z

@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 · 2024-02-05T14:29:05Z

Rebase and ping.

DianQK · 2024-04-15T13:50:49Z

Ping

DianQK requested review from nikic, dtcxzyw and goldsteinn November 24, 2023 00:18

DianQK mentioned this pull request Nov 24, 2023

[SimplifyCFG] Find the minimal table considering overflow in switchToLookupTable #67885

Open

khei4 marked this pull request as ready for review December 28, 2023 15:23

llvmbot added the llvm:transforms label Dec 28, 2023

dtcxzyw added a commit to dtcxzyw/llvm-opt-benchmark that referenced this pull request Dec 28, 2023

pre-commit: test PR73269

bad95e1

PR Link: llvm/llvm-project#73269

dtcxzyw mentioned this pull request Dec 28, 2023

pre-commit: test PR73269 dtcxzyw/llvm-opt-benchmark#61

Open

DianQK force-pushed the perf/covered-lookup-table-index branch from 2b8ba66 to 214e491 Compare December 28, 2023 22:48

dtcxzyw marked this pull request as draft December 29, 2023 07:42

DianQK force-pushed the perf/covered-lookup-table-index branch from 214e491 to 1c3c48f Compare December 30, 2023 13:21

DianQK marked this pull request as ready for review December 30, 2023 14:15

DianQK requested a review from khei4 December 30, 2023 14:15

dtcxzyw added a commit to dtcxzyw/llvm-opt-benchmark that referenced this pull request Jan 12, 2024

pre-commit: test PR73269

b9d9939

PR Link: llvm/llvm-project#73269

dtcxzyw added a commit to dtcxzyw/llvm-opt-benchmark that referenced this pull request Jan 12, 2024

pre-commit: test PR73269

9d74288

PR Link: llvm/llvm-project#73269

DianQK requested review from qiongsiwu and vfdff January 22, 2024 09:21

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

6d1d384

DianQK force-pushed the perf/covered-lookup-table-index branch from 1c3c48f to 6d1d384 Compare February 5, 2024 14:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

DianQK commented Nov 24, 2023 •

edited

llvmbot commented Dec 28, 2023

dtcxzyw commented Dec 28, 2023

DianQK commented Dec 28, 2023

DianQK commented Jan 12, 2024

khei4 commented Jan 22, 2024

DianQK commented Jan 22, 2024

DianQK commented Feb 5, 2024

DianQK commented Apr 15, 2024

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

Are you sure you want to change the base?

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

Conversation

DianQK commented Nov 24, 2023 • edited

llvmbot commented Dec 28, 2023

dtcxzyw commented Dec 28, 2023

DianQK commented Dec 28, 2023

DianQK commented Jan 12, 2024

khei4 commented Jan 22, 2024

DianQK commented Jan 22, 2024

DianQK commented Feb 5, 2024

DianQK commented Apr 15, 2024

DianQK commented Nov 24, 2023 •

edited