[SimplifyCFG] Allow some switch optimizations early in the pipeline

While we do not want to form actual lookup tables early, we do
want to perform some optimizations, as they may enable inlining
of the much simpler form.

Builds on llvm#156477, which
originally included this change as well.

Author: nikic

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -6642,6 +6642,9 @@ class SwitchReplacement {
   /// Return true if the replacement is a lookup table.
   bool isLookupTable();
 
+  /// Return true if the replacement is a bit map.
+  bool isBitMap();
+
 private:
   // Depending on the switch, there are different alternatives.
   enum {
@@ -6932,6 +6935,8 @@ Constant *SwitchReplacement::getDefaultValue() { return DefaultValue; }
 
 bool SwitchReplacement::isLookupTable() { return Kind == LookupTableKind; }
 
+bool SwitchReplacement::isBitMap() { return Kind == BitMapKind; }
+
 static bool isSwitchDense(uint64_t NumCases, uint64_t CaseRange) {
   // 40% is the default density for building a jump table in optsize/minsize
   // mode. See also TargetLoweringBase::isSuitableForJumpTable(), which this
@@ -7097,7 +7102,8 @@ static void reuseTableCompare(
 /// lookup tables.
 static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
                                  DomTreeUpdater *DTU, const DataLayout &DL,
-                                 const TargetTransformInfo &TTI) {
+                                 const TargetTransformInfo &TTI,
+                                 bool ConvertSwitchToLookupTable) {
   assert(SI->getNumCases() > 1 && "Degenerate switch?");
 
   BasicBlock *BB = SI->getParent();
@@ -7262,6 +7268,8 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
 
   bool AnyLookupTables = any_of(
       PhiToReplacementMap, [](auto &KV) { return KV.second.isLookupTable(); });
+  bool AnyBitMaps = any_of(PhiToReplacementMap,
+                           [](auto &KV) { return KV.second.isBitMap(); });
 
   // A few conditions prevent the generation of lookup tables:
   //     1. The target does not support lookup tables.
@@ -7274,6 +7282,11 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
        Fn->getFnAttribute("no-jump-tables").getValueAsBool()))
     return false;
 
+  // In the early optimization pipeline, disable formation of lookup tables
+  // and bit maps, as they may inhibit further optimization.
+  if (!ConvertSwitchToLookupTable && (AnyLookupTables || AnyBitMaps))
+    return false;
+
   Builder.SetInsertPoint(SI);
   // TableIndex is the switch condition - TableIndexOffset if we don't
   // use the condition directly
@@ -7929,13 +7942,8 @@ bool SimplifyCFGOpt::simplifySwitch(SwitchInst *SI, IRBuilder<> &Builder) {
   if (Options.ForwardSwitchCondToPhi && forwardSwitchConditionToPHI(SI))
     return requestResimplify();
 
-  // The conversion from switch to lookup tables results in difficult-to-analyze
-  // code and makes pruning branches much harder. This is a problem if the
-  // switch expression itself can still be restricted as a result of inlining or
-  // CVP. Therefore, only apply this transformation during late stages of the
-  // optimisation pipeline.
-  if (Options.ConvertSwitchToLookupTable &&
-      simplifySwitchLookup(SI, Builder, DTU, DL, TTI))
+  if (simplifySwitchLookup(SI, Builder, DTU, DL, TTI,
+                           Options.ConvertSwitchToLookupTable))
     return requestResimplify();
 
   if (simplifySwitchOfPowersOfTwo(SI, Builder, DL, TTI))

llvm/test/Transforms/SimplifyCFG/switch-transformations-no-lut.ll

@@ -7,23 +7,11 @@ target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 define i32 @linear_transform_with_default(i32 %x) {
 ; OPTNOLUT-LABEL: define i32 @linear_transform_with_default(
 ; OPTNOLUT-SAME: i32 [[X:%.*]]) {
-; OPTNOLUT-NEXT:  [[ENTRY:.*]]:
-; OPTNOLUT-NEXT:    switch i32 [[X]], label %[[END:.*]] [
-; OPTNOLUT-NEXT:      i32 0, label %[[CASE0:.*]]
-; OPTNOLUT-NEXT:      i32 1, label %[[CASE1:.*]]
-; OPTNOLUT-NEXT:      i32 2, label %[[CASE2:.*]]
-; OPTNOLUT-NEXT:      i32 3, label %[[CASE3:.*]]
-; OPTNOLUT-NEXT:    ]
-; OPTNOLUT:       [[CASE0]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE1]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE2]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE3]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[END]]:
-; OPTNOLUT-NEXT:    [[IDX:%.*]] = phi i32 [ 1, %[[CASE0]] ], [ 4, %[[CASE1]] ], [ 7, %[[CASE2]] ], [ 10, %[[CASE3]] ], [ 13, %[[ENTRY]] ]
+; OPTNOLUT-NEXT:  [[ENTRY:.*:]]
+; OPTNOLUT-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X]], 4
+; OPTNOLUT-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[X]], 3
+; OPTNOLUT-NEXT:    [[SWITCH_OFFSET:%.*]] = add nsw i32 [[SWITCH_IDX_MULT]], 1
+; OPTNOLUT-NEXT:    [[IDX:%.*]] = select i1 [[TMP0]], i32 [[SWITCH_OFFSET]], i32 13
 ; OPTNOLUT-NEXT:    ret i32 [[IDX]]
 ;
 ; TTINOLUT-LABEL: define i32 @linear_transform_with_default(
@@ -138,26 +126,8 @@ end:
 define i32 @linear_transform_no_default(i32 %x) {
 ; OPTNOLUT-LABEL: define i32 @linear_transform_no_default(
 ; OPTNOLUT-SAME: i32 [[X:%.*]]) {
-; OPTNOLUT-NEXT:  [[ENTRY:.*]]:
-; OPTNOLUT-NEXT:    switch i32 [[X]], label %[[DEFAULT:.*]] [
-; OPTNOLUT-NEXT:      i32 0, label %[[END:.*]]
-; OPTNOLUT-NEXT:      i32 1, label %[[CASE1:.*]]
-; OPTNOLUT-NEXT:      i32 2, label %[[CASE2:.*]]
-; OPTNOLUT-NEXT:      i32 3, label %[[CASE3:.*]]
-; OPTNOLUT-NEXT:      i32 4, label %[[CASE4:.*]]
-; OPTNOLUT-NEXT:    ]
-; OPTNOLUT:       [[CASE1]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE2]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE3]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[CASE4]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[DEFAULT]]:
-; OPTNOLUT-NEXT:    unreachable
-; OPTNOLUT:       [[END]]:
-; OPTNOLUT-NEXT:    [[SWITCH_IDX_MULT:%.*]] = phi i32 [ 3, %[[CASE1]] ], [ 6, %[[CASE2]] ], [ 9, %[[CASE3]] ], [ 12, %[[CASE4]] ], [ 0, %[[ENTRY]] ]
+; OPTNOLUT-NEXT:  [[ENTRY:.*:]]
+; OPTNOLUT-NEXT:    [[SWITCH_IDX_MULT:%.*]] = mul nsw i32 [[X]], 3
 ; OPTNOLUT-NEXT:    ret i32 [[SWITCH_IDX_MULT]]
 ;
 ; TTINOLUT-LABEL: define i32 @linear_transform_no_default(
@@ -350,18 +320,9 @@ end:
 define i32 @single_value_withdefault(i32 %x) {
 ; OPTNOLUT-LABEL: define i32 @single_value_withdefault(
 ; OPTNOLUT-SAME: i32 [[X:%.*]]) {
-; OPTNOLUT-NEXT:  [[ENTRY:.*]]:
-; OPTNOLUT-NEXT:    switch i32 [[X]], label %[[DEFAULT:.*]] [
-; OPTNOLUT-NEXT:      i32 0, label %[[END:.*]]
-; OPTNOLUT-NEXT:      i32 1, label %[[END]]
-; OPTNOLUT-NEXT:      i32 2, label %[[END]]
-; OPTNOLUT-NEXT:      i32 3, label %[[END]]
-; OPTNOLUT-NEXT:      i32 4, label %[[END]]
-; OPTNOLUT-NEXT:    ]
-; OPTNOLUT:       [[DEFAULT]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[END]]:
-; OPTNOLUT-NEXT:    [[DOT:%.*]] = phi i32 [ 3, %[[DEFAULT]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ]
+; OPTNOLUT-NEXT:  [[ENTRY:.*:]]
+; OPTNOLUT-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X]], 5
+; OPTNOLUT-NEXT:    [[DOT:%.*]] = select i1 [[TMP0]], i32 2, i32 3
 ; OPTNOLUT-NEXT:    ret i32 [[DOT]]
 ;
 ; TTINOLUT-LABEL: define i32 @single_value_withdefault(
@@ -401,18 +362,9 @@ end:
 define i32 @single_value_no_jump_tables(i32 %x) "no-jump-tables"="true" {
 ; OPTNOLUT-LABEL: define i32 @single_value_no_jump_tables(
 ; OPTNOLUT-SAME: i32 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
-; OPTNOLUT-NEXT:  [[ENTRY:.*]]:
-; OPTNOLUT-NEXT:    switch i32 [[X]], label %[[DEFAULT:.*]] [
-; OPTNOLUT-NEXT:      i32 0, label %[[END:.*]]
-; OPTNOLUT-NEXT:      i32 1, label %[[END]]
-; OPTNOLUT-NEXT:      i32 2, label %[[END]]
-; OPTNOLUT-NEXT:      i32 3, label %[[END]]
-; OPTNOLUT-NEXT:      i32 4, label %[[END]]
-; OPTNOLUT-NEXT:    ]
-; OPTNOLUT:       [[DEFAULT]]:
-; OPTNOLUT-NEXT:    br label %[[END]]
-; OPTNOLUT:       [[END]]:
-; OPTNOLUT-NEXT:    [[IDX:%.*]] = phi i32 [ 3, %[[DEFAULT]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ], [ 2, %[[ENTRY]] ]
+; OPTNOLUT-NEXT:  [[ENTRY:.*:]]
+; OPTNOLUT-NEXT:    [[TMP0:%.*]] = icmp ult i32 [[X]], 5
+; OPTNOLUT-NEXT:    [[IDX:%.*]] = select i1 [[TMP0]], i32 2, i32 3
 ; OPTNOLUT-NEXT:    ret i32 [[IDX]]
 ;
 ; TTINOLUT-LABEL: define i32 @single_value_no_jump_tables(