[CodeGen][NewPM] Port TailDuplicate pass to NPM (#113293)

Author: optimisan

llvm/include/llvm/CodeGen/Passes.h

@@ -261,11 +261,11 @@ namespace llvm {
 
   /// TailDuplicate - Duplicate blocks with unconditional branches
   /// into tails of their predecessors.
-  extern char &TailDuplicateID;
+  extern char &TailDuplicateLegacyID;
 
   /// Duplicate blocks with unconditional branches into tails of their
   /// predecessors. Variant that works before register allocation.
-  extern char &EarlyTailDuplicateID;
+  extern char &EarlyTailDuplicateLegacyID;
 
   /// MachineTraceMetrics - This pass computes critical path and CPU resource
   /// usage in an ensemble of traces.

llvm/include/llvm/CodeGen/TailDuplication.h

@@ -0,0 +1,47 @@
+//===- llvm/CodeGen/TailDuplication.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_TAILDUPLICATIONPASS_H
+#define LLVM_CODEGEN_TAILDUPLICATIONPASS_H
+
+#include "llvm/CodeGen/MBFIWrapper.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+template <typename DerivedT, bool PreRegAlloc>
+class TailDuplicatePassBase : public PassInfoMixin<DerivedT> {
+private:
+  std::unique_ptr<MBFIWrapper> MBFIW;
+
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+class EarlyTailDuplicatePass
+    : public TailDuplicatePassBase<EarlyTailDuplicatePass, true> {
+public:
+  MachineFunctionProperties getClearedProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+};
+
+class TailDuplicatePass
+    : public TailDuplicatePassBase<TailDuplicatePass, false> {};
+
+} // namespace llvm
+
+extern template class llvm::TailDuplicatePassBase<llvm::EarlyTailDuplicatePass,
+                                                  true>;
+extern template class llvm::TailDuplicatePassBase<llvm::TailDuplicatePass,
+                                                  false>;
+
+#endif // LLVM_CODEGEN_TAILDUPLICATIONPASS_H

llvm/include/llvm/InitializePasses.h

@@ -101,7 +101,7 @@ void initializeEarlyCSEMemSSALegacyPassPass(PassRegistry &);
 void initializeEarlyIfConverterLegacyPass(PassRegistry &);
 void initializeEarlyIfPredicatorPass(PassRegistry &);
 void initializeEarlyMachineLICMPass(PassRegistry &);
-void initializeEarlyTailDuplicatePass(PassRegistry &);
+void initializeEarlyTailDuplicateLegacyPass(PassRegistry &);
 void initializeEdgeBundlesPass(PassRegistry &);
 void initializeEHContGuardCatchretPass(PassRegistry &);
 void initializeExpandLargeFpConvertLegacyPassPass(PassRegistry &);
@@ -300,7 +300,7 @@ void initializeStraightLineStrengthReduceLegacyPassPass(PassRegistry &);
 void initializeStripDebugMachineModulePass(PassRegistry &);
 void initializeStructurizeCFGLegacyPassPass(PassRegistry &);
 void initializeTailCallElimPass(PassRegistry &);
-void initializeTailDuplicatePass(PassRegistry &);
+void initializeTailDuplicateLegacyPass(PassRegistry &);
 void initializeTargetLibraryInfoWrapperPassPass(PassRegistry &);
 void initializeTargetPassConfigPass(PassRegistry &);
 void initializeTargetTransformInfoWrapperPassPass(PassRegistry &);

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -60,6 +60,7 @@
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/StackColoring.h"
 #include "llvm/CodeGen/StackProtector.h"
+#include "llvm/CodeGen/TailDuplication.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TwoAddressInstructionPass.h"
 #include "llvm/CodeGen/UnreachableBlockElim.h"

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -133,6 +133,7 @@ MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis())
 MACHINE_FUNCTION_PASS("dead-mi-elimination", DeadMachineInstructionElimPass())
 MACHINE_FUNCTION_PASS("early-ifcvt", EarlyIfConverterPass())
 MACHINE_FUNCTION_PASS("early-machinelicm", EarlyMachineLICMPass())
+MACHINE_FUNCTION_PASS("early-tailduplication", EarlyTailDuplicatePass())
 MACHINE_FUNCTION_PASS("finalize-isel", FinalizeISelPass())
 MACHINE_FUNCTION_PASS("localstackalloc", LocalStackSlotAllocationPass())
 MACHINE_FUNCTION_PASS("machine-cse", MachineCSEPass())
@@ -157,6 +158,7 @@ MACHINE_FUNCTION_PASS("print<virtregmap>", VirtRegMapPrinterPass(dbgs()))
 MACHINE_FUNCTION_PASS("require-all-machine-function-properties",
                       RequireAllMachineFunctionPropertiesPass())
 MACHINE_FUNCTION_PASS("stack-coloring", StackColoringPass())
+MACHINE_FUNCTION_PASS("tailduplication", TailDuplicatePass())
 MACHINE_FUNCTION_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 MACHINE_FUNCTION_PASS("two-address-instruction", TwoAddressInstructionPass())
 MACHINE_FUNCTION_PASS("verify", MachineVerifierPass())
@@ -210,7 +212,6 @@ DUMMY_MACHINE_FUNCTION_PASS("cfi-fixup", CFIFixupPass)
 DUMMY_MACHINE_FUNCTION_PASS("cfi-instr-inserter", CFIInstrInserterPass)
 DUMMY_MACHINE_FUNCTION_PASS("detect-dead-lanes", DetectDeadLanesPass)
 DUMMY_MACHINE_FUNCTION_PASS("dot-machine-cfg", MachineCFGPrinter)
-DUMMY_MACHINE_FUNCTION_PASS("early-tailduplication", EarlyTailDuplicatePass)
 DUMMY_MACHINE_FUNCTION_PASS("fentry-insert", FEntryInserterPass)
 DUMMY_MACHINE_FUNCTION_PASS("fixup-statepoint-caller-saved", FixupStatepointCallerSavedPass)
 DUMMY_MACHINE_FUNCTION_PASS("fs-profile-loader", MIRProfileLoaderNewPass)
@@ -262,7 +263,6 @@ DUMMY_MACHINE_FUNCTION_PASS("simple-register-coalescing", RegisterCoalescerPass)
 DUMMY_MACHINE_FUNCTION_PASS("stack-frame-layout", StackFrameLayoutAnalysisPass)
 DUMMY_MACHINE_FUNCTION_PASS("stack-slot-coloring", StackSlotColoringPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
-DUMMY_MACHINE_FUNCTION_PASS("tailduplication", TailDuplicatePass)
 DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
 DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)
 DUMMY_MACHINE_FUNCTION_PASS("xray-instrumentation", XRayInstrumentationPass)

llvm/lib/CodeGen/CodeGen.cpp

@@ -38,7 +38,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeEarlyIfConverterLegacyPass(Registry);
   initializeEarlyIfPredicatorPass(Registry);
   initializeEarlyMachineLICMPass(Registry);
-  initializeEarlyTailDuplicatePass(Registry);
+  initializeEarlyTailDuplicateLegacyPass(Registry);
   initializeExpandLargeDivRemLegacyPassPass(Registry);
   initializeExpandLargeFpConvertLegacyPassPass(Registry);
   initializeExpandMemCmpLegacyPassPass(Registry);
@@ -131,7 +131,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeStackProtectorPass(Registry);
   initializeStackSlotColoringPass(Registry);
   initializeStripDebugMachineModulePass(Registry);
-  initializeTailDuplicatePass(Registry);
+  initializeTailDuplicateLegacyPass(Registry);
   initializeTargetPassConfigPass(Registry);
   initializeTwoAddressInstructionLegacyPassPass(Registry);
   initializeTypePromotionLegacyPass(Registry);

llvm/lib/CodeGen/TailDuplication.cpp

@@ -12,13 +12,16 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/TailDuplication.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/LazyMachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/TailDuplicator.h"
+#include "llvm/IR/Analysis.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
@@ -29,13 +32,13 @@ using namespace llvm;
 
 namespace {
 
-class TailDuplicateBase : public MachineFunctionPass {
+class TailDuplicateBaseLegacy : public MachineFunctionPass {
   TailDuplicator Duplicator;
   std::unique_ptr<MBFIWrapper> MBFIW;
   bool PreRegAlloc;
 public:
-  TailDuplicateBase(char &PassID, bool PreRegAlloc)
-    : MachineFunctionPass(PassID), PreRegAlloc(PreRegAlloc) {}
+  TailDuplicateBaseLegacy(char &PassID, bool PreRegAlloc)
+      : MachineFunctionPass(PassID), PreRegAlloc(PreRegAlloc) {}
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
@@ -47,19 +50,19 @@ class TailDuplicateBase : public MachineFunctionPass {
   }
 };
 
-class TailDuplicate : public TailDuplicateBase {
+class TailDuplicateLegacy : public TailDuplicateBaseLegacy {
 public:
   static char ID;
-  TailDuplicate() : TailDuplicateBase(ID, false) {
-    initializeTailDuplicatePass(*PassRegistry::getPassRegistry());
+  TailDuplicateLegacy() : TailDuplicateBaseLegacy(ID, false) {
+    initializeTailDuplicateLegacyPass(*PassRegistry::getPassRegistry());
   }
 };
 
-class EarlyTailDuplicate : public TailDuplicateBase {
+class EarlyTailDuplicateLegacy : public TailDuplicateBaseLegacy {
 public:
   static char ID;
-  EarlyTailDuplicate() : TailDuplicateBase(ID, true) {
-    initializeEarlyTailDuplicatePass(*PassRegistry::getPassRegistry());
+  EarlyTailDuplicateLegacy() : TailDuplicateBaseLegacy(ID, true) {
+    initializeEarlyTailDuplicateLegacyPass(*PassRegistry::getPassRegistry());
   }
 
   MachineFunctionProperties getClearedProperties() const override {
@@ -70,17 +73,18 @@ class EarlyTailDuplicate : public TailDuplicateBase {
 
 } // end anonymous namespace
 
-char TailDuplicate::ID;
-char EarlyTailDuplicate::ID;
+char TailDuplicateLegacy::ID;
+char EarlyTailDuplicateLegacy::ID;
 
-char &llvm::TailDuplicateID = TailDuplicate::ID;
-char &llvm::EarlyTailDuplicateID = EarlyTailDuplicate::ID;
+char &llvm::TailDuplicateLegacyID = TailDuplicateLegacy::ID;
+char &llvm::EarlyTailDuplicateLegacyID = EarlyTailDuplicateLegacy::ID;
 
-INITIALIZE_PASS(TailDuplicate, DEBUG_TYPE, "Tail Duplication", false, false)
-INITIALIZE_PASS(EarlyTailDuplicate, "early-tailduplication",
+INITIALIZE_PASS(TailDuplicateLegacy, DEBUG_TYPE, "Tail Duplication", false,
+                false)
+INITIALIZE_PASS(EarlyTailDuplicateLegacy, "early-tailduplication",
                 "Early Tail Duplication", false, false)
 
-bool TailDuplicateBase::runOnMachineFunction(MachineFunction &MF) {
+bool TailDuplicateBaseLegacy::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
@@ -100,3 +104,36 @@ bool TailDuplicateBase::runOnMachineFunction(MachineFunction &MF) {
 
   return MadeChange;
 }
+
+template <typename DerivedT, bool PreRegAlloc>
+PreservedAnalyses TailDuplicatePassBase<DerivedT, PreRegAlloc>::run(
+    MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) {
+  MFPropsModifier _(static_cast<DerivedT &>(*this), MF);
+
+  if (MF.getFunction().hasOptNone())
+    return PreservedAnalyses::all();
+
+  auto *MBPI = &MFAM.getResult<MachineBranchProbabilityAnalysis>(MF);
+  auto *PSI = MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF)
+                  .getCachedResult<ProfileSummaryAnalysis>(
+                      *MF.getFunction().getParent());
+  auto *MBFI = (PSI && PSI->hasProfileSummary()
+                    ? &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF)
+                    : nullptr);
+  if (MBFI)
+    MBFIW = std::make_unique<MBFIWrapper>(*MBFI);
+
+  TailDuplicator Duplicator;
+  Duplicator.initMF(MF, PreRegAlloc, MBPI, MBFI ? MBFIW.get() : nullptr, PSI,
+                    /*LayoutMode=*/false);
+  bool MadeChange = false;
+  while (Duplicator.tailDuplicateBlocks())
+    MadeChange = true;
+
+  if (!MadeChange)
+    return PreservedAnalyses::all();
+  return getMachineFunctionPassPreservedAnalyses();
+}
+
+template class llvm::TailDuplicatePassBase<TailDuplicatePass, false>;
+template class llvm::TailDuplicatePassBase<EarlyTailDuplicatePass, true>;

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -290,10 +290,10 @@ static IdentifyingPassPtr overridePass(AnalysisID StandardID,
   if (StandardID == &BranchFolderPassID)
     return applyDisable(TargetID, DisableBranchFold);
 
-  if (StandardID == &TailDuplicateID)
+  if (StandardID == &TailDuplicateLegacyID)
     return applyDisable(TargetID, DisableTailDuplicate);
 
-  if (StandardID == &EarlyTailDuplicateID)
+  if (StandardID == &EarlyTailDuplicateLegacyID)
     return applyDisable(TargetID, DisableEarlyTailDup);
 
   if (StandardID == &MachineBlockPlacementID)
@@ -1279,7 +1279,7 @@ void TargetPassConfig::addMachinePasses() {
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
-  addPass(&EarlyTailDuplicateID);
+  addPass(&EarlyTailDuplicateLegacyID);
 
   // Optimize PHIs before DCE: removing dead PHI cycles may make more
   // instructions dead.
@@ -1507,7 +1507,7 @@ void TargetPassConfig::addMachineLateOptimization() {
   // performance for targets that require Structured Control Flow.
   // In addition it can also make CFG irreducible. Thus we disable it.
   if (!TM->requiresStructuredCFG())
-    addPass(&TailDuplicateID);
+    addPass(&TailDuplicateLegacyID);
 
   // Copy propagation.
   addPass(&MachineCopyPropagationID);

llvm/lib/Passes/PassBuilder.cpp

@@ -125,6 +125,7 @@
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/StackColoring.h"
 #include "llvm/CodeGen/StackProtector.h"
+#include "llvm/CodeGen/TailDuplication.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TwoAddressInstructionPass.h"
 #include "llvm/CodeGen/TypePromotion.h"

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -332,7 +332,7 @@ void NVPTXPassConfig::addIRPasses() {
   disablePass(&PrologEpilogCodeInserterID);
   disablePass(&MachineLateInstrsCleanupID);
   disablePass(&MachineCopyPropagationID);
-  disablePass(&TailDuplicateID);
+  disablePass(&TailDuplicateLegacyID);
   disablePass(&StackMapLivenessID);
   disablePass(&PostRAMachineSinkingID);
   disablePass(&PostRASchedulerID);
@@ -461,7 +461,7 @@ void NVPTXPassConfig::addOptimizedRegAlloc() {
 
 void NVPTXPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
-  if (addPass(&EarlyTailDuplicateID))
+  if (addPass(&EarlyTailDuplicateLegacyID))
     printAndVerify("After Pre-RegAlloc TailDuplicate");
 
   // Optimize PHIs before DCE: removing dead PHI cycles may make more

llvm/test/CodeGen/AArch64/jump-table-duplicate.mir

@@ -1,4 +1,5 @@
 # RUN: llc -run-pass=tailduplication -tail-dup-size=4 %s -o - | FileCheck %s
+# RUN: llc -passes=tailduplication -tail-dup-size=4 %s -o - | FileCheck %s
 
 # JumpTableDest32 uses an `adr` to a temporary label (itself). If duplicated we
 # cannot guarantee reachability for any uses after the first.

llvm/test/CodeGen/AMDGPU/early-tailduplicator-nophis.mir

@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=early-tailduplication -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -passes=early-tailduplication -o - %s | FileCheck %s
 
  # There are no phis in this testcase. Early tail duplication introduces them,
  # so the NoPHIs property needs to be cleared to avoid verifier errors

llvm/test/CodeGen/AMDGPU/early-tailduplicator-terminator.mir

@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=early-tailduplication -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -passes=early-tailduplication -o - %s | FileCheck %s
 
 # Early tail duplication should not merge bb.6 into bb.5, adding a
 # non-terminator (S_SLEEP) after the terminator S_MOV_B32_term.

llvm/test/CodeGen/AMDGPU/stop-tail-duplicate-cfg-intrinsic.mir

@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=early-tailduplication -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -passes=early-tailduplication -o - %s | FileCheck %s
 
 ---
 name:            stop_duplicate_cfg_intrinsic