7 files changed, 46 insertions, 33 deletions

diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index 4370d116e08..bab86bdb01e 100644
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -212,6 +212,10 @@ namespace llvm {
   /// into tails of their predecessors.
   extern char &TailDuplicateID;
 
+  /// Duplicate blocks with unconditional branches into tails of their
+  /// predecessors. Variant that works before register allocation.
+  extern char &EarlyTailDuplicateID;
+
   /// MachineTraceMetrics - This pass computes critical path and CPU resource
   /// usage in an ensemble of traces.
   extern char &MachineTraceMetricsID;
diff --git a/llvm/include/llvm/CodeGen/TargetPassConfig.h b/llvm/include/llvm/CodeGen/TargetPassConfig.h
index c6987327c27..90480fbbbf2 100644
--- a/llvm/include/llvm/CodeGen/TargetPassConfig.h
+++ b/llvm/include/llvm/CodeGen/TargetPassConfig.h
@@ -90,10 +90,6 @@ public:
   /// TargetPassConfig APIs to identify multiple occurrences of the same pass.
   ///
 
-  /// EarlyTailDuplicate - A clone of the TailDuplicate pass that runs early
-  /// during codegen, on SSA form.
-  static char EarlyTailDuplicateID;
-
   /// PostRAMachineLICM - A clone of the LICM pass that runs during late machine
   /// optimization after regalloc.
   static char PostRAMachineLICMID;
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index dd7aa722ed2..56271fcd2f7 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -99,6 +99,7 @@ void initializeConstantMergeLegacyPassPass(PassRegistry&);
 void initializeConstantPropagationPass(PassRegistry&);
 void initializeCorrelatedValuePropagationPass(PassRegistry&);
 void initializeCostModelAnalysisPass(PassRegistry&);
+void initializeEarlyTailDuplicatePass(PassRegistry&);
 void initializeEntryExitInstrumenterPass(PassRegistry&);
 void initializePostInlineEntryExitInstrumenterPass(PassRegistry&);
 void initializeCrossDSOCFIPass(PassRegistry&);
@@ -362,7 +363,7 @@ void initializeStripSymbolsPass(PassRegistry&);
 void initializeStructurizeCFGPass(PassRegistry&);
 void initializeHWAddressSanitizerPass(PassRegistry&);
 void initializeTailCallElimPass(PassRegistry&);
-void initializeTailDuplicatePassPass(PassRegistry&);
+void initializeTailDuplicatePass(PassRegistry&);
 void initializeTargetLibraryInfoWrapperPassPass(PassRegistry&);
 void initializeTargetPassConfigPass(PassRegistry&);
 void initializeTargetTransformInfoWrapperPassPass(PassRegistry&);
diff --git a/llvm/lib/CodeGen/CodeGen.cpp b/llvm/lib/CodeGen/CodeGen.cpp
index c0d7eb4cf47..40f1599a2fa 100644
--- a/llvm/lib/CodeGen/CodeGen.cpp
+++ b/llvm/lib/CodeGen/CodeGen.cpp
@@ -28,6 +28,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeDetectDeadLanesPass(Registry);
   initializeDwarfEHPreparePass(Registry);
   initializeEarlyIfConverterPass(Registry);
+  initializeEarlyTailDuplicatePass(Registry);
   initializeExpandISelPseudosPass(Registry);
   initializeExpandMemCmpPassPass(Registry);
   initializeExpandPostRAPass(Registry);
@@ -89,7 +90,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeStackMapLivenessPass(Registry);
   initializeStackProtectorPass(Registry);
   initializeStackSlotColoringPass(Registry);
-  initializeTailDuplicatePassPass(Registry);
+  initializeTailDuplicatePass(Registry);
   initializeTargetPassConfigPass(Registry);
   initializeTwoAddressInstructionPassPass(Registry);
   initializeUnpackMachineBundlesPass(Registry);
diff --git a/llvm/lib/CodeGen/TailDuplication.cpp b/llvm/lib/CodeGen/TailDuplication.cpp
index df1eebf43b2..25cd7802264 100644
--- a/llvm/lib/CodeGen/TailDuplication.cpp
+++ b/llvm/lib/CodeGen/TailDuplication.cpp
@@ -7,8 +7,9 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This pass duplicates basic blocks ending in unconditional branches into
-// the tails of their predecessors, using the TailDuplicator utility class.
+/// \file This pass duplicates basic blocks ending in unconditional branches
+/// into the tails of their predecessors, using the TailDuplicator utility
+/// class.
 //
 //===----------------------------------------------------------------------===//
 
@@ -26,38 +27,55 @@ using namespace llvm;
 
 namespace {
 
-/// Perform tail duplication. Delegates to TailDuplicator
-class TailDuplicatePass : public MachineFunctionPass {
+class TailDuplicateBase : public MachineFunctionPass {
   TailDuplicator Duplicator;
-
+  bool PreRegAlloc;
 public:
-  static char ID;
-
-  explicit TailDuplicatePass() : MachineFunctionPass(ID) {}
+  TailDuplicateBase(char &PassID, bool PreRegAlloc)
+    : MachineFunctionPass(PassID), PreRegAlloc(PreRegAlloc) {}
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineBranchProbabilityInfo>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+class TailDuplicate : public TailDuplicateBase {
+public:
+  static char ID;
+  TailDuplicate() : TailDuplicateBase(ID, false) {
+    initializeTailDuplicatePass(*PassRegistry::getPassRegistry());
+  }
+};
+
+class EarlyTailDuplicate : public TailDuplicateBase {
+public:
+  static char ID;
+  EarlyTailDuplicate() : TailDuplicateBase(ID, true) {
+    initializeEarlyTailDuplicatePass(*PassRegistry::getPassRegistry());
+  }
 };
 
 } // end anonymous namespace
 
-char TailDuplicatePass::ID = 0;
+char TailDuplicate::ID;
+char EarlyTailDuplicate::ID;
 
-char &llvm::TailDuplicateID = TailDuplicatePass::ID;
+char &llvm::TailDuplicateID = TailDuplicate::ID;
+char &llvm::EarlyTailDuplicateID = EarlyTailDuplicate::ID;
 
-INITIALIZE_PASS(TailDuplicatePass, DEBUG_TYPE, "Tail Duplication", false, false)
+INITIALIZE_PASS(TailDuplicate, DEBUG_TYPE, "Tail Duplication", false, false)
+INITIALIZE_PASS(EarlyTailDuplicate, "early-tailduplication",
+                "Early Tail Duplication", false, false)
 
-bool TailDuplicatePass::runOnMachineFunction(MachineFunction &MF) {
+bool TailDuplicateBase::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   auto MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
-
-  // TODO: Querying isSSA() to determine pre-/post-regalloc is fragile, better
-  // split this into two passes instead.
-  bool PreRegAlloc = MF.getRegInfo().isSSA();
-  Duplicator.initMF(MF, PreRegAlloc, MBPI, /* LayoutMode */ false);
+  Duplicator.initMF(MF, PreRegAlloc, MBPI, /*LayoutMode=*/false);
 
   bool MadeChange = false;
   while (Duplicator.tailDuplicateBlocks())
@@ -65,8 +83,3 @@ bool TailDuplicatePass::runOnMachineFunction(MachineFunction &MF) {
 
   return MadeChange;
 }
-
-void TailDuplicatePass::getAnalysisUsage(AnalysisUsage &AU) const {
-  AU.addRequired<MachineBranchProbabilityInfo>();
-  MachineFunctionPass::getAnalysisUsage(AU);
-}
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
index 392f41ac438..e59ada53663 100644
--- a/llvm/lib/CodeGen/TargetPassConfig.cpp
+++ b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -226,7 +226,7 @@ static IdentifyingPassPtr overridePass(AnalysisID StandardID,
   if (StandardID == &TailDuplicateID)
     return applyDisable(TargetID, DisableTailDuplicate);
 
-  if (StandardID == &TargetPassConfig::EarlyTailDuplicateID)
+  if (StandardID == &EarlyTailDuplicateID)
     return applyDisable(TargetID, DisableEarlyTailDup);
 
   if (StandardID == &MachineBlockPlacementID)
@@ -268,7 +268,6 @@ INITIALIZE_PASS(TargetPassConfig, "targetpassconfig",
 char TargetPassConfig::ID = 0;
 
 // Pseudo Pass IDs.
-char TargetPassConfig::EarlyTailDuplicateID = 0;
 char TargetPassConfig::PostRAMachineLICMID = 0;
 
 namespace {
@@ -367,7 +366,6 @@ TargetPassConfig::TargetPassConfig(LLVMTargetMachine &TM, PassManagerBase &pm)
   initializeAAResultsWrapperPassPass(*PassRegistry::getPassRegistry());
 
   // Substitute Pseudo Pass IDs for real ones.
-  substitutePass(&EarlyTailDuplicateID, &TailDuplicateID);
   substitutePass(&PostRAMachineLICMID, &MachineLICMID);
 
   if (StringRef(PrintMachineInstrs.getValue()).equals(""))
diff --git a/llvm/test/CodeGen/X86/tail-dup-debugloc.ll b/llvm/test/CodeGen/X86/tail-dup-debugloc.ll
index 4907e5244b6..df1a8ee46b4 100644
--- a/llvm/test/CodeGen/X86/tail-dup-debugloc.ll
+++ b/llvm/test/CodeGen/X86/tail-dup-debugloc.ll
@@ -1,4 +1,4 @@
-; RUN: llc -stop-after=tailduplication < %s | FileCheck %s
+; RUN: llc -stop-after=early-tailduplication < %s | FileCheck %s
 ;
 ; Check that DebugLoc attached to the branch instruction of
 ; 'while.cond1.preheader.lr.ph' survives after tailduplication pass.