| 1 | //===- SwitchLoweringUtils.cpp - Switch Lowering --------------------------===// |
| 2 | // |
| 3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| 4 | // See https://llvm.org/LICENSE.txt for license information. |
| 5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | // |
| 7 | //===----------------------------------------------------------------------===// |
| 8 | // |
| 9 | // This file contains switch inst lowering optimizations and utilities for |
| 10 | // codegen, so that it can be used for both SelectionDAG and GlobalISel. |
| 11 | // |
| 12 | //===----------------------------------------------------------------------===// |
| 13 | |
| 14 | #include "llvm/CodeGen/SwitchLoweringUtils.h" |
| 15 | #include "llvm/CodeGen/FunctionLoweringInfo.h" |
| 16 | #include "llvm/CodeGen/MachineJumpTableInfo.h" |
| 17 | #include "llvm/CodeGen/TargetLowering.h" |
| 18 | #include "llvm/Target/TargetMachine.h" |
| 19 | |
| 20 | using namespace llvm; |
| 21 | using namespace SwitchCG; |
| 22 | |
| 23 | uint64_t SwitchCG::getJumpTableRange(const CaseClusterVector &Clusters, |
| 24 | unsigned First, unsigned Last) { |
| 25 | assert(Last >= First); |
| 26 | const APInt &LowCase = Clusters[First].Low->getValue(); |
| 27 | const APInt &HighCase = Clusters[Last].High->getValue(); |
| 28 | assert(LowCase.getBitWidth() == HighCase.getBitWidth()); |
| 29 | |
| 30 | // FIXME: A range of consecutive cases has 100% density, but only requires one |
| 31 | // comparison to lower. We should discriminate against such consecutive ranges |
| 32 | // in jump tables. |
| 33 | return (HighCase - LowCase).getLimitedValue(Limit: (UINT64_MAX - 1) / 100) + 1; |
| 34 | } |
| 35 | |
| 36 | uint64_t |
| 37 | SwitchCG::getJumpTableNumCases(const SmallVectorImpl<unsigned> &TotalCases, |
| 38 | unsigned First, unsigned Last) { |
| 39 | assert(Last >= First); |
| 40 | assert(TotalCases[Last] >= TotalCases[First]); |
| 41 | uint64_t NumCases = |
| 42 | TotalCases[Last] - (First == 0 ? 0 : TotalCases[First - 1]); |
| 43 | return NumCases; |
| 44 | } |
| 45 | |
| 46 | void SwitchCG::SwitchLowering::findJumpTables(CaseClusterVector &Clusters, |
| 47 | const SwitchInst *SI, |
| 48 | std::optional<SDLoc> SL, |
| 49 | MachineBasicBlock *DefaultMBB, |
| 50 | ProfileSummaryInfo *PSI, |
| 51 | BlockFrequencyInfo *BFI) { |
| 52 | #ifndef NDEBUG |
| 53 | // Clusters must be non-empty, sorted, and only contain Range clusters. |
| 54 | assert(!Clusters.empty()); |
| 55 | for (CaseCluster &C : Clusters) |
| 56 | assert(C.Kind == CC_Range); |
| 57 | for (unsigned i = 1, e = Clusters.size(); i < e; ++i) |
| 58 | assert(Clusters[i - 1].High->getValue().slt(Clusters[i].Low->getValue())); |
| 59 | #endif |
| 60 | |
| 61 | assert(TLI && "TLI not set!" ); |
| 62 | if (!TLI->areJTsAllowed(Fn: SI->getParent()->getParent())) |
| 63 | return; |
| 64 | |
| 65 | const unsigned MinJumpTableEntries = TLI->getMinimumJumpTableEntries(); |
| 66 | const unsigned SmallNumberOfEntries = MinJumpTableEntries / 2; |
| 67 | |
| 68 | // Bail if not enough cases. |
| 69 | const int64_t N = Clusters.size(); |
| 70 | if (N < 2 || N < MinJumpTableEntries) |
| 71 | return; |
| 72 | |
| 73 | // Accumulated number of cases in each cluster and those prior to it. |
| 74 | SmallVector<unsigned, 8> TotalCases(N); |
| 75 | for (unsigned i = 0; i < N; ++i) { |
| 76 | const APInt &Hi = Clusters[i].High->getValue(); |
| 77 | const APInt &Lo = Clusters[i].Low->getValue(); |
| 78 | TotalCases[i] = (Hi - Lo).getLimitedValue() + 1; |
| 79 | if (i != 0) |
| 80 | TotalCases[i] += TotalCases[i - 1]; |
| 81 | } |
| 82 | |
| 83 | uint64_t Range = getJumpTableRange(Clusters,First: 0, Last: N - 1); |
| 84 | uint64_t NumCases = getJumpTableNumCases(TotalCases, First: 0, Last: N - 1); |
| 85 | assert(NumCases < UINT64_MAX / 100); |
| 86 | assert(Range >= NumCases); |
| 87 | |
| 88 | // Cheap case: the whole range may be suitable for jump table. |
| 89 | if (TLI->isSuitableForJumpTable(SI, NumCases, Range, PSI, BFI)) { |
| 90 | CaseCluster JTCluster; |
| 91 | if (buildJumpTable(Clusters, First: 0, Last: N - 1, SI, SL, DefaultMBB, JTCluster)) { |
| 92 | Clusters[0] = JTCluster; |
| 93 | Clusters.resize(new_size: 1); |
| 94 | return; |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | // The algorithm below is not suitable for -O0. |
| 99 | if (TM->getOptLevel() == CodeGenOptLevel::None) |
| 100 | return; |
| 101 | |
| 102 | // Split Clusters into minimum number of dense partitions. The algorithm uses |
| 103 | // the same idea as Kannan & Proebsting "Correction to 'Producing Good Code |
| 104 | // for the Case Statement'" (1994), but builds the MinPartitions array in |
| 105 | // reverse order to make it easier to reconstruct the partitions in ascending |
| 106 | // order. In the choice between two optimal partitionings, it picks the one |
| 107 | // which yields more jump tables. The algorithm is described in |
| 108 | // https://arxiv.org/pdf/1910.02351v2 |
| 109 | |
| 110 | // MinPartitions[i] is the minimum nbr of partitions of Clusters[i..N-1]. |
| 111 | SmallVector<unsigned, 8> MinPartitions(N); |
| 112 | // LastElement[i] is the last element of the partition starting at i. |
| 113 | SmallVector<unsigned, 8> LastElement(N); |
| 114 | // PartitionsScore[i] is used to break ties when choosing between two |
| 115 | // partitionings resulting in the same number of partitions. |
| 116 | SmallVector<unsigned, 8> PartitionsScore(N); |
| 117 | // For PartitionsScore, a small number of comparisons is considered as good as |
| 118 | // a jump table and a single comparison is considered better than a jump |
| 119 | // table. |
| 120 | enum PartitionScores : unsigned { |
| 121 | NoTable = 0, |
| 122 | Table = 1, |
| 123 | FewCases = 1, |
| 124 | SingleCase = 2 |
| 125 | }; |
| 126 | |
| 127 | // Base case: There is only one way to partition Clusters[N-1]. |
| 128 | MinPartitions[N - 1] = 1; |
| 129 | LastElement[N - 1] = N - 1; |
| 130 | PartitionsScore[N - 1] = PartitionScores::SingleCase; |
| 131 | |
| 132 | // Note: loop indexes are signed to avoid underflow. |
| 133 | for (int64_t i = N - 2; i >= 0; i--) { |
| 134 | // Find optimal partitioning of Clusters[i..N-1]. |
| 135 | // Baseline: Put Clusters[i] into a partition on its own. |
| 136 | MinPartitions[i] = MinPartitions[i + 1] + 1; |
| 137 | LastElement[i] = i; |
| 138 | PartitionsScore[i] = PartitionsScore[i + 1] + PartitionScores::SingleCase; |
| 139 | |
| 140 | // Search for a solution that results in fewer partitions. |
| 141 | for (int64_t j = N - 1; j > i; j--) { |
| 142 | // Try building a partition from Clusters[i..j]. |
| 143 | Range = getJumpTableRange(Clusters, First: i, Last: j); |
| 144 | NumCases = getJumpTableNumCases(TotalCases, First: i, Last: j); |
| 145 | assert(NumCases < UINT64_MAX / 100); |
| 146 | assert(Range >= NumCases); |
| 147 | |
| 148 | if (TLI->isSuitableForJumpTable(SI, NumCases, Range, PSI, BFI)) { |
| 149 | unsigned NumPartitions = 1 + (j == N - 1 ? 0 : MinPartitions[j + 1]); |
| 150 | unsigned Score = j == N - 1 ? 0 : PartitionsScore[j + 1]; |
| 151 | int64_t NumEntries = j - i + 1; |
| 152 | |
| 153 | if (NumEntries == 1) |
| 154 | Score += PartitionScores::SingleCase; |
| 155 | else if (NumEntries <= SmallNumberOfEntries) |
| 156 | Score += PartitionScores::FewCases; |
| 157 | else if (NumEntries >= MinJumpTableEntries) |
| 158 | Score += PartitionScores::Table; |
| 159 | |
| 160 | // If this leads to fewer partitions, or to the same number of |
| 161 | // partitions with better score, it is a better partitioning. |
| 162 | if (NumPartitions < MinPartitions[i] || |
| 163 | (NumPartitions == MinPartitions[i] && Score > PartitionsScore[i])) { |
| 164 | MinPartitions[i] = NumPartitions; |
| 165 | LastElement[i] = j; |
| 166 | PartitionsScore[i] = Score; |
| 167 | } |
| 168 | } |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | // Iterate over the partitions, replacing some with jump tables in-place. |
| 173 | unsigned DstIndex = 0; |
| 174 | for (unsigned First = 0, Last; First < N; First = Last + 1) { |
| 175 | Last = LastElement[First]; |
| 176 | assert(Last >= First); |
| 177 | assert(DstIndex <= First); |
| 178 | unsigned NumClusters = Last - First + 1; |
| 179 | |
| 180 | CaseCluster JTCluster; |
| 181 | if (NumClusters >= MinJumpTableEntries && |
| 182 | buildJumpTable(Clusters, First, Last, SI, SL, DefaultMBB, JTCluster)) { |
| 183 | Clusters[DstIndex++] = JTCluster; |
| 184 | } else { |
| 185 | for (unsigned I = First; I <= Last; ++I) |
| 186 | std::memmove(dest: &Clusters[DstIndex++], src: &Clusters[I], n: sizeof(Clusters[I])); |
| 187 | } |
| 188 | } |
| 189 | Clusters.resize(new_size: DstIndex); |
| 190 | } |
| 191 | |
| 192 | bool SwitchCG::SwitchLowering::buildJumpTable(const CaseClusterVector &Clusters, |
| 193 | unsigned First, unsigned Last, |
| 194 | const SwitchInst *SI, |
| 195 | const std::optional<SDLoc> &SL, |
| 196 | MachineBasicBlock *DefaultMBB, |
| 197 | CaseCluster &JTCluster) { |
| 198 | assert(First <= Last); |
| 199 | |
| 200 | auto Prob = BranchProbability::getZero(); |
| 201 | unsigned NumCmps = 0; |
| 202 | std::vector<MachineBasicBlock*> Table; |
| 203 | DenseMap<MachineBasicBlock*, BranchProbability> JTProbs; |
| 204 | |
| 205 | // Initialize probabilities in JTProbs. |
| 206 | for (unsigned I = First; I <= Last; ++I) |
| 207 | JTProbs[Clusters[I].MBB] = BranchProbability::getZero(); |
| 208 | |
| 209 | for (unsigned I = First; I <= Last; ++I) { |
| 210 | assert(Clusters[I].Kind == CC_Range); |
| 211 | Prob += Clusters[I].Prob; |
| 212 | const APInt &Low = Clusters[I].Low->getValue(); |
| 213 | const APInt &High = Clusters[I].High->getValue(); |
| 214 | NumCmps += (Low == High) ? 1 : 2; |
| 215 | if (I != First) { |
| 216 | // Fill the gap between this and the previous cluster. |
| 217 | const APInt &PreviousHigh = Clusters[I - 1].High->getValue(); |
| 218 | assert(PreviousHigh.slt(Low)); |
| 219 | uint64_t Gap = (Low - PreviousHigh).getLimitedValue() - 1; |
| 220 | for (uint64_t J = 0; J < Gap; J++) |
| 221 | Table.push_back(x: DefaultMBB); |
| 222 | } |
| 223 | uint64_t ClusterSize = (High - Low).getLimitedValue() + 1; |
| 224 | for (uint64_t J = 0; J < ClusterSize; ++J) |
| 225 | Table.push_back(x: Clusters[I].MBB); |
| 226 | JTProbs[Clusters[I].MBB] += Clusters[I].Prob; |
| 227 | } |
| 228 | |
| 229 | unsigned NumDests = JTProbs.size(); |
| 230 | if (TLI->isSuitableForBitTests(NumDests, NumCmps, |
| 231 | Low: Clusters[First].Low->getValue(), |
| 232 | High: Clusters[Last].High->getValue(), DL: *DL)) { |
| 233 | // Clusters[First..Last] should be lowered as bit tests instead. |
| 234 | return false; |
| 235 | } |
| 236 | |
| 237 | // Create the MBB that will load from and jump through the table. |
| 238 | // Note: We create it here, but it's not inserted into the function yet. |
| 239 | MachineFunction *CurMF = FuncInfo.MF; |
| 240 | MachineBasicBlock *JumpTableMBB = |
| 241 | CurMF->CreateMachineBasicBlock(BB: SI->getParent()); |
| 242 | |
| 243 | // Add successors. Note: use table order for determinism. |
| 244 | SmallPtrSet<MachineBasicBlock *, 8> Done; |
| 245 | for (MachineBasicBlock *Succ : Table) { |
| 246 | if (Done.count(Ptr: Succ)) |
| 247 | continue; |
| 248 | addSuccessorWithProb(Src: JumpTableMBB, Dst: Succ, Prob: JTProbs[Succ]); |
| 249 | Done.insert(Ptr: Succ); |
| 250 | } |
| 251 | JumpTableMBB->normalizeSuccProbs(); |
| 252 | |
| 253 | unsigned JTI = CurMF->getOrCreateJumpTableInfo(JTEntryKind: TLI->getJumpTableEncoding()) |
| 254 | ->createJumpTableIndex(DestBBs: Table); |
| 255 | |
| 256 | // Set up the jump table info. |
| 257 | JumpTable JT(Register(), JTI, JumpTableMBB, nullptr, SL); |
| 258 | JumpTableHeader JTH(Clusters[First].Low->getValue(), |
| 259 | Clusters[Last].High->getValue(), SI->getCondition(), |
| 260 | nullptr, false); |
| 261 | JTCases.emplace_back(args: std::move(JTH), args: std::move(JT)); |
| 262 | |
| 263 | JTCluster = CaseCluster::jumpTable(Low: Clusters[First].Low, High: Clusters[Last].High, |
| 264 | JTCasesIndex: JTCases.size() - 1, Prob); |
| 265 | return true; |
| 266 | } |
| 267 | |
| 268 | void SwitchCG::SwitchLowering::findBitTestClusters(CaseClusterVector &Clusters, |
| 269 | const SwitchInst *SI) { |
| 270 | // Partition Clusters into as few subsets as possible, where each subset has a |
| 271 | // range that fits in a machine word and has <= 3 unique destinations. |
| 272 | |
| 273 | #ifndef NDEBUG |
| 274 | // Clusters must be sorted and contain Range or JumpTable clusters. |
| 275 | assert(!Clusters.empty()); |
| 276 | assert(Clusters[0].Kind == CC_Range || Clusters[0].Kind == CC_JumpTable); |
| 277 | for (const CaseCluster &C : Clusters) |
| 278 | assert(C.Kind == CC_Range || C.Kind == CC_JumpTable); |
| 279 | for (unsigned i = 1; i < Clusters.size(); ++i) |
| 280 | assert(Clusters[i-1].High->getValue().slt(Clusters[i].Low->getValue())); |
| 281 | #endif |
| 282 | |
| 283 | // The algorithm below is not suitable for -O0. |
| 284 | if (TM->getOptLevel() == CodeGenOptLevel::None) |
| 285 | return; |
| 286 | |
| 287 | // If target does not have legal shift left, do not emit bit tests at all. |
| 288 | EVT PTy = TLI->getPointerTy(DL: *DL); |
| 289 | if (!TLI->isOperationLegal(Op: ISD::SHL, VT: PTy)) |
| 290 | return; |
| 291 | |
| 292 | int BitWidth = PTy.getSizeInBits(); |
| 293 | const int64_t N = Clusters.size(); |
| 294 | |
| 295 | // MinPartitions[i] is the minimum nbr of partitions of Clusters[i..N-1]. |
| 296 | SmallVector<unsigned, 8> MinPartitions(N); |
| 297 | // LastElement[i] is the last element of the partition starting at i. |
| 298 | SmallVector<unsigned, 8> LastElement(N); |
| 299 | |
| 300 | // FIXME: This might not be the best algorithm for finding bit test clusters. |
| 301 | |
| 302 | // Base case: There is only one way to partition Clusters[N-1]. |
| 303 | MinPartitions[N - 1] = 1; |
| 304 | LastElement[N - 1] = N - 1; |
| 305 | |
| 306 | // Note: loop indexes are signed to avoid underflow. |
| 307 | for (int64_t i = N - 2; i >= 0; --i) { |
| 308 | // Find optimal partitioning of Clusters[i..N-1]. |
| 309 | // Baseline: Put Clusters[i] into a partition on its own. |
| 310 | MinPartitions[i] = MinPartitions[i + 1] + 1; |
| 311 | LastElement[i] = i; |
| 312 | |
| 313 | // Search for a solution that results in fewer partitions. |
| 314 | // Note: the search is limited by BitWidth, reducing time complexity. |
| 315 | for (int64_t j = std::min(a: N - 1, b: i + BitWidth - 1); j > i; --j) { |
| 316 | // Try building a partition from Clusters[i..j]. |
| 317 | |
| 318 | // Check the range. |
| 319 | if (!TLI->rangeFitsInWord(Low: Clusters[i].Low->getValue(), |
| 320 | High: Clusters[j].High->getValue(), DL: *DL)) |
| 321 | continue; |
| 322 | |
| 323 | // Check nbr of destinations and cluster types. |
| 324 | // FIXME: This works, but doesn't seem very efficient. |
| 325 | bool RangesOnly = true; |
| 326 | BitVector Dests(FuncInfo.MF->getNumBlockIDs()); |
| 327 | for (int64_t k = i; k <= j; k++) { |
| 328 | if (Clusters[k].Kind != CC_Range) { |
| 329 | RangesOnly = false; |
| 330 | break; |
| 331 | } |
| 332 | Dests.set(Clusters[k].MBB->getNumber()); |
| 333 | } |
| 334 | if (!RangesOnly || Dests.count() > 3) |
| 335 | break; |
| 336 | |
| 337 | // Check if it's a better partition. |
| 338 | unsigned NumPartitions = 1 + (j == N - 1 ? 0 : MinPartitions[j + 1]); |
| 339 | if (NumPartitions < MinPartitions[i]) { |
| 340 | // Found a better partition. |
| 341 | MinPartitions[i] = NumPartitions; |
| 342 | LastElement[i] = j; |
| 343 | } |
| 344 | } |
| 345 | } |
| 346 | |
| 347 | // Iterate over the partitions, replacing with bit-test clusters in-place. |
| 348 | unsigned DstIndex = 0; |
| 349 | for (unsigned First = 0, Last; First < N; First = Last + 1) { |
| 350 | Last = LastElement[First]; |
| 351 | assert(First <= Last); |
| 352 | assert(DstIndex <= First); |
| 353 | |
| 354 | CaseCluster BitTestCluster; |
| 355 | if (buildBitTests(Clusters, First, Last, SI, BTCluster&: BitTestCluster)) { |
| 356 | Clusters[DstIndex++] = BitTestCluster; |
| 357 | } else { |
| 358 | size_t NumClusters = Last - First + 1; |
| 359 | std::memmove(dest: &Clusters[DstIndex], src: &Clusters[First], |
| 360 | n: sizeof(Clusters[0]) * NumClusters); |
| 361 | DstIndex += NumClusters; |
| 362 | } |
| 363 | } |
| 364 | Clusters.resize(new_size: DstIndex); |
| 365 | } |
| 366 | |
| 367 | bool SwitchCG::SwitchLowering::buildBitTests(CaseClusterVector &Clusters, |
| 368 | unsigned First, unsigned Last, |
| 369 | const SwitchInst *SI, |
| 370 | CaseCluster &BTCluster) { |
| 371 | assert(First <= Last); |
| 372 | if (First == Last) |
| 373 | return false; |
| 374 | |
| 375 | BitVector Dests(FuncInfo.MF->getNumBlockIDs()); |
| 376 | unsigned NumCmps = 0; |
| 377 | for (int64_t I = First; I <= Last; ++I) { |
| 378 | assert(Clusters[I].Kind == CC_Range); |
| 379 | Dests.set(Clusters[I].MBB->getNumber()); |
| 380 | NumCmps += (Clusters[I].Low == Clusters[I].High) ? 1 : 2; |
| 381 | } |
| 382 | unsigned NumDests = Dests.count(); |
| 383 | |
| 384 | APInt Low = Clusters[First].Low->getValue(); |
| 385 | APInt High = Clusters[Last].High->getValue(); |
| 386 | assert(Low.slt(High)); |
| 387 | |
| 388 | if (!TLI->isSuitableForBitTests(NumDests, NumCmps, Low, High, DL: *DL)) |
| 389 | return false; |
| 390 | |
| 391 | APInt LowBound; |
| 392 | APInt CmpRange; |
| 393 | |
| 394 | const int BitWidth = TLI->getPointerTy(DL: *DL).getSizeInBits(); |
| 395 | assert(TLI->rangeFitsInWord(Low, High, *DL) && |
| 396 | "Case range must fit in bit mask!" ); |
| 397 | |
| 398 | // Check if the clusters cover a contiguous range such that no value in the |
| 399 | // range will jump to the default statement. |
| 400 | bool ContiguousRange = true; |
| 401 | for (int64_t I = First + 1; I <= Last; ++I) { |
| 402 | if (Clusters[I].Low->getValue() != Clusters[I - 1].High->getValue() + 1) { |
| 403 | ContiguousRange = false; |
| 404 | break; |
| 405 | } |
| 406 | } |
| 407 | |
| 408 | if (Low.isStrictlyPositive() && High.slt(RHS: BitWidth)) { |
| 409 | // Optimize the case where all the case values fit in a word without having |
| 410 | // to subtract minValue. In this case, we can optimize away the subtraction. |
| 411 | LowBound = APInt::getZero(numBits: Low.getBitWidth()); |
| 412 | CmpRange = High; |
| 413 | ContiguousRange = false; |
| 414 | } else { |
| 415 | LowBound = Low; |
| 416 | CmpRange = High - Low; |
| 417 | } |
| 418 | |
| 419 | CaseBitsVector CBV; |
| 420 | auto TotalProb = BranchProbability::getZero(); |
| 421 | for (unsigned i = First; i <= Last; ++i) { |
| 422 | // Find the CaseBits for this destination. |
| 423 | unsigned j; |
| 424 | for (j = 0; j < CBV.size(); ++j) |
| 425 | if (CBV[j].BB == Clusters[i].MBB) |
| 426 | break; |
| 427 | if (j == CBV.size()) |
| 428 | CBV.push_back( |
| 429 | x: CaseBits(0, Clusters[i].MBB, 0, BranchProbability::getZero())); |
| 430 | CaseBits *CB = &CBV[j]; |
| 431 | |
| 432 | // Update Mask, Bits and ExtraProb. |
| 433 | uint64_t Lo = (Clusters[i].Low->getValue() - LowBound).getZExtValue(); |
| 434 | uint64_t Hi = (Clusters[i].High->getValue() - LowBound).getZExtValue(); |
| 435 | assert(Hi >= Lo && Hi < 64 && "Invalid bit case!" ); |
| 436 | CB->Mask |= (-1ULL >> (63 - (Hi - Lo))) << Lo; |
| 437 | CB->Bits += Hi - Lo + 1; |
| 438 | CB->ExtraProb += Clusters[i].Prob; |
| 439 | TotalProb += Clusters[i].Prob; |
| 440 | } |
| 441 | |
| 442 | BitTestInfo BTI; |
| 443 | llvm::sort(C&: CBV, Comp: [](const CaseBits &a, const CaseBits &b) { |
| 444 | // Sort by probability first, number of bits second, bit mask third. |
| 445 | if (a.ExtraProb != b.ExtraProb) |
| 446 | return a.ExtraProb > b.ExtraProb; |
| 447 | if (a.Bits != b.Bits) |
| 448 | return a.Bits > b.Bits; |
| 449 | return a.Mask < b.Mask; |
| 450 | }); |
| 451 | |
| 452 | for (auto &CB : CBV) { |
| 453 | MachineBasicBlock *BitTestBB = |
| 454 | FuncInfo.MF->CreateMachineBasicBlock(BB: SI->getParent()); |
| 455 | BTI.push_back(Elt: BitTestCase(CB.Mask, BitTestBB, CB.BB, CB.ExtraProb)); |
| 456 | } |
| 457 | BitTestCases.emplace_back(args: std::move(LowBound), args: std::move(CmpRange), |
| 458 | args: SI->getCondition(), args: Register(), args: MVT::Other, args: false, |
| 459 | args&: ContiguousRange, args: nullptr, args: nullptr, args: std::move(BTI), |
| 460 | args&: TotalProb); |
| 461 | |
| 462 | BTCluster = CaseCluster::bitTests(Low: Clusters[First].Low, High: Clusters[Last].High, |
| 463 | BTCasesIndex: BitTestCases.size() - 1, Prob: TotalProb); |
| 464 | return true; |
| 465 | } |
| 466 | |
| 467 | void SwitchCG::sortAndRangeify(CaseClusterVector &Clusters) { |
| 468 | #ifndef NDEBUG |
| 469 | for (const CaseCluster &CC : Clusters) |
| 470 | assert(CC.Low == CC.High && "Input clusters must be single-case" ); |
| 471 | #endif |
| 472 | |
| 473 | llvm::sort(C&: Clusters, Comp: [](const CaseCluster &a, const CaseCluster &b) { |
| 474 | return a.Low->getValue().slt(RHS: b.Low->getValue()); |
| 475 | }); |
| 476 | |
| 477 | // Merge adjacent clusters with the same destination. |
| 478 | const unsigned N = Clusters.size(); |
| 479 | unsigned DstIndex = 0; |
| 480 | for (unsigned SrcIndex = 0; SrcIndex < N; ++SrcIndex) { |
| 481 | CaseCluster &CC = Clusters[SrcIndex]; |
| 482 | const ConstantInt *CaseVal = CC.Low; |
| 483 | MachineBasicBlock *Succ = CC.MBB; |
| 484 | |
| 485 | if (DstIndex != 0 && Clusters[DstIndex - 1].MBB == Succ && |
| 486 | (CaseVal->getValue() - Clusters[DstIndex - 1].High->getValue()) == 1) { |
| 487 | // If this case has the same successor and is a neighbour, merge it into |
| 488 | // the previous cluster. |
| 489 | Clusters[DstIndex - 1].High = CaseVal; |
| 490 | Clusters[DstIndex - 1].Prob += CC.Prob; |
| 491 | } else { |
| 492 | std::memmove(dest: &Clusters[DstIndex++], src: &Clusters[SrcIndex], |
| 493 | n: sizeof(Clusters[SrcIndex])); |
| 494 | } |
| 495 | } |
| 496 | Clusters.resize(new_size: DstIndex); |
| 497 | } |
| 498 | |
| 499 | unsigned SwitchCG::SwitchLowering::caseClusterRank(const CaseCluster &CC, |
| 500 | CaseClusterIt First, |
| 501 | CaseClusterIt Last) { |
| 502 | return std::count_if(first: First, last: Last + 1, pred: [&](const CaseCluster &X) { |
| 503 | if (X.Prob != CC.Prob) |
| 504 | return X.Prob > CC.Prob; |
| 505 | |
| 506 | // Ties are broken by comparing the case value. |
| 507 | return X.Low->getValue().slt(RHS: CC.Low->getValue()); |
| 508 | }); |
| 509 | } |
| 510 | |
| 511 | llvm::SwitchCG::SwitchLowering::SplitWorkItemInfo |
| 512 | SwitchCG::SwitchLowering::computeSplitWorkItemInfo( |
| 513 | const SwitchWorkListItem &W) { |
| 514 | CaseClusterIt LastLeft = W.FirstCluster; |
| 515 | CaseClusterIt FirstRight = W.LastCluster; |
| 516 | auto LeftProb = LastLeft->Prob + W.DefaultProb / 2; |
| 517 | auto RightProb = FirstRight->Prob + W.DefaultProb / 2; |
| 518 | |
| 519 | // Move LastLeft and FirstRight towards each other from opposite directions to |
| 520 | // find a partitioning of the clusters which balances the probability on both |
| 521 | // sides. If LeftProb and RightProb are equal, alternate which side is |
| 522 | // taken to ensure 0-probability nodes are distributed evenly. |
| 523 | unsigned I = 0; |
| 524 | while (LastLeft + 1 < FirstRight) { |
| 525 | if (LeftProb < RightProb || (LeftProb == RightProb && (I & 1))) |
| 526 | LeftProb += (++LastLeft)->Prob; |
| 527 | else |
| 528 | RightProb += (--FirstRight)->Prob; |
| 529 | I++; |
| 530 | } |
| 531 | |
| 532 | while (true) { |
| 533 | // Our binary search tree differs from a typical BST in that ours can have |
| 534 | // up to three values in each leaf. The pivot selection above doesn't take |
| 535 | // that into account, which means the tree might require more nodes and be |
| 536 | // less efficient. We compensate for this here. |
| 537 | |
| 538 | unsigned NumLeft = LastLeft - W.FirstCluster + 1; |
| 539 | unsigned NumRight = W.LastCluster - FirstRight + 1; |
| 540 | |
| 541 | if (std::min(a: NumLeft, b: NumRight) < 3 && std::max(a: NumLeft, b: NumRight) > 3) { |
| 542 | // If one side has less than 3 clusters, and the other has more than 3, |
| 543 | // consider taking a cluster from the other side. |
| 544 | |
| 545 | if (NumLeft < NumRight) { |
| 546 | // Consider moving the first cluster on the right to the left side. |
| 547 | CaseCluster &CC = *FirstRight; |
| 548 | unsigned RightSideRank = caseClusterRank(CC, First: FirstRight, Last: W.LastCluster); |
| 549 | unsigned LeftSideRank = caseClusterRank(CC, First: W.FirstCluster, Last: LastLeft); |
| 550 | if (LeftSideRank <= RightSideRank) { |
| 551 | // Moving the cluster to the left does not demote it. |
| 552 | ++LastLeft; |
| 553 | ++FirstRight; |
| 554 | continue; |
| 555 | } |
| 556 | } else { |
| 557 | assert(NumRight < NumLeft); |
| 558 | // Consider moving the last element on the left to the right side. |
| 559 | CaseCluster &CC = *LastLeft; |
| 560 | unsigned LeftSideRank = caseClusterRank(CC, First: W.FirstCluster, Last: LastLeft); |
| 561 | unsigned RightSideRank = caseClusterRank(CC, First: FirstRight, Last: W.LastCluster); |
| 562 | if (RightSideRank <= LeftSideRank) { |
| 563 | // Moving the cluster to the right does not demot it. |
| 564 | --LastLeft; |
| 565 | --FirstRight; |
| 566 | continue; |
| 567 | } |
| 568 | } |
| 569 | } |
| 570 | break; |
| 571 | } |
| 572 | |
| 573 | assert(LastLeft + 1 == FirstRight); |
| 574 | assert(LastLeft >= W.FirstCluster); |
| 575 | assert(FirstRight <= W.LastCluster); |
| 576 | |
| 577 | return SplitWorkItemInfo{.LastLeft: LastLeft, .FirstRight: FirstRight, .LeftProb: LeftProb, .RightProb: RightProb}; |
| 578 | } |
| 579 | |