llvm/lib/CodeGen/IndirectBrExpandPass.cpp - llvm-project - Git at Google

 //===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 ///
 /// Implements an expansion pass to turn `indirectbr` instructions in the IR
 /// into `switch` instructions. This works by enumerating the basic blocks in
 /// a dense range of integers, replacing each `blockaddr` constant with the
 /// corresponding integer constant, and then building a switch that maps from
 /// the integers to the actual blocks. All of the indirectbr instructions in the
 /// function are redirected to this common switch.
 ///
 /// While this is generically useful if a target is unable to codegen
 /// `indirectbr` natively, it is primarily useful when there is some desire to
 /// get the builtin non-jump-table lowering of a switch even when the input
 /// source contained an explicit indirect branch construct.
 ///
 /// Note that it doesn't make any sense to enable this pass unless a target also
 /// disables jump-table lowering of switches. Doing that is likely to pessimize
 /// the code.
 ///
 //===----------------------------------------------------------------------===//

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/LazyBlockFrequencyInfo.h"
 #include "llvm/CodeGen/IndirectBrExpand.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/ProfDataUtils.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ScaledNumber.h"
 #include "llvm/Target/TargetMachine.h"
 #include <optional>

 using namespace llvm;

 #define DEBUG_TYPE "indirectbr-expand"

 namespace llvm {
 extern cl::opt<bool> ProfcheckDisableMetadataFixes;
 } // namespace llvm

 namespace {

 class IndirectBrExpandLegacyPass : public FunctionPass {
 public:
   static char ID; // Pass identification, replacement for typeid

   IndirectBrExpandLegacyPass() : FunctionPass(ID) {}

   void getAnalysisUsage(AnalysisUsage &AU) const override {
     LazyBlockFrequencyInfoPass::getLazyBFIAnalysisUsage(AU);
     AU.addPreserved<DominatorTreeWrapperPass>();
   }

   bool runOnFunction(Function &F) override;
 };

 } // end anonymous namespace

 static bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU,
                     function_ref<BlockFrequencyInfo *()> GetBFI);

 PreservedAnalyses IndirectBrExpandPass::run(Function &F,
                                             FunctionAnalysisManager &FAM) {
   auto *STI = TM->getSubtargetImpl(F);
   if (!STI->enableIndirectBrExpand())
     return PreservedAnalyses::all();

   auto *TLI = STI->getTargetLowering();
   auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);

   bool Changed = runImpl(F, TLI, DT ? &DTU : nullptr, [&]() {
     return &FAM.getResult<BlockFrequencyAnalysis>(F);
   });
   if (!Changed)
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserve<DominatorTreeAnalysis>();
   return PA;
 }

 char IndirectBrExpandLegacyPass::ID = 0;

 INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
                       "Expand indirectbr instructions", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
                     "Expand indirectbr instructions", false, false)

 FunctionPass *llvm::createIndirectBrExpandPass() {
   return new IndirectBrExpandLegacyPass();
 }

 bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU,
              function_ref<BlockFrequencyInfo *()> GetBFI) {
   auto &DL = F.getDataLayout();

   SmallVector<IndirectBrInst *, 1> IndirectBrs;
   SmallVector<uint64_t, 1> IndirectBrsBlockFrequencies;
   SmallVector<uint64_t, 1> IndirectBrsBranchWeightSums;
   bool SkipProfileUpdates = false;
   BlockFrequencyInfo *BFI = nullptr;

   struct IndirectBrSuccessor {
     // The index into the IndirectBrs, IndirectBrsBlockFrequencies, and
     // IndirectBrsBranchWeightSums vectors.
     size_t IndirectBrIndex = 0;
     uint64_t SuccessorBranchWeight = 0;
   };

   // Set of all potential successors for indirectbr instructions.
   DenseMap<const BasicBlock *, SmallVector<IndirectBrSuccessor>>
       IndirectBrSuccToIndirectBr;

   // Build a list of indirectbrs that we want to rewrite.
   for (BasicBlock &BB : F)
     if (auto *IBr = dyn_cast<IndirectBrInst>(BB.getTerminator())) {
       // Handle the degenerate case of no successors by replacing the indirectbr
       // with unreachable as there is no successor available.
       if (IBr->getNumSuccessors() == 0) {
         (void)new UnreachableInst(F.getContext(), IBr->getIterator());
         IBr->eraseFromParent();
         continue;
       }

       IndirectBrs.push_back(IBr);
       const size_t CurrentIndirectBrIndex = IndirectBrs.size() - 1;
       for (const BasicBlock *SuccessorBB : IBr->successors())
         IndirectBrSuccToIndirectBr.insert({SuccessorBB, {}});

       if (SkipProfileUpdates)
         continue;
       if (!BFI)
         BFI = GetBFI();
       std::optional<uint64_t> BlockFrequency = BFI->getBlockProfileCount(&BB);
       if (!BlockFrequency.has_value()) {
         SkipProfileUpdates = true;
         continue;
       }
       IndirectBrsBlockFrequencies.push_back(*BlockFrequency);
       SmallVector<uint32_t> IndirectBrBranchWeights;
       bool HasBranchWeights =
           extractBranchWeights(*IBr, IndirectBrBranchWeights);
       if (!HasBranchWeights) {
         SkipProfileUpdates = true;
         continue;
       }
       for (const auto [SuccessorBB, SuccessorBranchWeight] :
            zip_equal(IBr->successors(), IndirectBrBranchWeights))
         IndirectBrSuccToIndirectBr[SuccessorBB].push_back(
             {CurrentIndirectBrIndex, SuccessorBranchWeight});
       IndirectBrsBranchWeightSums.push_back(sum_of(IndirectBrBranchWeights));
       assert(IndirectBrsBranchWeightSums.size() == IndirectBrs.size() &&
              "expected an identical number of blocks in both vectors");
     }

   if (IndirectBrs.empty())
     return false;

   // If we need to replace any indirectbrs we need to establish integer
   // constants that will correspond to each of the basic blocks in the function
   // whose address escapes. We do that here and rewrite all the blockaddress
   // constants to just be those integer constants cast to a pointer type.
   SmallVector<BasicBlock *, 4> BBs;
   SmallVector<ScaledNumber<uint64_t>, 4> BBWeights;

   for (BasicBlock &BB : F) {
     // Skip blocks that aren't successors to an indirectbr we're going to
     // rewrite.
     auto IndirectBrSuccToIndirectBrIt = IndirectBrSuccToIndirectBr.find(&BB);
     if (IndirectBrSuccToIndirectBrIt == IndirectBrSuccToIndirectBr.end())
       continue;

     auto *BA = BlockAddress::lookup(&BB);

     // Skip if the constant was formed but ended up not being used (due to DCE
     // or whatever).
     if (!BA || !BA->isConstantUsed())
       continue;

     // Compute the index we want to use for this basic block. We can't use zero
     // because null can be compared with block addresses.
     int BBIndex = BBs.size() + 1;
     BBs.push_back(&BB);

     auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
     ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);

     // Now rewrite the blockaddress to an integer constant based on the index.
     // FIXME: This part doesn't properly recognize other uses of blockaddress
     // expressions, for instance, where they are used to pass labels to
     // asm-goto. This part of the pass needs a rework.
     BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));

     if (SkipProfileUpdates)
       continue;
     ScaledNumber<uint64_t> BranchWeightSumsProduct(1, 0);
     for (uint64_t BranchWeightSum : IndirectBrsBranchWeightSums)
       BranchWeightSumsProduct *= ScaledNumber<uint64_t>(BranchWeightSum, 0);
     ScaledNumber<uint64_t> BlockWeight(0, 0);
     for (const auto &[IndirectBrIndex, BlockBranchProbability] :
          IndirectBrSuccToIndirectBrIt->second) {
       // If the branch weight sum is zero, skip adding the block weight or
       // otherwise we end up dividing by zero.
       const uint64_t CurrentBranchWeightSum =
           IndirectBrsBranchWeightSums[IndirectBrIndex];
       if (CurrentBranchWeightSum == 0)
         continue;
       BlockWeight += ScaledNumber<uint64_t>(
                          IndirectBrsBlockFrequencies[IndirectBrIndex], 0) *
                      ScaledNumber<uint64_t>(BlockBranchProbability, 0) *
                      (BranchWeightSumsProduct /
                       ScaledNumber<uint64_t>(CurrentBranchWeightSum, 0));
     }
     BBWeights.push_back(BlockWeight);
   }

   if (BBs.empty()) {
     // There are no blocks whose address is taken, so any indirectbr instruction
     // cannot get a valid input and we can replace all of them with unreachable.
     SmallVector<DominatorTree::UpdateType, 8> Updates;
     if (DTU)
       Updates.reserve(IndirectBrSuccToIndirectBr.size());
     for (auto *IBr : IndirectBrs) {
       if (DTU) {
         for (BasicBlock *SuccBB : IBr->successors())
           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
       }
       (void)new UnreachableInst(F.getContext(), IBr->getIterator());
       IBr->eraseFromParent();
     }
     if (DTU) {
       assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
              "Got unexpected update count.");
       DTU->applyUpdates(Updates);
     }
     return true;
   }

   BasicBlock *SwitchBB;
   Value *SwitchValue;

   // Compute a common integer type across all the indirectbr instructions.
   IntegerType *CommonITy = nullptr;
   for (auto *IBr : IndirectBrs) {
     auto *ITy =
         cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));
     if (!CommonITy || ITy->getBitWidth() > CommonITy->getBitWidth())
       CommonITy = ITy;
   }

   auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {
     return CastInst::CreatePointerCast(IBr->getAddress(), CommonITy,
                                        Twine(IBr->getAddress()->getName()) +
                                            ".switch_cast",
                                        IBr->getIterator());
   };

   SmallVector<DominatorTree::UpdateType, 8> Updates;

   if (IndirectBrs.size() == 1) {
     // If we only have one indirectbr, we can just directly replace it within
     // its block.
     IndirectBrInst *IBr = IndirectBrs[0];
     SwitchBB = IBr->getParent();
     SwitchValue = GetSwitchValue(IBr);
     if (DTU) {
       Updates.reserve(IndirectBrSuccToIndirectBr.size());
       for (BasicBlock *SuccBB : IBr->successors())
         Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
       assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
              "Got unexpected update count.");
     }
     IBr->eraseFromParent();
   } else {
     // Otherwise we need to create a new block to hold the switch across BBs,
     // jump to that block instead of each indirectbr, and phi together the
     // values for the switch.
     SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
     auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
                                      "switch_value_phi", SwitchBB);
     SwitchValue = SwitchPN;

     // Now replace the indirectbr instructions with direct branches to the
     // switch block and fill out the PHI operands.
     if (DTU)
       Updates.reserve(IndirectBrs.size() +
                       2 * IndirectBrSuccToIndirectBr.size());
     for (auto *IBr : IndirectBrs) {
       SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
       UncondBrInst::Create(SwitchBB, IBr->getIterator());
       if (DTU) {
         Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
         for (BasicBlock *SuccBB : IBr->successors())
           Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
       }
       IBr->eraseFromParent();
     }
   }

   // Now build the switch in the block. The block will have no terminator
   // already.
   auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);

   // Add a case for each block.
   for (int i : llvm::seq<int>(1, BBs.size()))
     SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);

   if (DTU) {
     // If there were multiple indirectbr's, they may have common successors,
     // but in the dominator tree, we only track unique edges.
     SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
     Updates.reserve(Updates.size() + BBs.size());
     for (BasicBlock *BB : BBs) {
       if (UniqueSuccessors.insert(BB).second)
         Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
     }
     DTU->applyUpdates(Updates);
   }

   if (SkipProfileUpdates || ProfcheckDisableMetadataFixes) {
     setExplicitlyUnknownBranchWeightsIfProfiled(*SI, DEBUG_TYPE);
     return true;
   }

   // We need to convert the ScaledNumber weights (which might not be
   // representable in 64 bits) back to normal 64 bit integers so we can apply
   // them as metadata. They might not have the same scale though, so we find the
   // max scale and then scale down any weights that have a scale less than the
   // max scale. This ensures that all the weights have the same scale.
   int16_t MaxScale = 0;
   for (const ScaledNumber<uint64_t> &BBWeight : BBWeights)
     MaxScale = std::max(MaxScale, BBWeight.getScale());
   SmallVector<uint64_t, 4> ExtractedBBWeights;
   ExtractedBBWeights.reserve(BBWeights.size());
   for (ScaledNumber<uint64_t> &BBWeight : BBWeights) {
     int16_t Shift = MaxScale - BBWeight.getScale();
     assert(Shift >= 0 && "expected non-negative shift");
     ExtractedBBWeights.push_back(BBWeight.getDigits() >> Shift);
   }
   setFittedBranchWeights(*SI, ExtractedBBWeights, false);

   return true;
 }

 bool IndirectBrExpandLegacyPass::runOnFunction(Function &F) {
   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
   if (!TPC)
     return false;

   auto &TM = TPC->getTM<TargetMachine>();
   auto &STI = *TM.getSubtargetImpl(F);
   if (!STI.enableIndirectBrExpand())
     return false;
   auto *TLI = STI.getTargetLowering();

   std::optional<DomTreeUpdater> DTU;
   if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
     DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);

   return runImpl(F, TLI, DTU ? &*DTU : nullptr, [&]() {
     return &getAnalysis<LazyBlockFrequencyInfoPass>().getBFI();
   });
 }
	//===- IndirectBrExpandPass.cpp - Expand indirectbr to switch -------------===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	/// \file
	///
	/// Implements an expansion pass to turn `indirectbr` instructions in the IR
	/// into `switch` instructions. This works by enumerating the basic blocks in
	/// a dense range of integers, replacing each `blockaddr` constant with the
	/// corresponding integer constant, and then building a switch that maps from
	/// the integers to the actual blocks. All of the indirectbr instructions in the
	/// function are redirected to this common switch.
	///
	/// While this is generically useful if a target is unable to codegen
	/// `indirectbr` natively, it is primarily useful when there is some desire to
	/// get the builtin non-jump-table lowering of a switch even when the input
	/// source contained an explicit indirect branch construct.
	///
	/// Note that it doesn't make any sense to enable this pass unless a target also
	/// disables jump-table lowering of switches. Doing that is likely to pessimize
	/// the code.
	///
	//===----------------------------------------------------------------------===//

	#include "llvm/ADT/STLExtras.h"
	#include "llvm/ADT/Sequence.h"
	#include "llvm/ADT/SmallVector.h"
	#include "llvm/Analysis/BlockFrequencyInfo.h"
	#include "llvm/Analysis/DomTreeUpdater.h"
	#include "llvm/Analysis/LazyBlockFrequencyInfo.h"
	#include "llvm/CodeGen/IndirectBrExpand.h"
	#include "llvm/CodeGen/TargetPassConfig.h"
	#include "llvm/CodeGen/TargetSubtargetInfo.h"
	#include "llvm/IR/BasicBlock.h"
	#include "llvm/IR/Constants.h"
	#include "llvm/IR/Dominators.h"
	#include "llvm/IR/Function.h"
	#include "llvm/IR/Instructions.h"
	#include "llvm/IR/ProfDataUtils.h"
	#include "llvm/InitializePasses.h"
	#include "llvm/Pass.h"
	#include "llvm/Support/ErrorHandling.h"
	#include "llvm/Support/ScaledNumber.h"
	#include "llvm/Target/TargetMachine.h"
	#include <optional>

	using namespace llvm;

	#define DEBUG_TYPE "indirectbr-expand"

	namespace llvm {
	extern cl::opt<bool> ProfcheckDisableMetadataFixes;
	} // namespace llvm

	namespace {

	class IndirectBrExpandLegacyPass : public FunctionPass {
	public:
	static char ID; // Pass identification, replacement for typeid

	IndirectBrExpandLegacyPass() : FunctionPass(ID) {}

	void getAnalysisUsage(AnalysisUsage &AU) const override {
	LazyBlockFrequencyInfoPass::getLazyBFIAnalysisUsage(AU);
	AU.addPreserved<DominatorTreeWrapperPass>();
	}

	bool runOnFunction(Function &F) override;
	};

	} // end anonymous namespace

	static bool runImpl(Function &F, const TargetLowering TLI, DomTreeUpdater DTU,
	function_ref<BlockFrequencyInfo *()> GetBFI);

	PreservedAnalyses IndirectBrExpandPass::run(Function &F,
	FunctionAnalysisManager &FAM) {
	auto *STI = TM->getSubtargetImpl(F);
	if (!STI->enableIndirectBrExpand())
	return PreservedAnalyses::all();

	auto *TLI = STI->getTargetLowering();
	auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
	DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);

	bool Changed = runImpl(F, TLI, DT ? &DTU : nullptr, [&]() {
	return &FAM.getResult<BlockFrequencyAnalysis>(F);
	});
	if (!Changed)
	return PreservedAnalyses::all();
	PreservedAnalyses PA;
	PA.preserve<DominatorTreeAnalysis>();
	return PA;
	}

	char IndirectBrExpandLegacyPass::ID = 0;

	INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
	"Expand indirectbr instructions", false, false)
	INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
	INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
	"Expand indirectbr instructions", false, false)

	FunctionPass *llvm::createIndirectBrExpandPass() {
	return new IndirectBrExpandLegacyPass();
	}

	bool runImpl(Function &F, const TargetLowering TLI, DomTreeUpdater DTU,
	function_ref<BlockFrequencyInfo *()> GetBFI) {
	auto &DL = F.getDataLayout();

	SmallVector<IndirectBrInst *, 1> IndirectBrs;
	SmallVector<uint64_t, 1> IndirectBrsBlockFrequencies;
	SmallVector<uint64_t, 1> IndirectBrsBranchWeightSums;
	bool SkipProfileUpdates = false;
	BlockFrequencyInfo *BFI = nullptr;

	struct IndirectBrSuccessor {
	// The index into the IndirectBrs, IndirectBrsBlockFrequencies, and
	// IndirectBrsBranchWeightSums vectors.
	size_t IndirectBrIndex = 0;
	uint64_t SuccessorBranchWeight = 0;
	};

	// Set of all potential successors for indirectbr instructions.
	DenseMap<const BasicBlock *, SmallVector<IndirectBrSuccessor>>
	IndirectBrSuccToIndirectBr;

	// Build a list of indirectbrs that we want to rewrite.
	for (BasicBlock &BB : F)
	if (auto *IBr = dyn_cast<IndirectBrInst>(BB.getTerminator())) {
	// Handle the degenerate case of no successors by replacing the indirectbr
	// with unreachable as there is no successor available.
	if (IBr->getNumSuccessors() == 0) {
	(void)new UnreachableInst(F.getContext(), IBr->getIterator());
	IBr->eraseFromParent();
	continue;
	}

	IndirectBrs.push_back(IBr);
	const size_t CurrentIndirectBrIndex = IndirectBrs.size() - 1;
	for (const BasicBlock *SuccessorBB : IBr->successors())
	IndirectBrSuccToIndirectBr.insert({SuccessorBB, {}});

	if (SkipProfileUpdates)
	continue;
	if (!BFI)
	BFI = GetBFI();
	std::optional<uint64_t> BlockFrequency = BFI->getBlockProfileCount(&BB);
	if (!BlockFrequency.has_value()) {
	SkipProfileUpdates = true;
	continue;
	}
	IndirectBrsBlockFrequencies.push_back(*BlockFrequency);
	SmallVector<uint32_t> IndirectBrBranchWeights;
	bool HasBranchWeights =
	extractBranchWeights(*IBr, IndirectBrBranchWeights);
	if (!HasBranchWeights) {
	SkipProfileUpdates = true;
	continue;
	}
	for (const auto [SuccessorBB, SuccessorBranchWeight] :
	zip_equal(IBr->successors(), IndirectBrBranchWeights))
	IndirectBrSuccToIndirectBr[SuccessorBB].push_back(
	{CurrentIndirectBrIndex, SuccessorBranchWeight});
	IndirectBrsBranchWeightSums.push_back(sum_of(IndirectBrBranchWeights));
	assert(IndirectBrsBranchWeightSums.size() == IndirectBrs.size() &&
	"expected an identical number of blocks in both vectors");
	}

	if (IndirectBrs.empty())
	return false;

	// If we need to replace any indirectbrs we need to establish integer
	// constants that will correspond to each of the basic blocks in the function
	// whose address escapes. We do that here and rewrite all the blockaddress
	// constants to just be those integer constants cast to a pointer type.
	SmallVector<BasicBlock *, 4> BBs;
	SmallVector<ScaledNumber<uint64_t>, 4> BBWeights;

	for (BasicBlock &BB : F) {
	// Skip blocks that aren't successors to an indirectbr we're going to
	// rewrite.
	auto IndirectBrSuccToIndirectBrIt = IndirectBrSuccToIndirectBr.find(&BB);
	if (IndirectBrSuccToIndirectBrIt == IndirectBrSuccToIndirectBr.end())
	continue;

	auto *BA = BlockAddress::lookup(&BB);

	// Skip if the constant was formed but ended up not being used (due to DCE
	// or whatever).
	if (!BA \|\| !BA->isConstantUsed())
	continue;

	// Compute the index we want to use for this basic block. We can't use zero
	// because null can be compared with block addresses.
	int BBIndex = BBs.size() + 1;
	BBs.push_back(&BB);

	auto *ITy = cast<IntegerType>(DL.getIntPtrType(BA->getType()));
	ConstantInt *BBIndexC = ConstantInt::get(ITy, BBIndex);

	// Now rewrite the blockaddress to an integer constant based on the index.
	// FIXME: This part doesn't properly recognize other uses of blockaddress
	// expressions, for instance, where they are used to pass labels to
	// asm-goto. This part of the pass needs a rework.
	BA->replaceAllUsesWith(ConstantExpr::getIntToPtr(BBIndexC, BA->getType()));

	if (SkipProfileUpdates)
	continue;
	ScaledNumber<uint64_t> BranchWeightSumsProduct(1, 0);
	for (uint64_t BranchWeightSum : IndirectBrsBranchWeightSums)
	BranchWeightSumsProduct *= ScaledNumber<uint64_t>(BranchWeightSum, 0);
	ScaledNumber<uint64_t> BlockWeight(0, 0);
	for (const auto &[IndirectBrIndex, BlockBranchProbability] :
	IndirectBrSuccToIndirectBrIt->second) {
	// If the branch weight sum is zero, skip adding the block weight or
	// otherwise we end up dividing by zero.
	const uint64_t CurrentBranchWeightSum =
	IndirectBrsBranchWeightSums[IndirectBrIndex];
	if (CurrentBranchWeightSum == 0)
	continue;
	BlockWeight += ScaledNumber<uint64_t>(
	IndirectBrsBlockFrequencies[IndirectBrIndex], 0) *
	ScaledNumber<uint64_t>(BlockBranchProbability, 0) *
	(BranchWeightSumsProduct /
	ScaledNumber<uint64_t>(CurrentBranchWeightSum, 0));
	}
	BBWeights.push_back(BlockWeight);
	}

	if (BBs.empty()) {
	// There are no blocks whose address is taken, so any indirectbr instruction
	// cannot get a valid input and we can replace all of them with unreachable.
	SmallVector<DominatorTree::UpdateType, 8> Updates;
	if (DTU)
	Updates.reserve(IndirectBrSuccToIndirectBr.size());
	for (auto *IBr : IndirectBrs) {
	if (DTU) {
	for (BasicBlock *SuccBB : IBr->successors())
	Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
	}
	(void)new UnreachableInst(F.getContext(), IBr->getIterator());
	IBr->eraseFromParent();
	}
	if (DTU) {
	assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
	"Got unexpected update count.");
	DTU->applyUpdates(Updates);
	}
	return true;
	}

	BasicBlock *SwitchBB;
	Value *SwitchValue;

	// Compute a common integer type across all the indirectbr instructions.
	IntegerType *CommonITy = nullptr;
	for (auto *IBr : IndirectBrs) {
	auto *ITy =
	cast<IntegerType>(DL.getIntPtrType(IBr->getAddress()->getType()));
	if (!CommonITy \|\| ITy->getBitWidth() > CommonITy->getBitWidth())
	CommonITy = ITy;
	}

	auto GetSwitchValue = [CommonITy](IndirectBrInst *IBr) {
	return CastInst::CreatePointerCast(IBr->getAddress(), CommonITy,
	Twine(IBr->getAddress()->getName()) +
	".switch_cast",
	IBr->getIterator());
	};

	SmallVector<DominatorTree::UpdateType, 8> Updates;

	if (IndirectBrs.size() == 1) {
	// If we only have one indirectbr, we can just directly replace it within
	// its block.
	IndirectBrInst *IBr = IndirectBrs[0];
	SwitchBB = IBr->getParent();
	SwitchValue = GetSwitchValue(IBr);
	if (DTU) {
	Updates.reserve(IndirectBrSuccToIndirectBr.size());
	for (BasicBlock *SuccBB : IBr->successors())
	Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
	assert(Updates.size() == IndirectBrSuccToIndirectBr.size() &&
	"Got unexpected update count.");
	}
	IBr->eraseFromParent();
	} else {
	// Otherwise we need to create a new block to hold the switch across BBs,
	// jump to that block instead of each indirectbr, and phi together the
	// values for the switch.
	SwitchBB = BasicBlock::Create(F.getContext(), "switch_bb", &F);
	auto *SwitchPN = PHINode::Create(CommonITy, IndirectBrs.size(),
	"switch_value_phi", SwitchBB);
	SwitchValue = SwitchPN;

	// Now replace the indirectbr instructions with direct branches to the
	// switch block and fill out the PHI operands.
	if (DTU)
	Updates.reserve(IndirectBrs.size() +
	2 * IndirectBrSuccToIndirectBr.size());
	for (auto *IBr : IndirectBrs) {
	SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
	UncondBrInst::Create(SwitchBB, IBr->getIterator());
	if (DTU) {
	Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
	for (BasicBlock *SuccBB : IBr->successors())
	Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
	}
	IBr->eraseFromParent();
	}
	}

	// Now build the switch in the block. The block will have no terminator
	// already.
	auto *SI = SwitchInst::Create(SwitchValue, BBs[0], BBs.size(), SwitchBB);

	// Add a case for each block.
	for (int i : llvm::seq<int>(1, BBs.size()))
	SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);

	if (DTU) {
	// If there were multiple indirectbr's, they may have common successors,
	// but in the dominator tree, we only track unique edges.
	SmallPtrSet<BasicBlock *, 8> UniqueSuccessors;
	Updates.reserve(Updates.size() + BBs.size());
	for (BasicBlock *BB : BBs) {
	if (UniqueSuccessors.insert(BB).second)
	Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
	}
	DTU->applyUpdates(Updates);
	}

	if (SkipProfileUpdates \|\| ProfcheckDisableMetadataFixes) {
	setExplicitlyUnknownBranchWeightsIfProfiled(*SI, DEBUG_TYPE);
	return true;
	}

	// We need to convert the ScaledNumber weights (which might not be
	// representable in 64 bits) back to normal 64 bit integers so we can apply
	// them as metadata. They might not have the same scale though, so we find the
	// max scale and then scale down any weights that have a scale less than the
	// max scale. This ensures that all the weights have the same scale.
	int16_t MaxScale = 0;
	for (const ScaledNumber<uint64_t> &BBWeight : BBWeights)
	MaxScale = std::max(MaxScale, BBWeight.getScale());
	SmallVector<uint64_t, 4> ExtractedBBWeights;
	ExtractedBBWeights.reserve(BBWeights.size());
	for (ScaledNumber<uint64_t> &BBWeight : BBWeights) {
	int16_t Shift = MaxScale - BBWeight.getScale();
	assert(Shift >= 0 && "expected non-negative shift");
	ExtractedBBWeights.push_back(BBWeight.getDigits() >> Shift);
	}
	setFittedBranchWeights(*SI, ExtractedBBWeights, false);

	return true;
	}

	bool IndirectBrExpandLegacyPass::runOnFunction(Function &F) {
	auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
	if (!TPC)
	return false;

	auto &TM = TPC->getTM<TargetMachine>();
	auto &STI = *TM.getSubtargetImpl(F);
	if (!STI.enableIndirectBrExpand())
	return false;
	auto *TLI = STI.getTargetLowering();

	std::optional<DomTreeUpdater> DTU;
	if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
	DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);

	return runImpl(F, TLI, DTU ? &*DTU : nullptr, [&]() {
	return &getAnalysis<LazyBlockFrequencyInfoPass>().getBFI();
	});
	}