[llvm] Add a pass to convert jump tables to switches. (PR #77709)

Alexander Shaposhnikov via llvm-commits llvm-commits at lists.llvm.org
Sun Jan 21 02:44:14 PST 2024


https://github.com/alexander-shaposhnikov updated https://github.com/llvm/llvm-project/pull/77709

>From e03138380cf44fb723e954a8a20b9dc1fbc7a636 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <ashaposhnikov at google.com>
Date: Thu, 11 Jan 2024 00:58:53 +0000
Subject: [PATCH] Add a pass to convert jump tables to switches

---
 .../Transforms/Scalar/JumpTableToSwitch.h     |  24 +++
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 llvm/lib/Passes/PassBuilderPipelines.cpp      |   2 +
 llvm/lib/Passes/PassRegistry.def              |   1 +
 llvm/lib/Transforms/Scalar/CMakeLists.txt     |   1 +
 .../Transforms/Scalar/JumpTableToSwitch.cpp   | 193 ++++++++++++++++++
 llvm/test/Other/new-pm-defaults.ll            |   1 +
 .../Other/new-pm-thinlto-postlink-defaults.ll |   1 +
 .../new-pm-thinlto-postlink-pgo-defaults.ll   |   1 +
 ...-pm-thinlto-postlink-samplepgo-defaults.ll |   1 +
 .../Other/new-pm-thinlto-prelink-defaults.ll  |   1 +
 .../new-pm-thinlto-prelink-pgo-defaults.ll    |   1 +
 ...w-pm-thinlto-prelink-samplepgo-defaults.ll |   1 +
 .../Transforms/JumpTableToSwitch/basic.ll     | 169 +++++++++++++++
 .../test/Transforms/JumpTableToSwitch/skip.ll | 131 ++++++++++++
 .../Transforms/JumpTableToSwitch/struct.ll    |  42 ++++
 16 files changed, 571 insertions(+)
 create mode 100644 llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h
 create mode 100644 llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
 create mode 100644 llvm/test/Transforms/JumpTableToSwitch/basic.ll
 create mode 100644 llvm/test/Transforms/JumpTableToSwitch/skip.ll
 create mode 100644 llvm/test/Transforms/JumpTableToSwitch/struct.ll

diff --git a/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h b/llvm/include/llvm/Transforms/Scalar/JumpTableToSwitch.h
new file mode 100644
index 00000000000000..61786227d7a335
--- /dev/null
+++ b/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
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 8d3f69be503831..eebc5b362420e9 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -198,6 +198,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"
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 5c6c391049a7b2..a82be03bed7cd7 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/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"
@@ -558,6 +559,7 @@ PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
   // Optimize based on known information about branches, and cleanup afterward.
   FPM.addPass(JumpThreadingPass());
   FPM.addPass(CorrelatedValuePropagationPass());
+  FPM.addPass(JumpTableToSwitchPass());
 
   FPM.addPass(
       SimplifyCFGPass(SimplifyCFGOptions().convertSwitchRangeToICmp(true)));
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 756f4ab23ab66c..303c5feee55db2 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -342,6 +342,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())
diff --git a/llvm/lib/Transforms/Scalar/CMakeLists.txt b/llvm/lib/Transforms/Scalar/CMakeLists.txt
index 2dd27037a17de7..7b49c8010c874a 100644
--- a/llvm/lib/Transforms/Scalar/CMakeLists.txt
+++ b/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
diff --git a/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
new file mode 100644
index 00000000000000..85c23066dfd134
--- /dev/null
+++ b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
@@ -0,0 +1,193 @@
+//===- 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/ConstantFolding.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"
+
+namespace {
+struct JumpTableTy {
+  Value *Index;
+  SmallVector<Function *, 5> Funcs;
+};
+} // anonymous namespace
+
+static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP,
+                                                 PointerType *PtrTy) {
+  if (!GEP->isInBounds())
+    return std::nullopt;
+
+  Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand());
+  if (!Ptr || !Ptr->getNumOperands())
+    return std::nullopt;
+
+  GlobalVariable *GV = dyn_cast<GlobalVariable>(Ptr);
+  if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
+    return std::nullopt;
+
+  Function &F = *GEP->getParent()->getParent();
+  const DataLayout &DL = F.getParent()->getDataLayout();
+  const unsigned PtrSizeBytes = DL.getPointerTypeSize(PtrTy);
+  const uint64_t JumpTableSizeBytes = DL.getTypeAllocSize(GV->getValueType());
+  if (JumpTableSizeBytes % PtrSizeBytes != 0)
+    return std::nullopt;
+  const uint64_t N = JumpTableSizeBytes / PtrSizeBytes;
+  if (N > JumpTableSizeThreshold)
+    return std::nullopt;
+
+  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;
+  // TODO: consider supporting more general patterns
+  if (!ConstantOffset.isZero())
+    return std::nullopt;
+
+  JumpTableTy JumpTable;
+  JumpTable.Index = VariableOffsets.front().first;
+  JumpTable.Funcs.assign(N, nullptr);
+  const unsigned PtrSizeBits = DL.getPointerTypeSizeInBits(PtrTy);
+  for (uint64_t Index = 0; Index < N; ++Index) {
+    APInt Offset(PtrSizeBits, Index * PtrSizeBytes);
+    Offset += ConstantOffset;
+    Constant *C = ConstantFoldLoadFromConst(cast<Constant>(GV->getOperand(0)),
+                                            PtrTy, Offset, DL);
+    auto *Func = dyn_cast_or_null<Function>(C);
+    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);
+  }
+  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);
+  DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Lazy);
+  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)) {
+        auto *Call = dyn_cast<CallInst>(&I);
+        if (!Call || isa<IntrinsicInst>(Call) || Call->getCalledFunction() ||
+            Call->isMustTailCall())
+          continue;
+
+        Value *V;
+        if (!match(Call->getCalledOperand(), m_Load(m_Value(V))))
+          continue;
+        auto *L = cast<LoadInst>(Call->getCalledOperand());
+        // Skip atomic or volatile loads.
+        if (!L->isSimple())
+          continue;
+        auto *GEP = dyn_cast<GetElementPtrInst>(V);
+        if (!GEP)
+          continue;
+        PointerType *PtrTy = dyn_cast<PointerType>(L->getType());
+        if (!PtrTy)
+          continue;
+        std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy);
+        if (!JumpTable)
+          continue;
+        SplittedOutTail = split(Call, *JumpTable, DTU);
+        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;
+}
diff --git a/llvm/test/Other/new-pm-defaults.ll b/llvm/test/Other/new-pm-defaults.ll
index ecdb5a5e010d92..285077ff8e31a9 100644
--- a/llvm/test/Other/new-pm-defaults.ll
+++ b/llvm/test/Other/new-pm-defaults.ll
@@ -151,6 +151,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Other/new-pm-thinlto-postlink-defaults.ll b/llvm/test/Other/new-pm-thinlto-postlink-defaults.ll
index 064362eabbf839..29a4d79037427e 100644
--- a/llvm/test/Other/new-pm-thinlto-postlink-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-postlink-defaults.ll
@@ -90,6 +90,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll b/llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll
index 19a44867e434ac..bf06782c86f862 100644
--- a/llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll
@@ -78,6 +78,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll b/llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll
index ac80a31d8fd4bc..0cc61121de01ce 100644
--- a/llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll
@@ -86,6 +86,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Other/new-pm-thinlto-prelink-defaults.ll b/llvm/test/Other/new-pm-thinlto-prelink-defaults.ll
index 6486639e07b49c..0e5839797afe9e 100644
--- a/llvm/test/Other/new-pm-thinlto-prelink-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-prelink-defaults.ll
@@ -121,6 +121,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll b/llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll
index 09f9f0f48baddb..68c2e581463000 100644
--- a/llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll
@@ -118,6 +118,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O-NEXT: Running analysis: BlockFrequencyAnalysis on foo
diff --git a/llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll b/llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll
index 47bdbfd2d357d4..8311a009711d1e 100644
--- a/llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll
+++ b/llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll
@@ -90,6 +90,7 @@
 ; CHECK-O23SZ-NEXT: Running analysis: LazyValueAnalysis
 ; CHECK-O23SZ-NEXT: Running pass: CorrelatedValuePropagationPass
 ; CHECK-O23SZ-NEXT: Invalidating analysis: LazyValueAnalysis
+; CHECK-O23SZ-NEXT: Running pass: JumpTableToSwitchPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: InstCombinePass
 ; CHECK-O23SZ-NEXT: Running pass: AggressiveInstCombinePass
diff --git a/llvm/test/Transforms/JumpTableToSwitch/basic.ll b/llvm/test/Transforms/JumpTableToSwitch/basic.ll
new file mode 100644
index 00000000000000..534419e52a8852
--- /dev/null
+++ b/llvm/test/Transforms/JumpTableToSwitch/basic.ll
@@ -0,0 +1,169 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=jump-table-to-switch -verify-dom-info -S | FileCheck %s
+; RUN: opt < %s -passes=jump-table-to-switch -jump-table-to-switch-size-threshold=0 -verify-dom-info -S | FileCheck %s --check-prefix=THRESHOLD-0
+
+ at func_array = constant [2 x ptr] [ptr @func0, ptr @func1]
+
+define i32 @func0() {
+  ret i32 1
+}
+
+define i32 @func1() {
+  ret i32 2
+}
+
+define i32 @function_with_jump_table(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @func_array, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    switch i32 [[INDEX]], label [[DEFAULT_SWITCH_CASE_UNREACHABLE:%.*]] [
+; CHECK-NEXT:      i32 0, label [[CALL_0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CALL_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default.switch.case.unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       call.0:
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @func0()
+; CHECK-NEXT:    br label [[DOTTAIL:%.*]]
+; CHECK:       call.1:
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @func1()
+; CHECK-NEXT:    br label [[DOTTAIL]]
+; CHECK:       .tail:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i32 [ [[TMP1]], [[CALL_0]] ], [ [[TMP2]], [[CALL_1]] ]
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+; THRESHOLD-0-LABEL: define i32 @function_with_jump_table(
+; THRESHOLD-0-SAME: i32 [[INDEX:%.*]]) {
+; THRESHOLD-0-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @func_array, i32 0, i32 [[INDEX]]
+; THRESHOLD-0-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; THRESHOLD-0-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]]()
+; THRESHOLD-0-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @func_array, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep
+  %result = call i32 %func_ptr()
+  ret i32 %result
+}
+
+define void @void_func0() {
+  ret void
+}
+
+define void @void_func1() {
+  ret void
+}
+
+ at void_func_array = constant [2 x ptr] [ptr @void_func0, ptr @void_func1]
+
+define void @void_function_with_jump_table(i32 %index) {
+; CHECK-LABEL: define void @void_function_with_jump_table(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    switch i32 [[INDEX]], label [[DEFAULT_SWITCH_CASE_UNREACHABLE:%.*]] [
+; CHECK-NEXT:      i32 0, label [[CALL_0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CALL_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default.switch.case.unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       call.0:
+; CHECK-NEXT:    call void @void_func0()
+; CHECK-NEXT:    br label [[DOTTAIL:%.*]]
+; CHECK:       call.1:
+; CHECK-NEXT:    call void @void_func1()
+; CHECK-NEXT:    br label [[DOTTAIL]]
+; CHECK:       .tail:
+; CHECK-NEXT:    ret void
+;
+; THRESHOLD-0-LABEL: define void @void_function_with_jump_table(
+; THRESHOLD-0-SAME: i32 [[INDEX:%.*]]) {
+; THRESHOLD-0-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 [[INDEX]]
+; THRESHOLD-0-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; THRESHOLD-0-NEXT:    call void [[FUNC_PTR]]()
+; THRESHOLD-0-NEXT:    ret void
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep
+  call void %func_ptr()
+  ret void
+}
+
+define void @void_function_with_jump_table_and_call_site_attr(i32 %index) {
+; CHECK-LABEL: define void @void_function_with_jump_table_and_call_site_attr(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    switch i32 [[INDEX]], label [[DEFAULT_SWITCH_CASE_UNREACHABLE:%.*]] [
+; CHECK-NEXT:      i32 0, label [[CALL_0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CALL_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default.switch.case.unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       call.0:
+; CHECK-NEXT:    call void @void_func0() #[[ATTR0:[0-9]+]]
+; CHECK-NEXT:    br label [[DOTTAIL:%.*]]
+; CHECK:       call.1:
+; CHECK-NEXT:    call void @void_func1() #[[ATTR0]]
+; CHECK-NEXT:    br label [[DOTTAIL]]
+; CHECK:       .tail:
+; CHECK-NEXT:    ret void
+;
+; THRESHOLD-0-LABEL: define void @void_function_with_jump_table_and_call_site_attr(
+; THRESHOLD-0-SAME: i32 [[INDEX:%.*]]) {
+; THRESHOLD-0-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 [[INDEX]]
+; THRESHOLD-0-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; THRESHOLD-0-NEXT:    call void [[FUNC_PTR]]() #[[ATTR0:[0-9]+]]
+; THRESHOLD-0-NEXT:    ret void
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @void_func_array, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep
+  call void %func_ptr() nounwind
+  ret void
+}
+
+
+define i32 @func0_addrspace_42() addrspace(42) {
+  ret i32 1
+}
+
+define i32 @func1_addrspace_42() addrspace(42) {
+  ret i32 2
+}
+
+ at func_array_addrspace_42 = addrspace(42) constant [2 x ptr addrspace(42)] [ptr addrspace(42) @func0_addrspace_42, ptr addrspace(42) @func1_addrspace_42]
+
+define i32 @function_with_jump_table_addrspace_42(i32 %index) addrspace(42) {
+; CHECK-LABEL: define i32 @function_with_jump_table_addrspace_42(
+; CHECK-SAME: i32 [[INDEX:%.*]]) addrspace(42) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr addrspace(42)], ptr addrspace(42) @func_array_addrspace_42, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr addrspace(42), ptr addrspace(42) [[GEP]], align 8
+; CHECK-NEXT:    switch i32 [[INDEX]], label [[DEFAULT_SWITCH_CASE_UNREACHABLE:%.*]] [
+; CHECK-NEXT:      i32 0, label [[CALL_0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CALL_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default.switch.case.unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       call.0:
+; CHECK-NEXT:    [[TMP1:%.*]] = call addrspace(42) i32 @func0_addrspace_42()
+; CHECK-NEXT:    br label [[DOTTAIL:%.*]]
+; CHECK:       call.1:
+; CHECK-NEXT:    [[TMP2:%.*]] = call addrspace(42) i32 @func1_addrspace_42()
+; CHECK-NEXT:    br label [[DOTTAIL]]
+; CHECK:       .tail:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i32 [ [[TMP1]], [[CALL_0]] ], [ [[TMP2]], [[CALL_1]] ]
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+; THRESHOLD-0-LABEL: define i32 @function_with_jump_table_addrspace_42(
+; THRESHOLD-0-SAME: i32 [[INDEX:%.*]]) addrspace(42) {
+; THRESHOLD-0-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr addrspace(42)], ptr addrspace(42) @func_array_addrspace_42, i32 0, i32 [[INDEX]]
+; THRESHOLD-0-NEXT:    [[FUNC_PTR:%.*]] = load ptr addrspace(42), ptr addrspace(42) [[GEP]], align 8
+; THRESHOLD-0-NEXT:    [[RESULT:%.*]] = call addrspace(42) i32 [[FUNC_PTR]]()
+; THRESHOLD-0-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [2 x ptr addrspace(42)], ptr addrspace(42) @func_array_addrspace_42, i32 0, i32 %index
+  %func_ptr = load ptr addrspace(42), ptr addrspace(42) %gep, align 8
+  %result = call addrspace(42) i32 %func_ptr()
+  ret i32 %result
+}
+
diff --git a/llvm/test/Transforms/JumpTableToSwitch/skip.ll b/llvm/test/Transforms/JumpTableToSwitch/skip.ll
new file mode 100644
index 00000000000000..4504423a015236
--- /dev/null
+++ b/llvm/test/Transforms/JumpTableToSwitch/skip.ll
@@ -0,0 +1,131 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=jump-table-to-switch -verify-dom-info -S | FileCheck %s
+
+ at func_array0 = constant [2 x ptr] [ptr @func0, ptr @declared_only_func1]
+
+define i32 @func0() {
+  ret i32 1
+}
+
+declare i32 @declared_only_func1()
+
+define i32 @function_with_jump_table_with_a_declared_only_func(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table_with_a_declared_only_func(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @func_array0, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]]()
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @func_array0, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep, align 8
+  %result = call i32 %func_ptr()
+  ret i32 %result
+}
+
+declare i32 @__gxx_personality_v0(...)
+
+define i32 @function_with_jump_table_invoke(i32 %index) personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define i32 @function_with_jump_table_invoke(
+; CHECK-SAME: i32 [[INDEX:%.*]]) personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @func_array0, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = invoke i32 [[FUNC_PTR]]()
+; CHECK-NEXT:            to label [[NORMAL:%.*]] unwind label [[EXCEPTIONAL:%.*]]
+; CHECK:       normal:
+; CHECK-NEXT:    ret i32 [[RESULT]]
+; CHECK:       exceptional:
+; CHECK-NEXT:    [[LANDING_PAD:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            catch ptr null
+; CHECK-NEXT:    resume { ptr, i32 } [[LANDING_PAD]]
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @func_array0, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep, align 8
+  %result = invoke i32 %func_ptr() to label %normal unwind label %exceptional
+normal:
+  ret i32 %result
+exceptional:
+  %landing_pad = landingpad { ptr, i32 } catch ptr null
+  resume { ptr, i32 } %landing_pad
+}
+
+ at func_array1 = constant [1 x ptr] [ptr @func2]
+
+define i32 @func2(i32 %arg) {
+  ret i32 %arg
+}
+
+define i32 @function_with_jump_table_musttail_call(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table_musttail_call(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = musttail call i32 [[FUNC_PTR]](i32 [[INDEX]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep, align 8
+  %result = musttail call i32 %func_ptr(i32 %index)
+  ret i32 %result
+}
+
+define i32 @function_with_jump_table_and_volatile_load(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table_and_volatile_load(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load volatile ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]](i32 [[INDEX]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 %index
+  %func_ptr = load volatile ptr, ptr %gep, align 8
+  %result = call i32 %func_ptr(i32 %index)
+  ret i32 %result
+}
+
+define i32 @function_with_jump_table_and_atomic_load(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table_and_atomic_load(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load atomic ptr, ptr [[GEP]] monotonic, align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]](i32 [[INDEX]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [1 x ptr], ptr @func_array1, i32 0, i32 %index
+  %func_ptr = load atomic ptr, ptr %gep monotonic, align 8
+  %result = call i32 %func_ptr(i32 %index)
+  ret i32 %result
+}
+
+ at func_array2 = global [1 x ptr] [ptr @func2]
+
+define i32 @function_with_nonconstant_jump_table(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_nonconstant_jump_table(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [1 x ptr], ptr @func_array2, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]](i32 [[INDEX]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [1 x ptr], ptr @func_array2, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep, align 8
+  %result = call i32 %func_ptr(i32 %index)
+  ret i32 %result
+}
+
+ at func_array3 = weak constant [1 x ptr] [ptr @func2]
+
+define i32 @function_with_constant_weak_jump_table(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_constant_weak_jump_table(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [1 x ptr], ptr @func_array3, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 [[FUNC_PTR]](i32 [[INDEX]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %gep = getelementptr inbounds [1 x ptr], ptr @func_array3, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep, align 8
+  %result = call i32 %func_ptr(i32 %index)
+  ret i32 %result
+}
+
diff --git a/llvm/test/Transforms/JumpTableToSwitch/struct.ll b/llvm/test/Transforms/JumpTableToSwitch/struct.ll
new file mode 100644
index 00000000000000..7aa709c402fe1f
--- /dev/null
+++ b/llvm/test/Transforms/JumpTableToSwitch/struct.ll
@@ -0,0 +1,42 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=jump-table-to-switch -verify-dom-info -S | FileCheck %s
+
+%"struct_ty" = type { [2 x ptr] }
+
+ at func_array = constant %"struct_ty" { [2 x ptr] [ptr @func0, ptr @func1] }
+
+define i32 @func0() {
+  ret i32 1
+}
+
+define i32 @func1() {
+  ret i32 2
+}
+
+define i32 @function_with_jump_table(i32 %index) {
+; CHECK-LABEL: define i32 @function_with_jump_table(
+; CHECK-SAME: i32 [[INDEX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [2 x ptr], ptr @func_array, i32 0, i32 [[INDEX]]
+; CHECK-NEXT:    [[FUNC_PTR:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    switch i32 [[INDEX]], label [[DEFAULT_SWITCH_CASE_UNREACHABLE:%.*]] [
+; CHECK-NEXT:      i32 0, label [[CALL_0:%.*]]
+; CHECK-NEXT:      i32 1, label [[CALL_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default.switch.case.unreachable:
+; CHECK-NEXT:    unreachable
+; CHECK:       call.0:
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @func0()
+; CHECK-NEXT:    br label [[DOTTAIL:%.*]]
+; CHECK:       call.1:
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @func1()
+; CHECK-NEXT:    br label [[DOTTAIL]]
+; CHECK:       .tail:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i32 [ [[TMP1]], [[CALL_0]] ], [ [[TMP2]], [[CALL_1]] ]
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+  %gep = getelementptr inbounds [2 x ptr], ptr @func_array, i32 0, i32 %index
+  %func_ptr = load ptr, ptr %gep
+  %result = call i32 %func_ptr()
+  ret i32 %result
+}
+



More information about the llvm-commits mailing list