[CodeGen] Port IndirectBrExpand to new pass manager (#75287)

Author: paperchalice

llvm/include/llvm/CodeGen/CodeGenPassBuilder.h

@@ -27,6 +27,7 @@
 #include "llvm/CodeGen/DwarfEHPrepare.h"
 #include "llvm/CodeGen/ExpandReductions.h"
 #include "llvm/CodeGen/GCMetadata.h"
+#include "llvm/CodeGen/IndirectBrExpand.h"
 #include "llvm/CodeGen/InterleavedAccess.h"
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"

llvm/include/llvm/CodeGen/IndirectBrExpand.h

@@ -0,0 +1,28 @@
+//===- llvm/CodeGen/IndirectBrExpand.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_CODEGEN_INDIRECTBREXPAND_H
+#define LLVM_CODEGEN_INDIRECTBREXPAND_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class TargetMachine;
+
+class IndirectBrExpandPass : public PassInfoMixin<IndirectBrExpandPass> {
+  const TargetMachine *TM;
+
+public:
+  IndirectBrExpandPass(const TargetMachine *TM) : TM(TM) {}
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_INDIRECTBREXPAND_H

llvm/include/llvm/CodeGen/MachinePassRegistry.def

@@ -48,6 +48,7 @@ FUNCTION_PASS("expand-large-div-rem", ExpandLargeDivRemPass, ())
 FUNCTION_PASS("expand-large-fp-convert", ExpandLargeFpConvertPass, ())
 FUNCTION_PASS("expand-reductions", ExpandReductionsPass, ())
 FUNCTION_PASS("expandvp", ExpandVectorPredicationPass, ())
+FUNCTION_PASS("indirectbr-expand", IndirectBrExpandPass, (TM))
 FUNCTION_PASS("interleaved-access", InterleavedAccessPass, (TM))
 FUNCTION_PASS("interleaved-load-combine", InterleavedLoadCombinePass, (TM))
 FUNCTION_PASS("lower-constant-intrinsics", LowerConstantIntrinsicsPass, ())
@@ -131,7 +132,6 @@ DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
 DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
 DUMMY_FUNCTION_PASS("expandmemcmp", ExpandMemCmpPass, ())
 DUMMY_FUNCTION_PASS("gc-lowering", GCLoweringPass, ())
-DUMMY_FUNCTION_PASS("indirectbr-expand", IndirectBrExpandPass, ())
 DUMMY_FUNCTION_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
 DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
 #undef DUMMY_FUNCTION_PASS

llvm/include/llvm/InitializePasses.h

@@ -129,7 +129,7 @@ void initializeIVUsersWrapperPassPass(PassRegistry&);
 void initializeIfConverterPass(PassRegistry&);
 void initializeImmutableModuleSummaryIndexWrapperPassPass(PassRegistry&);
 void initializeImplicitNullChecksPass(PassRegistry&);
-void initializeIndirectBrExpandPassPass(PassRegistry&);
+void initializeIndirectBrExpandLegacyPassPass(PassRegistry &);
 void initializeInferAddressSpacesPass(PassRegistry&);
 void initializeInstSimplifyLegacyPassPass(PassRegistry &);
 void initializeInstructionCombiningPassPass(PassRegistry&);

llvm/lib/CodeGen/CodeGen.cpp

@@ -53,7 +53,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeHardwareLoopsLegacyPass(Registry);
   initializeIfConverterPass(Registry);
   initializeImplicitNullChecksPass(Registry);
-  initializeIndirectBrExpandPassPass(Registry);
+  initializeIndirectBrExpandLegacyPassPass(Registry);
   initializeInterleavedLoadCombinePass(Registry);
   initializeInterleavedAccessPass(Registry);
   initializeJMCInstrumenterPass(Registry);

llvm/lib/CodeGen/IndirectBrExpandPass.cpp

@@ -29,6 +29,7 @@
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/CodeGen/IndirectBrExpand.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
@@ -48,14 +49,12 @@ using namespace llvm;
 
 namespace {
 
-class IndirectBrExpandPass : public FunctionPass {
-  const TargetLowering *TLI = nullptr;
-
+class IndirectBrExpandLegacyPass : public FunctionPass {
 public:
   static char ID; // Pass identification, replacement for typeid
 
-  IndirectBrExpandPass() : FunctionPass(ID) {
-    initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry());
+  IndirectBrExpandLegacyPass() : FunctionPass(ID) {
+    initializeIndirectBrExpandLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -67,33 +66,41 @@ class IndirectBrExpandPass : public FunctionPass {
 
 } // end anonymous namespace
 
-char IndirectBrExpandPass::ID = 0;
+static bool runImpl(Function &F, const TargetLowering *TLI,
+                    DomTreeUpdater *DTU);
+
+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);
+  if (!Changed)
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<DominatorTreeAnalysis>();
+  return PA;
+}
+
+char IndirectBrExpandLegacyPass::ID = 0;
 
-INITIALIZE_PASS_BEGIN(IndirectBrExpandPass, DEBUG_TYPE,
+INITIALIZE_PASS_BEGIN(IndirectBrExpandLegacyPass, DEBUG_TYPE,
                       "Expand indirectbr instructions", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
-INITIALIZE_PASS_END(IndirectBrExpandPass, DEBUG_TYPE,
+INITIALIZE_PASS_END(IndirectBrExpandLegacyPass, DEBUG_TYPE,
                     "Expand indirectbr instructions", false, false)
 
 FunctionPass *llvm::createIndirectBrExpandPass() {
-  return new IndirectBrExpandPass();
+  return new IndirectBrExpandLegacyPass();
 }
 
-bool IndirectBrExpandPass::runOnFunction(Function &F) {
+bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU) {
   auto &DL = F.getParent()->getDataLayout();
-  auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
-  if (!TPC)
-    return false;
-
-  auto &TM = TPC->getTM<TargetMachine>();
-  auto &STI = *TM.getSubtargetImpl(F);
-  if (!STI.enableIndirectBrExpand())
-    return false;
-  TLI = STI.getTargetLowering();
-
-  std::optional<DomTreeUpdater> DTU;
-  if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
-    DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
 
   SmallVector<IndirectBrInst *, 1> IndirectBrs;
 
@@ -268,3 +275,21 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
 
   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);
+}

llvm/lib/Passes/PassBuilder.cpp

@@ -78,6 +78,7 @@
 #include "llvm/CodeGen/ExpandLargeFpConvert.h"
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/HardwareLoops.h"
+#include "llvm/CodeGen/IndirectBrExpand.h"
 #include "llvm/CodeGen/InterleavedAccess.h"
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"

llvm/lib/Passes/PassRegistry.def

@@ -315,6 +315,7 @@ FUNCTION_PASS("guard-widening", GuardWideningPass())
 FUNCTION_PASS("gvn-hoist", GVNHoistPass())
 FUNCTION_PASS("gvn-sink", GVNSinkPass())
 FUNCTION_PASS("helloworld", HelloWorldPass())
+FUNCTION_PASS("indirectbr-expand", IndirectBrExpandPass(TM))
 FUNCTION_PASS("infer-address-spaces", InferAddressSpacesPass())
 FUNCTION_PASS("infer-alignment", InferAlignmentPass())
 FUNCTION_PASS("inject-tli-mappings", InjectTLIMappings())

llvm/test/Transforms/IndirectBrExpand/basic.ll

@@ -1,4 +1,5 @@
 ; RUN: opt < %s -indirectbr-expand -S | FileCheck %s
+; RUN: opt < %s -passes=indirectbr-expand -S | FileCheck %s
 ;
 ; REQUIRES: x86-registered-target
 

llvm/tools/opt/opt.cpp

@@ -434,7 +434,7 @@ int main(int argc, char **argv) {
   initializeSjLjEHPreparePass(Registry);
   initializePreISelIntrinsicLoweringLegacyPassPass(Registry);
   initializeGlobalMergePass(Registry);
-  initializeIndirectBrExpandPassPass(Registry);
+  initializeIndirectBrExpandLegacyPassPass(Registry);
   initializeInterleavedLoadCombinePass(Registry);
   initializeInterleavedAccessPass(Registry);
   initializeUnreachableBlockElimLegacyPassPass(Registry);