[llvm] [SimplifyCFG] Find an arrayless index for the covered lookup table (PR #73269)

Quentin Dian via llvm-commits llvm-commits at lists.llvm.org
Mon Feb 5 06:28:59 PST 2024


https://github.com/DianQK updated https://github.com/llvm/llvm-project/pull/73269

>From 6d1d384df0db1f7e4dc11d9f59c9f85beffd1010 Mon Sep 17 00:00:00 2001
From: DianQK <dianqk at dianqk.net>
Date: Thu, 23 Nov 2023 20:54:53 +0800
Subject: [PATCH] [SimplifyCFG] Find an arrayless index for the covered lookup
 table

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 256 ++++++++++++------
 .../SimplifyCFG/X86/CoveredLookupTable.ll     |   9 +-
 .../SimplifyCFG/X86/switch-table-bug.ll       |   7 +-
 .../SimplifyCFG/X86/switch_to_lookup_table.ll |   9 +-
 4 files changed, 185 insertions(+), 96 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index f5b398cae04ed..384cebb37c829 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6186,6 +6186,21 @@ class SwitchLookupTable {
   static bool WouldFitInRegister(const DataLayout &DL, uint64_t TableSize,
                                  Type *ElementType);
 
+  static SmallVector<Constant *, 64> buildTableContents(
+      uint64_t TableSize, ConstantInt *Offset,
+      const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
+      Constant *DefaultValue);
+  static bool
+  canBeSingleValueKind(const SmallVectorImpl<Constant *> &TableContents);
+  static bool
+  canBeLinearMapKind(const SmallVectorImpl<Constant *> &TableContents,
+                     bool &NonMonotonic, APInt &DistToPrev);
+  static bool canBeBitMapKind(const SmallVectorImpl<Constant *> &TableContents,
+                              const DataLayout &DL);
+  static bool
+  canOnlyFallbackToArrayKind(const SmallVectorImpl<Constant *> &TableContents,
+                             const DataLayout &DL);
+
 private:
   // Depending on the contents of the table, it can be represented in
   // different ways.
@@ -6231,97 +6246,36 @@ SwitchLookupTable::SwitchLookupTable(
     Module &M, uint64_t TableSize, ConstantInt *Offset,
     const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
     Constant *DefaultValue, const DataLayout &DL, const StringRef &FuncName) {
-  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.
-  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;
-  }
+  SmallVector<Constant *, 64> TableContents =
+      buildTableContents(TableSize, Offset, Values, DefaultValue);
 
   // If each element in the table contains the same value, we only need to store
   // that single value.
-  if (SingleValue) {
+  if (canBeSingleValueKind(TableContents)) {
+    SingleValue = TableContents[0];
     Kind = SingleValueKind;
     return;
   }
-
-  // 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;
-    // 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.
-        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;
-        }
-        NonMonotonic |=
-            Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);
-      }
-      PrevVal = Val;
-    }
-    if (LinearMappingPossible) {
-      LinearOffset = cast<ConstantInt>(TableContents[0]);
-      LinearMultiplier = ConstantInt::get(M.getContext(), DistToPrev);
-      bool MayWrap = false;
-      APInt M = LinearMultiplier->getValue();
-      (void)M.smul_ov(APInt(M.getBitWidth(), TableSize - 1), MayWrap);
-      LinearMapValWrapped = NonMonotonic || MayWrap;
-      Kind = LinearMapKind;
-      ++NumLinearMaps;
-      return;
-    }
+  // 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;
+  APInt DistToPrev;
+  if (canBeLinearMapKind(TableContents, NonMonotonic, DistToPrev)) {
+    LinearOffset = cast<ConstantInt>(TableContents[0]);
+    LinearMultiplier = ConstantInt::get(M.getContext(), DistToPrev);
+    bool MayWrap = false;
+    APInt M = LinearMultiplier->getValue();
+    (void)M.smul_ov(APInt(M.getBitWidth(), TableSize - 1), MayWrap);
+    LinearMapValWrapped = NonMonotonic || MayWrap;
+    Kind = LinearMapKind;
+    ++NumLinearMaps;
+    return;
   }
 
+  Type *ValueType = Values.begin()->second->getType();
   // If the type is integer and the table fits in a register, build a bitmap.
-  if (WouldFitInRegister(DL, TableSize, ValueType)) {
+  if (canBeBitMapKind(TableContents, DL)) {
     IntegerType *IT = cast<IntegerType>(ValueType);
     APInt TableInt(TableSize * IT->getBitWidth(), 0);
     for (uint64_t I = TableSize; I > 0; --I) {
@@ -6430,6 +6384,110 @@ bool SwitchLookupTable::WouldFitInRegister(const DataLayout &DL,
   return DL.fitsInLegalInteger(TableSize * IT->getBitWidth());
 }
 
+SmallVector<Constant *, 64> SwitchLookupTable::buildTableContents(
+    uint64_t TableSize, ConstantInt *Offset,
+    const SmallVectorImpl<std::pair<ConstantInt *, Constant *>> &Values,
+    Constant *DefaultValue) {
+  assert(Values.size() && "Can't build lookup table without values!");
+  assert(TableSize >= Values.size() && "Can't fit values in table!");
+
+  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;
+  }
+
+  // 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;
+    }
+  }
+  return TableContents;
+}
+
+bool SwitchLookupTable::canBeSingleValueKind(
+    const SmallVectorImpl<Constant *> &TableContents) {
+  // If all values in the table are equal, this is that value.
+  const Constant *SingleValue = TableContents[0];
+  for (const Constant *Value : TableContents) {
+    if (Value != SingleValue)
+      return false;
+  }
+  return true;
+}
+
+bool SwitchLookupTable::canBeLinearMapKind(
+    const SmallVectorImpl<Constant *> &TableContents, bool &NonMonotonic,
+    APInt &DistToPrev) {
+  Type *ValueType = TableContents[0]->getType();
+  // Check if we can derive the value with a linear transformation from the
+  // table index.
+  if (!isa<IntegerType>(ValueType))
+    return false;
+  bool LinearMappingPossible = true;
+  APInt PrevVal;
+  auto TableSize = TableContents.size();
+
+  // When linear map is monotonic and signed overflow doesn't happen on
+  // maximum index, we can attach nsw on Add and Mul.
+  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;
+      }
+      NonMonotonic |=
+          Dist.isStrictlyPositive() ? Val.sle(PrevVal) : Val.sgt(PrevVal);
+    }
+    PrevVal = Val;
+  }
+  return LinearMappingPossible;
+}
+
+bool SwitchLookupTable::canBeBitMapKind(
+    const SmallVectorImpl<Constant *> &TableContents, const DataLayout &DL) {
+  return WouldFitInRegister(DL, TableContents.size(),
+                            TableContents[0]->getType());
+}
+
+bool SwitchLookupTable::canOnlyFallbackToArrayKind(
+    const SmallVectorImpl<Constant *> &TableContents, const DataLayout &DL) {
+  if (canBeSingleValueKind(TableContents))
+    return false;
+  bool NonMonotonic = false;
+  APInt DistToPrev;
+  if (canBeLinearMapKind(TableContents, NonMonotonic, DistToPrev))
+    return false;
+  if (canBeBitMapKind(TableContents, DL))
+    return false;
+  return true;
+}
+
 static bool isTypeLegalForLookupTable(Type *Ty, const TargetTransformInfo &TTI,
                                       const DataLayout &DL) {
   // Allow any legal type.
@@ -6743,6 +6801,42 @@ 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
+  // doesn't 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++) {
+      auto *TableIndexOffset =
+          ConstantInt::get(MaxCaseVal->getIntegerType(), Offset);
+      bool CanOnlyFallbackToArrayKind = false;
+      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];
+        SmallVector<Constant *, 64> TableContents =
+            SwitchLookupTable::buildTableContents(TableSize, TableIndexOffset,
+                                                  ResultList, DV);
+        if (SwitchLookupTable::canOnlyFallbackToArrayKind(TableContents, DL)) {
+          CanOnlyFallbackToArrayKind = true;
+          break;
+        }
+      }
+      if (!CanOnlyFallbackToArrayKind) {
+        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 56fa249d23bba..a333a0dc70752 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 37001f4fba2aa..75cee53d23037 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 3873f0c0ae0bb..1fbabce97e857 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



More information about the llvm-commits mailing list