Add a pass to convert jump tables to switches

Author: alexander-shaposhnikov

llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h

@@ -0,0 +1,24 @@
+//===- JumpTableToSwitch.h - ------------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H
+#define LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class Function;
+
+struct JumpTableToSwitchPass : PassInfoMixin<JumpTableToSwitchPass> {
+  /// Run the pass over the function.
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCALAR_JUMP_TABLE_TO_SWITCH_H

llvm/lib/Passes/PassBuilder.cpp

@@ -192,6 +192,7 @@
 #include "llvm/Transforms/Scalar/InferAddressSpaces.h"
 #include "llvm/Transforms/Scalar/InferAlignment.h"
 #include "llvm/Transforms/Scalar/InstSimplifyPass.h"
+#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Scalar/LICM.h"
 #include "llvm/Transforms/Scalar/LoopAccessAnalysisPrinter.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -91,6 +91,7 @@
 #include "llvm/Transforms/Scalar/IndVarSimplify.h"
 #include "llvm/Transforms/Scalar/InferAlignment.h"
 #include "llvm/Transforms/Scalar/InstSimplifyPass.h"
+#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
 #include "llvm/Transforms/Scalar/LICM.h"
 #include "llvm/Transforms/Scalar/LoopDeletion.h"
@@ -557,6 +558,7 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
 
   // Optimize based on known information about branches, and cleanup afterward.
   FPM.addPass(JumpThreadingPass());
+  FPM.addPass(JumpTableToSwitchPass());
   FPM.addPass(CorrelatedValuePropagationPass());
 
   FPM.addPass(
@@ -695,6 +697,7 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
     FPM.addPass(DFAJumpThreadingPass());
 
   FPM.addPass(JumpThreadingPass());
+
   FPM.addPass(CorrelatedValuePropagationPass());
 
   // Finally, do an expensive DCE pass to catch all the dead code exposed by
@@ -1926,6 +1929,7 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
 
   invokePeepholeEPCallbacks(MainFPM, Level);
   MainFPM.addPass(JumpThreadingPass());
+
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(MainFPM),
                                                 PTO.EagerlyInvalidateAnalyses));
 

llvm/lib/Passes/PassRegistry.def

@@ -330,6 +330,7 @@ FUNCTION_PASS("interleaved-load-combine", InterleavedLoadCombinePass(TM))
 FUNCTION_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 FUNCTION_PASS("irce", IRCEPass())
 FUNCTION_PASS("jump-threading", JumpThreadingPass())
+FUNCTION_PASS("jump-table-to-switch", JumpTableToSwitchPass());
 FUNCTION_PASS("kcfi", KCFIPass())
 FUNCTION_PASS("lcssa", LCSSAPass())
 FUNCTION_PASS("libcalls-shrinkwrap", LibCallsShrinkWrapPass())

llvm/lib/Transforms/Scalar/CMakeLists.txt

@@ -25,6 +25,7 @@ add_llvm_component_library(LLVMScalarOpts
   InferAlignment.cpp
   InstSimplifyPass.cpp
   JumpThreading.cpp
+  JumpTableToSwitch.cpp
   LICM.cpp
   LoopAccessAnalysisPrinter.cpp
   LoopBoundSplit.cpp

llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp

@@ -0,0 +1,205 @@
+//===- JumpTableToSwitch.cpp ----------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/Local.h"
+
+using namespace llvm;
+using namespace PatternMatch;
+
+static cl::opt<unsigned>
+    JumpTableSizeThreshold("jump-table-to-switch-size-threshold", cl::Hidden,
+                           cl::desc("Only split jump tables with size less or "
+                                    "equal than JumpTableSizeThreshold."),
+                           cl::init(10));
+
+#define DEBUG_TYPE "jump-table-to-switch"
+
+static Constant *getElementWithGivenTypeAtOffset(Constant *C, const Type *Ty,
+                                                 uint64_t Offset, Module &M) {
+  if (Offset == 0 && C->getType() == Ty)
+    return C;
+  if (auto *CS = dyn_cast<ConstantStruct>(C)) {
+    const DataLayout &DL = M.getDataLayout();
+    const StructLayout *SL = DL.getStructLayout(CS->getType());
+    if (Offset >= SL->getSizeInBytes())
+      return nullptr;
+    const unsigned Op = SL->getElementContainingOffset(Offset);
+    const uint64_t AdjustedOffset = Offset - SL->getElementOffset(Op);
+    Constant *Element = cast<Constant>(CS->getOperand(Op));
+    return getElementWithGivenTypeAtOffset(Element, Ty, AdjustedOffset, M);
+  }
+  // TODO: add support for arrays.
+  return nullptr;
+}
+
+namespace {
+struct JumpTableTy {
+  Value *Index;
+  SmallVector<Function *, 5> Funcs;
+};
+} // anonymous namespace
+
+static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP) {
+  if (!GEP || !GEP->isInBounds())
+    return std::nullopt;
+  ArrayType *ArrayTy = dyn_cast<ArrayType>(GEP->getSourceElementType());
+  if (!ArrayTy || ArrayTy->getArrayNumElements() > JumpTableSizeThreshold)
+    return std::nullopt;
+  Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand());
+  if (!Ptr)
+    return std::nullopt;
+
+  Function &F = *GEP->getParent()->getParent();
+  const DataLayout &DL = F.getParent()->getDataLayout();
+  const unsigned BitWidth =
+      DL.getIndexSizeInBits(GEP->getPointerAddressSpace());
+  MapVector<Value *, APInt> VariableOffsets;
+  APInt ConstantOffset(BitWidth, 0);
+  if (!GEP->collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset))
+    return std::nullopt;
+  if (VariableOffsets.empty() || VariableOffsets.size() > 1)
+    return std::nullopt;
+
+  Module &M = *F.getParent();
+  unsigned Offset = ConstantOffset.getZExtValue();
+  // TODO: support more general patterns
+  // (see also TODO in getElementWithGivenTypeAtOffset).
+  if (Offset != 0)
+    return std::nullopt;
+  if (!Ptr->getNumOperands())
+    return std::nullopt;
+  Constant *ConstArray = getElementWithGivenTypeAtOffset(
+      cast<Constant>(Ptr->getOperand(0)), ArrayTy, Offset, M);
+  if (!ConstArray)
+    return std::nullopt;
+
+  JumpTableTy JumpTable;
+  JumpTable.Index = VariableOffsets.front().first;
+
+  const uint64_t N = ArrayTy->getArrayNumElements();
+  JumpTable.Funcs.assign(N, nullptr);
+  for (uint64_t Index = 0; Index < N; ++Index) {
+    auto *Func =
+        dyn_cast_or_null<Function>(ConstArray->getAggregateElement(Index));
+    if (!Func || Func->isDeclaration())
+      return std::nullopt;
+    JumpTable.Funcs[Index] = Func;
+  }
+  return JumpTable;
+}
+
+static BasicBlock *split(CallBase *CB, const JumpTableTy &JT,
+                         DomTreeUpdater *DTU) {
+  const bool IsVoid = CB->getType() == Type::getVoidTy(CB->getContext());
+
+  SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
+  BasicBlock *BB = CB->getParent();
+  BasicBlock *Tail =
+      SplitBlock(BB, CB, DTU, nullptr, nullptr, BB->getName() + Twine(".tail"));
+  DTUpdates.push_back({DominatorTree::Delete, BB, Tail});
+  BB->getTerminator()->eraseFromParent();
+
+  Function &F = *BB->getParent();
+  BasicBlock *BBUnreachable = BasicBlock::Create(
+      F.getContext(), "default.switch.case.unreachable", &F, Tail);
+  IRBuilder<> BuilderUnreachable(BBUnreachable);
+  BuilderUnreachable.CreateUnreachable();
+
+  IRBuilder<> Builder(BB);
+  SwitchInst *Switch = Builder.CreateSwitch(JT.Index, BBUnreachable);
+  DTUpdates.push_back({DominatorTree::Insert, BB, BBUnreachable});
+
+  IRBuilder<> BuilderTail(CB);
+  PHINode *PHI =
+      IsVoid ? nullptr : BuilderTail.CreatePHI(CB->getType(), JT.Funcs.size());
+
+  for (auto [Index, Func] : llvm::enumerate(JT.Funcs)) {
+    BasicBlock *B = BasicBlock::Create(Func->getContext(),
+                                       "call." + Twine(Index), &F, Tail);
+    DTUpdates.push_back({DominatorTree::Insert, BB, B});
+    DTUpdates.push_back({DominatorTree::Insert, B, Tail});
+
+    CallBase *Call = cast<CallBase>(CB->clone());
+    Call->setCalledFunction(Func);
+    Call->insertInto(B, B->end());
+    Switch->addCase(
+        cast<ConstantInt>(ConstantInt::get(JT.Index->getType(), Index)), B);
+    BranchInst::Create(Tail, B);
+    if (PHI)
+      PHI->addIncoming(Call, B);
+  }
+  if (DTU)
+    DTU->applyUpdates(DTUpdates);
+  if (PHI)
+    CB->replaceAllUsesWith(PHI);
+  CB->eraseFromParent();
+  return Tail;
+}
+
+PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
+                                             FunctionAnalysisManager &AM) {
+  DominatorTree *DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
+  PostDominatorTree *PDT = AM.getCachedResult<PostDominatorTreeAnalysis>(F);
+  std::unique_ptr<DomTreeUpdater> DTU;
+  bool Changed = false;
+  for (BasicBlock &BB : make_early_inc_range(F)) {
+    BasicBlock *CurrentBB = &BB;
+    while (CurrentBB) {
+      BasicBlock *SplittedOutTail = nullptr;
+      for (Instruction &I : make_early_inc_range(*CurrentBB)) {
+        CallBase *CB = dyn_cast<CallBase>(&I);
+        if (!CB || isa<IntrinsicInst>(CB) || CB->getCalledFunction() ||
+            isa<InvokeInst>(CB) || CB->isMustTailCall())
+          continue;
+
+        Value *V;
+        if (!match(CB->getCalledOperand(), m_Load(m_Value(V))))
+          continue;
+        auto *GEP = dyn_cast<GetElementPtrInst>(V);
+        if (!GEP)
+          continue;
+
+        std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP);
+        if (!JumpTable)
+          continue;
+        if ((DT || PDT) && !DTU)
+          DTU = std::make_unique<DomTreeUpdater>(
+              DT, PDT, DomTreeUpdater::UpdateStrategy::Lazy);
+        SplittedOutTail = split(CB, *JumpTable, DTU.get());
+        Changed = true;
+        break;
+      }
+      CurrentBB = SplittedOutTail ? SplittedOutTail : nullptr;
+    }
+  }
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  if (DT)
+    PA.preserve<DominatorTreeAnalysis>();
+  if (PDT)
+    PA.preserve<PostDominatorTreeAnalysis>();
+  return PA;
+}

llvm/test/Other/new-pm-defaults.ll

@@ -149,6 +149,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

@@ -88,6 +88,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll

@@ -76,6 +76,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll

@@ -84,6 +84,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-prelink-defaults.ll

@@ -119,6 +119,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll

@@ -116,6 +116,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass

llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll

@@ -88,6 +88,7 @@
 ; CHECK-O23SZ-NEXT: Running pass: SpeculativeExecutionPass
 ; CHECK-O23SZ-NEXT: Running pass: JumpThreadingPass
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass