[SimplifyCFG][ConstantMerge] Test non-table cases ahead of a switch's lookup table

[nazar-pc]

Two commits. The first is the change this is about; the second is what stops it from costing binary size.

Fixes #222292.

Despite LLM wrote most of it under my guidance, I edited, audited, tested and benchmarked the changes to the best of my ability before opening the PR.

---

The problem

simplifySwitchLookup() turns a switch whose cases feed constants into phis in one common successor into a lookup table. If a single case does not fit that shape (because it reaches a destination of its own rather than the one the others share), the whole switch is given up and every case pays for the one that did not fit.

Profile-guided builds create exactly that shape. Given a dispatch that loads a handler out of a switch and calls it, indirect call promotion rewrites the call as

if (h == @handler_7) handler_7(...) else h(...)

On the path where the switch took case 7, h is already known to be @handler_7: the switch answered that question. Jump threading sees this and rewires case 7 straight to the promoted call, deleting the test. Case 7 now reaches a destination of its own and contributes no table entry, so the table is refused. By then inlining has usually replicated the dispatch into every handler, so the loss is paid once per copy, which is how a profile makes a threaded interpreter slower than no profile at all.

The change

Set those cases aside and test them ahead of the table, leaving the rest of the switch in the shape a table can replace:

if (x == guarded) goto guarded_dest;
... lookup table for the remaining cases ...

A guarded case inside the range the table spans keeps its slot, holding poison. Nothing can index the table with a guarded value (the guard is reached first), so the slot needs no bitmask check, and no value has to be invented for it. Guarded cases are also excluded from the counts that decide whether a table is worth building.

Each guard is a test that every value reaching the table runs first, so more than one has to be paid for. Without branch weights only one case is guarded (the single one indirect call promotion leaves behind), with nothing to say it is hot. Note that this one case is guarded whatever its weight, including zero: there is no table without it either way. With weights, up to -simplifycfg-max-guarded-cases (default 3) are guarded, and only those taken at least as often as the average case that does reach the table (an integer average, so ties go to guarding).

This matters more than it sounds: indirect call promotion promotes up to three targets per site, and guarding only one leaves the switch without a table anyway. On the workload below nothing promotes a third target, so two is where the benefit was.

Behavior is unchanged where no case is set aside. The common destination is still pinned on the first case that yields results, the majority is consulted only once a case is actually being guarded. Running opt over every file in test/Transforms/SimplifyCFG with -simplifycfg-max-guarded-cases=0 produces byte-identical output to the parent commit for all 316 of them, and a separately built unpatched compiler produced a binary matching the guard-off one on the workload below.

Why the second commit

Inlining replicates a dispatch, and each copy guards a different set of cases, so each copy's table holds poison in different slots. Identical everywhere else, they cannot be merged, and each costs a table and its relocations. On the workload below that was 50 extra 445-entry tables: +178 KB of .data.rel.ro and +532 KB of .rela.dyn.

poison may be replaced by any value, so two arrays that agree wherever both are defined describe the same thing, and unifying them into the more defined of the two refines both. That needs no reasoning about which slots are read, so ConstantMerge can do it: a second phase compares poison-holding arrays against the others and merges them. The 50 tables collapse back into one.

That phase cannot use the hash map the existing merging relies on. Two constants that unify need not be equal, and no hash of one finds the other, since they disagree exactly where a hash would read. Only a constant holding poison has anything to gain, so those drive the search, each compared against the constants of its type. It merges into the one needing the fewest elements filled in, so the result does not depend on the order the globals appear in. The search is bounded by -constmerge-max-poison-candidates (1024), and giving that up costs an optimization, never correctness.

Measurements

A RISC-V interpreter in Rust: a match over the opcode yields a handler function pointer, then a tail call, with the dispatch inlined into all 448 handlers (supports vector extensions). Without a profile they share one 445-entry table. Rustc was built from source against patched and unpatched LLVM, so the only difference between columns is these commits.

	before, no PGO	before, PGO	first commit, PGO	both commits, PGO
ed25519_verify	728.1 us	770.4 us	583.2 us	597.2 us
blake3_hash_chunk	16.05 us	20.17 us	13.98 us	13.78 us
.text	3,173,358	3,230,190	2,779,246	2,785,518
.rodata	248,564	404,020	252,532	251,764
.data.rel.ro	55,384	57,352	235,360	60,928
.rela.dyn	101,016	105,048	636,960	115,776
total	3,578,322	3,796,610	3,904,098	3,213,986

With this PR, PGO finally produces both smaller and faster binaries.

ed25519_verify represents workload with a wider set of instructions (auto-vectorized), while blake3_hash_chunk is narrower (didn't auto-vectorize well).

The first commit on its own is a 2.8% size regression: 50 copies of a 445-entry table, each identical to the rest apart from the handful of slots its guards left as poison, is +178 KB of .data.rel.ro and +532 KB of .rela.dyn.

And the second commit recovers 690 KB of size with a small time cost: merging costs 2.4% on ed25519_verify while gaining 1.4% on blake3_hash_chunk.

It seemed like a good tradeoff overall to have both.

[llvmorg-github-actions]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-transforms

Author: Nazar Mokrynskyi (nazar-pc)

Changes

Two commits. The first is the change this is about; the second is what stops it from costing binary size.

Fixes #222292.

Despite LLM wrote most of it under my guidance, I edited, audited, tested and benchmarked the changes to the best of my ability before opening the PR.

---

The problem

simplifySwitchLookup() turns a switch whose cases feed constants into phis in one common successor into a lookup table. If a single case does not fit that shape (because it reaches a destination of its own rather than the one the others share), the whole switch is given up and every case pays for the one that did not fit.

Profile-guided builds create exactly that shape. Given a dispatch that loads a handler out of a switch and calls it, indirect call promotion rewrites the call as

if (h == @<!-- -->handler_7) handler_7(...) else h(...)

On the path where the switch took case 7, h is already known to be @<!-- -->handler_7: the switch answered that question. Jump threading sees this and rewires case 7 straight to the promoted call, deleting the test. Case 7 now reaches a destination of its own and contributes no table entry, so the table is refused. By then inlining has usually replicated the dispatch into every handler, so the loss is paid once per copy, which is how a profile makes a threaded interpreter slower than no profile at all.

The change

Set those cases aside and test them ahead of the table, leaving the rest of the switch in the shape a table can replace:

if (x == guarded) goto guarded_dest;
... lookup table for the remaining cases ...

A guarded case inside the range the table spans keeps its slot, holding poison. Nothing can index the table with a guarded value (the guard is reached first), so the slot needs no bitmask check, and no value has to be invented for it. Guarded cases are also excluded from the counts that decide whether a table is worth building.

Each guard is a test that every value reaching the table runs first, so more than one has to be paid for. Without branch weights only one case is guarded (the single one indirect call promotion leaves behind), with nothing to say it is hot. Note that this one case is guarded whatever its weight, including zero: there is no table without it either way. With weights, up to -simplifycfg-max-guarded-cases (default 3) are guarded, and only those taken at least as often as the average case that does reach the table (an integer average, so ties go to guarding).

This matters more than it sounds: indirect call promotion promotes up to three targets per site, and guarding only one leaves the switch without a table anyway. On the workload below nothing promotes a third target, so two is where the benefit was.

Behavior is unchanged where no case is set aside. The common destination is still pinned on the first case that yields results, the majority is consulted only once a case is actually being guarded. Running opt over every file in test/Transforms/SimplifyCFG with -simplifycfg-max-guarded-cases=0 produces byte-identical output to the parent commit for all 316 of them, and a separately built unpatched compiler produced a binary matching the guard-off one on the workload below.

Why the second commit

Inlining replicates a dispatch, and each copy guards a different set of cases, so each copy's table holds poison in different slots. Identical everywhere else, they cannot be merged, and each costs a table and its relocations. On the workload below that was 50 extra 445-entry tables: +178 KB of .data.rel.ro and +532 KB of .rela.dyn.

poison may be replaced by any value, so two arrays that agree wherever both are defined describe the same thing, and unifying them into the more defined of the two refines both. That needs no reasoning about which slots are read, so ConstantMerge can do it: a second phase compares poison-holding arrays against the others and merges them. The 50 tables collapse back into one.

That phase cannot use the hash map the existing merging relies on. Two constants that unify need not be equal, and no hash of one finds the other, since they disagree exactly where a hash would read. Only a constant holding poison has anything to gain, so those drive the search, each compared against the constants of its type. It merges into the one needing the fewest elements filled in, so the result does not depend on the order the globals appear in. The search is bounded by -constmerge-max-poison-candidates (1024), and giving that up costs an optimization, never correctness.

Measurements

A RISC-V interpreter in Rust: a match over the opcode yields a handler function pointer, then a tail call, with the dispatch inlined into all 448 handlers (supports vector extensions). Without a profile they share one 445-entry table. Rustc was built from source against patched and unpatched LLVM, so the only difference between columns is these commits.

	before, no PGO	before, PGO	first commit, PGO	both commits, PGO
ed25519_verify	728.1 us	770.4 us	583.2 us	597.2 us
blake3_hash_chunk	16.05 us	20.17 us	13.98 us	13.78 us
.text	3,173,358	3,230,190	2,779,246	2,785,518
.rodata	248,564	404,020	252,532	251,764
.data.rel.ro	55,384	57,352	235,360	60,928
.rela.dyn	101,016	105,048	636,960	115,776
total	3,578,322	3,796,610	3,904,098	3,213,986

With this PR, PGO finally produces both smaller and faster binaries.

ed25519_verify represents workload with a wider set of instructions (auto-vectorized), while blake3_hash_chunk is narrower (didn't auto-vectorize well).

The first commit on its own is a 2.8% size regression: 50 copies of a 445-entry table, each identical to the rest apart from the handful of slots its guards left as poison, is +178 KB of .data.rel.ro and +532 KB of .rela.dyn.

And the second commit recovers 690 KB of size with a small time cost: merging costs 2.4% on ed25519_verify while gaining 1.4% on blake3_hash_chunk.

It seemed like a good tradeoff overall to have both.

---

Patch is 61.82 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/223255.diff

5 Files Affected:

• (modified) llvm/lib/Transforms/IPO/ConstantMerge.cpp (+124)
• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+247-28)
• (added) llvm/test/Transforms/ConstantMerge/merge-poison-compatible.ll (+62)
• (modified) llvm/test/Transforms/SimplifyCFG/X86/switch_to_lookup_table.ll (+10-15)
• (added) llvm/test/Transforms/SimplifyCFG/switch-guard-non-table-case.ll (+961)

diff --git a/llvm/lib/Transforms/IPO/ConstantMerge.cpp b/llvm/lib/Transforms/IPO/ConstantMerge.cpp
index 480f4c5f68f790..2bfc46478300e2 100644
--- a/llvm/lib/Transforms/IPO/ConstantMerge.cpp
+++ b/llvm/lib/Transforms/IPO/ConstantMerge.cpp
@@ -18,6 +18,7 @@
 
 #include "llvm/Transforms/IPO/ConstantMerge.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
@@ -29,6 +30,7 @@
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/IPO.h"
 #include <algorithm>
@@ -39,7 +41,14 @@ using namespace llvm;
 
 #define DEBUG_TYPE "constmerge"
 
+static cl::opt<unsigned> MaxPoisonCandidates(
+    "constmerge-max-poison-candidates", cl::Hidden, cl::init(1024),
+    cl::desc("Largest set of constants of one type that a constant holding "
+             "poison is compared against. A larger set is left alone"));
+
 STATISTIC(NumIdenticalMerged, "Number of identical global constants merged");
+STATISTIC(NumPoisonMerged,
+          "Number of global constants merged that differed only in poison");
 
 /// Find values that are marked as llvm.used.
 static void FindUsedValues(GlobalVariable *LLVMUsed,
@@ -135,6 +144,119 @@ static void replace(Module &M, GlobalVariable *Old, GlobalVariable *New) {
   Old->eraseFromParent();
 }
 
+/// Unify two array constants that differ only where one of them holds poison.
+///
+/// Poison may be replaced by any value, so a table holding poison at an index
+/// and a table holding a defined value there describe the same thing: taking
+/// the defined value refines both. Returns the unified constant, or null if the
+/// two disagree anywhere they are both defined.
+///
+/// Copies of an inlined switch produce exactly this shape when each copy tests
+/// a different case ahead of its lookup table: every copy's table is the same
+/// but for the slots standing in for the cases that copy tested, which nothing
+/// loads and which are therefore poison.
+static Constant *unifyPoison(Constant *A, Constant *B) {
+  if (A == B)
+    return A;
+  auto *ATy = dyn_cast<ArrayType>(A->getType());
+  if (!ATy || A->getType() != B->getType())
+    return nullptr;
+
+  uint64_t N = ATy->getNumElements();
+  SmallVector<Constant *, 64> Unified(N);
+  for (uint64_t I = 0; I != N; ++I) {
+    Constant *EA = A->getAggregateElement(I);
+    Constant *EB = B->getAggregateElement(I);
+    if (!EA || !EB)
+      return nullptr;
+    if (EA == EB) {
+      Unified[I] = EA;
+      continue;
+    }
+    if (isa<PoisonValue>(EA)) {
+      Unified[I] = EB;
+      continue;
+    }
+    if (isa<PoisonValue>(EB)) {
+      Unified[I] = EA;
+      continue;
+    }
+    return nullptr;
+  }
+  return ConstantArray::get(ATy, Unified);
+}
+
+/// Merge globals whose initializers differ only in poison elements.
+///
+/// Kept apart from the identical-initializer merging above because it cannot
+/// use a hash map. Two constants that unify need not be equal, and no hash of
+/// one of them can find the other, since they disagree exactly where a hash
+/// would read. Only a constant that holds poison has anything to gain, so those
+/// drive the search, and each is compared against the constants of its type,
+/// bounded by MaxPoisonCandidates. Giving that bound up costs an optimisation,
+/// never correctness.
+static size_t
+mergePoisonCompatible(Module &M,
+                      const SmallPtrSetImpl<const GlobalValue *> &UsedGlobals) {
+  // Candidates, grouped by type so only plausible pairs are compared.
+  MapVector<Type *, SmallVector<GlobalVariable *, 8>> ByType;
+  SmallVector<GlobalVariable *, 8> HoldsPoison;
+  for (GlobalVariable &GV : M.globals()) {
+    // The same conditions the merging above puts on a constant before it may
+    // stand in for another. A global it refuses to make canonical must not
+    // become one here either, since replace() assumes it could have.
+    if (isUnmergeableGlobal(&GV, UsedGlobals) || !GV.hasLocalLinkage() ||
+        GV.isWeakForLinker() || GV.hasMetadataOtherThanDebugLocAndGuid())
+      continue;
+    // Only ConstantArray can hold poison. An array of defined integers is a
+    // ConstantDataArray, which nothing here can unify with anyway.
+    auto *Init = dyn_cast<ConstantArray>(GV.getInitializer());
+    if (!Init)
+      continue;
+    ByType[GV.getValueType()].push_back(&GV);
+    if (any_of(Init->operands(),
+               [](const Use &U) { return isa<PoisonValue>(U.get()); }))
+      HoldsPoison.push_back(&GV);
+  }
+
+  size_t Merged = 0;
+  for (GlobalVariable *GV : HoldsPoison) {
+    SmallVectorImpl<GlobalVariable *> &Candidates = ByType[GV->getValueType()];
+    if (Candidates.size() > MaxPoisonCandidates)
+      continue;
+
+    // Merge into the constant that needs the fewest elements filled in, so that
+    // the outcome does not depend on the order the globals happen to appear in.
+    // Ties go to the earlier one, which keeps it deterministic.
+    GlobalVariable *Into = nullptr;
+    Constant *Unified = nullptr;
+    unsigned FewestPoison = 0;
+    for (GlobalVariable *C : Candidates) {
+      if (C == GV || C->getParent() != &M)
+        continue;
+      Constant *U = unifyPoison(C->getInitializer(), GV->getInitializer());
+      if (!U)
+        continue;
+      unsigned Poison =
+          count_if(cast<ConstantArray>(C->getInitializer())->operands(),
+                   [](const Use &Op) { return isa<PoisonValue>(Op.get()); });
+      if (Into && Poison >= FewestPoison)
+        continue;
+      Into = C;
+      Unified = U;
+      FewestPoison = Poison;
+    }
+    if (!Into || makeMergeable(GV, Into) == CanMerge::No)
+      continue;
+
+    Into->setInitializer(Unified);
+    replace(M, GV, Into);
+    ++Merged;
+    ++NumPoisonMerged;
+  }
+  return Merged;
+}
+
 static bool mergeConstants(Module &M) {
   // Find all the globals that are marked "used".  These cannot be merged.
   SmallPtrSet<const GlobalValue*, 8> UsedGlobals;
@@ -244,6 +366,8 @@ static bool mergeConstants(Module &M) {
     CMap.clear();
   }
 
+  ChangesMade += mergePoisonCompatible(M, UsedGlobals);
+
   return ChangesMade;
 }
 
diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index ca96f2e70d8106..458465e2a14010 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -200,6 +200,11 @@ static cl::opt<unsigned> MaxSwitchCasesPerResult(
     "max-switch-cases-per-result", cl::Hidden, cl::init(16),
     cl::desc("Limit cases to analyze when converting a switch to select"));
 
+static cl::opt<unsigned> MaxGuardedCases(
+    "simplifycfg-max-guarded-cases", cl::Hidden, cl::init(3),
+    cl::desc("Most cases a switch may have tested ahead of its lookup table; "
+             "zero gives up the table instead, as before"));
+
 static cl::opt<unsigned> MaxJumpThreadingLiveBlocks(
     "max-jump-threading-live-blocks", cl::Hidden, cl::init(24),
     cl::desc("Limit number of blocks a define in a threaded block is allowed "
@@ -217,6 +222,9 @@ STATISTIC(NumLookupTables,
 STATISTIC(
     NumLookupTablesHoles,
     "Number of switch instructions turned into lookup tables (holes checked)");
+STATISTIC(NumLookupTablesGuardedCase,
+          "Number of switch lookup tables that needed a case tested ahead "
+          "of them");
 STATISTIC(NumTableCmpReuses, "Number of reused switch table lookup compares");
 STATISTIC(NumFoldValueComparisonIntoPredecessors,
           "Number of value comparisons folded into predecessor basic blocks");
@@ -6627,15 +6635,16 @@ getCaseResults(SwitchInst *SI, ConstantInt *CaseVal, BasicBlock *CaseDest,
     }
   }
 
-  // If we did not have a CommonDest before, use the current one.
-  if (!*CommonDest)
-    *CommonDest = CaseDest;
+  // Hold off on adopting this case's destination as the common one until the
+  // case is known to be usable, so that a case which fails below leaves the
+  // caller no worse off than if it had never been looked at.
+  BasicBlock *NewCommonDest = *CommonDest ? *CommonDest : CaseDest;
   // If the destination isn't the common one, abort.
-  if (CaseDest != *CommonDest)
+  if (CaseDest != NewCommonDest)
     return false;
 
   // Get the values for this case from phi nodes in the destination block.
-  for (PHINode &PHI : (*CommonDest)->phis()) {
+  for (PHINode &PHI : NewCommonDest->phis()) {
     int Idx = PHI.getBasicBlockIndex(Pred);
     if (Idx == -1)
       continue;
@@ -6652,7 +6661,11 @@ getCaseResults(SwitchInst *SI, ConstantInt *CaseVal, BasicBlock *CaseDest,
     Res.push_back(std::make_pair(&PHI, ConstVal));
   }
 
-  return Res.size() > 0;
+  if (Res.empty())
+    return false;
+
+  *CommonDest = NewCommonDest;
+  return true;
 }
 
 // Helper function used to add CaseVal to the list of cases that generate
@@ -7366,11 +7379,12 @@ getDenseSwitchRangeReductionShift(ArrayRef<int64_t> Values, int64_t Base,
 // TODO: We could support larger than legal types by limiting based on the
 // number of loads required and/or table size. If the constants are small we
 // could use smaller table entries and extend after the load.
-static bool shouldBuildLookupTable(SwitchInst *SI, uint64_t TableSize,
+static bool shouldBuildLookupTable(SwitchInst *SI, uint64_t NumCases,
+                                   uint64_t TableSize,
                                    const TargetTransformInfo &TTI,
                                    const DataLayout &DL,
                                    const SmallVector<Type *> &ResultTypes) {
-  if (SI->getNumCases() > TableSize)
+  if (NumCases > TableSize)
     return false; // TableSize overflowed.
 
   bool AllTablesFitInRegister = true;
@@ -7399,8 +7413,7 @@ static bool shouldBuildLookupTable(SwitchInst *SI, uint64_t TableSize,
   if (HasIllegalType)
     return false;
 
-  return isSwitchDense(SI->getNumCases(), TableSize,
-                       SI->getFunction()->hasOptSize());
+  return isSwitchDense(NumCases, TableSize, SI->getFunction()->hasOptSize());
 }
 
 static bool shouldUseSwitchConditionAsTableIndex(
@@ -7529,45 +7542,179 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
   // common destination, as well as the min and max case values.
   assert(!SI->cases().empty());
   SwitchInst::CaseIt CI = SI->case_begin();
-  ConstantInt *MinCaseVal = CI->getCaseValue();
-  ConstantInt *MaxCaseVal = CI->getCaseValue();
+  // Over the cases that make it into the table, so that a guarded case far away
+  // from the others does not stretch the table to cover a range nothing reads.
+  ConstantInt *MinCaseVal = nullptr;
+  ConstantInt *MaxCaseVal = nullptr;
 
   BasicBlock *CommonDest = nullptr;
 
   using ResultListTy = SmallVector<std::pair<ConstantInt *, Constant *>, 4>;
   SmallDenseMap<PHINode *, ResultListTy> ResultLists;
 
+  // Cases a table cannot hold, to be tested ahead of it.
+  struct GuardedCase {
+    ConstantInt *Val;
+    BasicBlock *Dest;
+    bool HasSlot = false;
+  };
+  SmallVector<GuardedCase, 4> Guarded;
+
   SmallDenseMap<PHINode *, Constant *> DefaultResults;
   SmallVector<Type *> ResultTypes;
   SmallVector<PHINode *, 4> PHIs;
 
+  // Resolve each case on its own first. Letting the first case that resolves
+  // decide the common destination would make the transform depend on the order
+  // the cases happen to be written in, since the case left over is then
+  // whichever one disagrees with that choice rather than the odd one out.
+  using ResultsTy = SmallVector<std::pair<PHINode *, Constant *>, 4>;
+  struct ResolvedCase {
+    ConstantInt *Val;
+    BasicBlock *Succ;
+    BasicBlock *Dest; // Null if the case has no constant to contribute.
+    ResultsTy Results;
+  };
+  SmallVector<ResolvedCase> Cases;
+  SmallMapVector<BasicBlock *, unsigned, 8> DestCounts;
   for (SwitchInst::CaseIt E = SI->case_end(); CI != E; ++CI) {
-    ConstantInt *CaseVal = CI->getCaseValue();
-    if (CaseVal->getValue().slt(MinCaseVal->getValue()))
-      MinCaseVal = CaseVal;
-    if (CaseVal->getValue().sgt(MaxCaseVal->getValue()))
-      MaxCaseVal = CaseVal;
+    ResolvedCase RC{CI->getCaseValue(), CI->getCaseSuccessor(), nullptr, {}};
+    BasicBlock *Dest = nullptr;
+    if (getCaseResults(SI, RC.Val, RC.Succ, &Dest, RC.Results, DL, TTI)) {
+      RC.Dest = Dest;
+      ++DestCounts[Dest];
+    } else {
+      RC.Results.clear();
+    }
+    Cases.push_back(std::move(RC));
+  }
 
-    // Resulting value at phi nodes for this case value.
-    using ResultsTy = SmallVector<std::pair<PHINode *, Constant *>, 4>;
-    ResultsTy Results;
-    if (!getCaseResults(SI, CaseVal, CI->getCaseSuccessor(), &CommonDest,
-                        Results, DL, TTI))
-      return false;
+  // The destination most of the cases agree on is the one a table can stand in
+  // for. Where nothing is guarded they all agree, so this is the destination
+  // the first case reaches, exactly as before. Where one is, taking the
+  // majority keeps an odd first case from pinning a destination the rest
+  // disagree with and rejecting a switch that is nearly all table.
+  unsigned BestCount = 0;
+  for (const auto &[Dest, Count] : DestCounts)
+    if (Count > BestCount) {
+      BestCount = Count;
+      CommonDest = Dest;
+    }
+  if (!CommonDest)
+    return false;
+
+  for (const ResolvedCase &RC : Cases) {
+    if (RC.Dest != CommonDest) {
+      // This case does not reach the common destination with a constant, so a
+      // table cannot stand in for it: it has to go somewhere else and do
+      // something else. Set it aside rather than give up the table. If it is
+      // the only such case it is tested ahead of the table, which leaves the
+      // rest of the switch in the shape a table can replace. Indirect call
+      // promotion creates exactly this shape, by giving the profiled target of
+      // a dispatch a destination of its own.
+      if (Guarded.size() >= (size_t)MaxGuardedCases)
+        return false;
+      // A case that cannot be reached is not worth a test of its own, and
+      // SimplifyCFG removes it in its own time, leaving the hole the table
+      // already knows how to fill.
+      if (isa<UnreachableInst>(RC.Succ->getFirstNonPHIIt()))
+        return false;
+      Guarded.push_back({RC.Val, RC.Succ});
+      continue;
+    }
+
+    if (!MinCaseVal || RC.Val->getValue().slt(MinCaseVal->getValue()))
+      MinCaseVal = RC.Val;
+    if (!MaxCaseVal || RC.Val->getValue().sgt(MaxCaseVal->getValue()))
+      MaxCaseVal = RC.Val;
 
     // Append the result and result types from this case to the list for each
     // phi.
-    for (const auto &I : Results) {
+    for (const auto &I : RC.Results) {
       PHINode *PHI = I.first;
       Constant *Value = I.second;
       auto [It, Inserted] = ResultLists.try_emplace(PHI);
       if (Inserted)
         PHIs.push_back(PHI);
-      It->second.push_back(std::make_pair(CaseVal, Value));
+      It->second.push_back(std::make_pair(RC.Val, Value));
       ResultTypes.push_back(PHI->getType());
     }
   }
 
+  // A guarded case inside the range the others span owns a slot in the table.
+  // One outside it does not, and needs nothing done for it.
+  for (GuardedCase &G : Guarded)
+    G.HasSlot = MinCaseVal && G.Val->getValue().sgt(MinCaseVal->getValue()) &&
+                G.Val->getValue().slt(MaxCaseVal->getValue());
+
+  // Decide whether the case set aside above can really be guarded. Nothing has
+  // been mutated yet, so this can still give up on the whole transform.
+  SmallDenseMap<ConstantInt *, uint64_t> CaseWeights;
+  uint64_t DefaultWeight = 0;
+  bool HaveWeights = false;
+  bool WeightsAreExpected = false;
+  if (!Guarded.empty()) {
+    // The table is the point of the guard, so there has to be one worth
+    // building without the guarded cases. Everything else is left to the
+    // decisions the table already has to pass below.
+    if (SI->getNumCases() - Guarded.size() < 3)
+      return false;
+
+    SmallVector<uint32_t> Weights;
+    HaveWeights = extractBranchWeights(*SI, Weights) &&
+                  Weights.size() == SI->getNumSuccessors();
+    if (HaveWeights) {
+      WeightsAreExpected = hasBranchWeightOrigin(*SI);
+      DefaultWeight = Weights[0];
+      for (const auto &Case : SI->cases())
+        CaseWeights[Case.getCaseValue()] = Weights[Case.getSuccessorIndex()];
+    }
+
+    // Each guarded case is a test every value that does reach the table has to
+    // run first, so more than one has to be paid for. Without a profile there
+    // is nothing to say the case is taken often enough to be worth the test, so
+    // only the single case indirect call promotion leaves behind is guarded.
+    // With one, a case earns its test by being taken more often than the
+    // average case that reaches the table.
+    if (Guarded.size() > 1) {
+      if (!HaveWeights)
+        return false;
+      uint64_t TableWeight = 0, TableCases = 0;
+      for (const auto &Case : SI->cases())
+        if (none_of(Guarded, [&](const GuardedCase &G) {
+              return G.Val == Case.getCaseValue();
+            })) {
+          TableWeight += CaseWeights.lookup(Case.getCaseValue());
+          ++TableCases;
+        }
+      uint64_t Average = TableCases ? TableWeight / TableCases : 0;
+      for (const GuardedCase &G : Guarded)
+        if (CaseWeights.lookup(G.Val) < Average)
+          return false;
+    }
+
+    // Keep the rewrite simple by requiring the two halves to be disjoint, so no
+    // phi ends up with edges from both switches.
+    for (const GuardedCase &G : Guarded) {
+      if (G.Dest == SI->getDefaultDest())
+        return false;
+      for (const auto &Case : SI->cases())
+        if (none_of(Guarded,
+                    [&](const GuardedCase &O) {
+                      return O.Val == Case.getCaseValue();
+                    }) &&
+            Case.getCaseSuccessor() == G.Dest)
+          return false;
+    }
+  }
+
+  assert(MinCaseVal && MaxCaseVal && "No case reaches the common destination?");
+
+  // The guarded case never reaches the table, so it does not count towards the
+  // table's density, its size, or the case count that decides whether a hole
+  // check is worth paying for.
+  uint64_t NumTableCases = SI->getNumCases() - Guarded.size();
+
   // If the table has holes, we need a constant result for the default case
   // or a bitmask that fits in a register.
   SmallVector<std::pair<PHINode *, Constant *>, 4> DefaultResultsList;
@@ -7580,6 +7727,22 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
     DefaultResults[PHI] = Result;
   }
 
+  // Give each guarded slot poison. Nothing can index the table with a guarded
+  // value, since the guard is reached first, so the slot is unreachable rather
+  // than a hole: it needs no bitmask check, and no value has to be invented for
+  // it. Copies of a dispatch guard different cases, so it also leaves their
+  // tables agreeing wherever both are defined, which is what lets ConstantMerge
+  // fold them back into one.
+  for (const GuardedCase &G : Guarded) {
+    if (!G.HasSlot)
+      continue;
+    for (PHINode *PHI : PHIs) {
+      ResultListTy &ResultList = ResultLists[PHI];
+      Type *ResultType = ResultList.front().second->getType();
+      ResultList.emplace_back(G.Val, PoisonValue::get(ResultType));
+    }
+  }
+
   bool UseSwitchConditionAsTableIndex = shouldUseSwitchConditionAsTableIndex(
       *MinCaseVal, *MaxCaseVal, HasDefaultResults, ResultTypes, DL, TTI);
   uint64_t TableSize;
@@ -7616,13 +7779,14 @@ static bool simplifySwitchLookup(SwitchInst *SI, IRBuilder<> &Builder,
   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).
+    if (NumTableCases < 4) // FIXME: Find best threshold value (benchmark).
       return false;
     if (!DL.fitsInLegalInteger(TableSize))
       return false;
   }
 
-  if (!shouldBuildLookupTable(SI, TableSize, TTI, DL, ResultTypes))
+  if (!shouldBuildLookupTable(SI, NumTableCases, TableSize, TTI, DL,
+                              ResultTypes))
     return false;
 
   // Compute the table index valu...
[truncated]

[nazar-pc]

I can split two commits into separate PRs after review, but it makes the most logical sense to land them together.

[nazar-pc]

@dtcxzyw, assuming https://github.com/dtcxzyw/llvm-opt-benchmark-nightly uses PGO, this PR would be useful to test there. Otherwise it might be difficult to measure its impact.