[CodeGen] Port CallBrPrepare to new pass manager

[paperchalice]

IIUC in the new pass manager infrastructure, the analysis result is always computed lazily. So just use getResult here.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: None (paperchalice)

Changes

IIUC in the new pass manager infrastructure, the analysis result is always computed lazily. So just use getResult here.

---

Full diff: https://github.com/llvm/llvm-project/pull/73630.diff

5 Files Affected:

• (added) llvm/include/llvm/CodeGen/CallBrPrepare.h (+23)
• (modified) llvm/lib/CodeGen/CallBrPrepare.cpp (+33-14)
• (modified) llvm/lib/Passes/PassBuilder.cpp (+1)
• (modified) llvm/lib/Passes/PassRegistry.def (+1)
• (modified) llvm/test/CodeGen/AArch64/callbr-prepare.ll (+1)

diff --git a/llvm/include/llvm/CodeGen/CallBrPrepare.h b/llvm/include/llvm/CodeGen/CallBrPrepare.h
new file mode 100644
index 000000000000000..989343b02d024ad
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/CallBrPrepare.h
@@ -0,0 +1,23 @@
+//===-- CallBrPrepare - Prepare callbr for code generation ------*- 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_CALLBRPREPARE_H
+#define LLVM_CODEGEN_CALLBRPREPARE_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class CallBrPreparePass : public PassInfoMixin<CallBrPreparePass> {
+public:
+  PreservedAnalyses run(Function &Fn, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_CALLBRPREPARE_H
diff --git a/llvm/lib/CodeGen/CallBrPrepare.cpp b/llvm/lib/CodeGen/CallBrPrepare.cpp
index db243a0bfebe1a7..fddc4d74b2da918 100644
--- a/llvm/lib/CodeGen/CallBrPrepare.cpp
+++ b/llvm/lib/CodeGen/CallBrPrepare.cpp
@@ -31,6 +31,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/CallBrPrepare.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -53,15 +54,16 @@ using namespace llvm;
 
 #define DEBUG_TYPE "callbrprepare"
 
+static bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);
+static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,
+                                 DominatorTree &DT);
+static void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
+                      SSAUpdater &SSAUpdate);
+static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn);
+
 namespace {
 
 class CallBrPrepare : public FunctionPass {
-  bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);
-  bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,
-                            DominatorTree &DT) const;
-  void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
-                 SSAUpdater &SSAUpdate) const;
-
 public:
   CallBrPrepare() : FunctionPass(ID) {}
   void getAnalysisUsage(AnalysisUsage &AU) const override;
@@ -71,6 +73,26 @@ class CallBrPrepare : public FunctionPass {
 
 } // end anonymous namespace
 
+PreservedAnalyses CallBrPreparePass::run(Function &Fn,
+                                         FunctionAnalysisManager &FAM) {
+  bool Changed = false;
+  SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn);
+
+  if (CBRs.empty())
+    return PreservedAnalyses::all();
+
+  auto &DT = FAM.getResult<DominatorTreeAnalysis>(Fn);
+
+  Changed |= SplitCriticalEdges(CBRs, DT);
+  Changed |= InsertIntrinsicCalls(CBRs, DT);
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<DominatorTreeAnalysis>();
+  return PA;
+}
+
 char CallBrPrepare::ID = 0;
 INITIALIZE_PASS_BEGIN(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
@@ -82,7 +104,7 @@ void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addPreserved<DominatorTreeWrapperPass>();
 }
 
-static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) {
+SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) {
   SmallVector<CallBrInst *, 2> CBRs;
   for (BasicBlock &BB : Fn)
     if (auto *CBR = dyn_cast<CallBrInst>(BB.getTerminator()))
@@ -91,8 +113,7 @@ static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) {
   return CBRs;
 }
 
-bool CallBrPrepare::SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs,
-                                       DominatorTree &DT) {
+bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {
   bool Changed = false;
   CriticalEdgeSplittingOptions Options(&DT);
   Options.setMergeIdenticalEdges();
@@ -114,8 +135,7 @@ bool CallBrPrepare::SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs,
   return Changed;
 }
 
-bool CallBrPrepare::InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,
-                                         DominatorTree &DT) const {
+bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {
   bool Changed = false;
   SmallPtrSet<const BasicBlock *, 4> Visited;
   IRBuilder<> Builder(CBRs[0]->getContext());
@@ -160,9 +180,8 @@ static void PrintDebugDomInfo(const DominatorTree &DT, const Use &U,
 }
 #endif
 
-void CallBrPrepare::UpdateSSA(DominatorTree &DT, CallBrInst *CBR,
-                              CallInst *Intrinsic,
-                              SSAUpdater &SSAUpdate) const {
+void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
+               SSAUpdater &SSAUpdate) {
 
   SmallPtrSet<Use *, 4> Visited;
   BasicBlock *DefaultDest = CBR->getDefaultDest();
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index e06370c5463fa09..aeb9726a186b51e 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -72,6 +72,7 @@
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
 #include "llvm/Analysis/UniformityAnalysis.h"
+#include "llvm/CodeGen/CallBrPrepare.h"
 #include "llvm/CodeGen/DwarfEHPrepare.h"
 #include "llvm/CodeGen/ExpandLargeDivRem.h"
 #include "llvm/CodeGen/ExpandLargeFpConvert.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 3a9a37a231f8d2f..f64aed317034279 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -448,6 +448,7 @@ FUNCTION_PASS("declare-to-assign", llvm::AssignmentTrackingPass())
 FUNCTION_PASS("expand-large-div-rem", ExpandLargeDivRemPass(TM));
 FUNCTION_PASS("expand-large-fp-convert", ExpandLargeFpConvertPass(TM));
 FUNCTION_PASS("dwarfehprepare", DwarfEHPreparePass(TM));
+FUNCTION_PASS("callbrprepare", CallBrPreparePass());
 #undef FUNCTION_PASS
 
 #ifndef FUNCTION_PASS_WITH_PARAMS
diff --git a/llvm/test/CodeGen/AArch64/callbr-prepare.ll b/llvm/test/CodeGen/AArch64/callbr-prepare.ll
index 08b48d65c89a0cf..701355793d1cae6 100644
--- a/llvm/test/CodeGen/AArch64/callbr-prepare.ll
+++ b/llvm/test/CodeGen/AArch64/callbr-prepare.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -callbrprepare -S -o - | FileCheck %s
+; RUN: opt %s -passes=callbrprepare -S -o - | FileCheck %s
 
 define i32 @test0() {
 ; CHECK-LABEL: @test0(

[paperchalice]

CC @aeubanks @arsenm @nikic.

[paperchalice]

Not sure why MLIR test failed, try to re-run it.