2 files changed, 50 insertions, 22 deletions

diff --git a/llvm/lib/CodeGen/Passes.cpp b/llvm/lib/CodeGen/Passes.cpp
index 5b6e6f0441c..f36e5134b5a 100644
--- a/llvm/lib/CodeGen/Passes.cpp
+++ b/llvm/lib/CodeGen/Passes.cpp
@@ -53,9 +53,6 @@ static cl::opt<bool> DisableMachineLICM("disable-machine-licm", cl::Hidden,
     cl::desc("Disable Machine LICM"));
 static cl::opt<bool> DisableMachineCSE("disable-machine-cse", cl::Hidden,
     cl::desc("Disable Machine Common Subexpression Elimination"));
-static cl::opt<cl::boolOrDefault>
-    EnableShrinkWrapOpt("enable-shrink-wrap", cl::Hidden,
-                        cl::desc("enable the shrink-wrapping pass"));
 static cl::opt<cl::boolOrDefault> OptimizeRegAlloc(
     "optimize-regalloc", cl::Hidden,
     cl::desc("Enable optimized register allocation compilation path."));
@@ -218,7 +215,7 @@ TargetPassConfig::TargetPassConfig(TargetMachine *tm, PassManagerBase &pm)
     : ImmutablePass(ID), PM(&pm), StartBefore(nullptr), StartAfter(nullptr),
       StopAfter(nullptr), Started(true), Stopped(false),
       AddingMachinePasses(false), TM(tm), Impl(nullptr), Initialized(false),
-      DisableVerify(false), EnableTailMerge(true), EnableShrinkWrap(false) {
+      DisableVerify(false), EnableTailMerge(true) {
 
   Impl = new PassConfigImpl();
 
@@ -540,8 +537,8 @@ void TargetPassConfig::addMachinePasses() {
   addPostRegAlloc();
 
   // Insert prolog/epilog code.  Eliminate abstract frame index references...
-  if (getEnableShrinkWrap())
-    addPass(&ShrinkWrapID);
+  if (getOptLevel() != CodeGenOpt::None)
+    addPass(createShrinkWrapPass());
   addPass(&PrologEpilogCodeInserterID);
 
   /// Add passes that optimize machine instructions after register allocation.
@@ -620,21 +617,6 @@ void TargetPassConfig::addMachineSSAOptimization() {
   addPass(&DeadMachineInstructionElimID);
 }
 
-bool TargetPassConfig::getEnableShrinkWrap() const {
-  switch (EnableShrinkWrapOpt) {
-  case cl::BOU_UNSET:
-    return EnableShrinkWrap && getOptLevel() != CodeGenOpt::None;
-  // If EnableShrinkWrap is set, it takes precedence on whatever the
-  // target sets. The rational is that we assume we want to test
-  // something related to shrink-wrapping.
-  case cl::BOU_TRUE:
-    return true;
-  case cl::BOU_FALSE:
-    return false;
-  }
-  llvm_unreachable("Invalid shrink-wrapping state");
-}
-
 //===---------------------------------------------------------------------===//
 /// Register Allocation Pass Configuration
 //===---------------------------------------------------------------------===//
diff --git a/llvm/lib/CodeGen/ShrinkWrap.cpp b/llvm/lib/CodeGen/ShrinkWrap.cpp
index 4463cc7d3c5..f8ec1f358c4 100644
--- a/llvm/lib/CodeGen/ShrinkWrap.cpp
+++ b/llvm/lib/CodeGen/ShrinkWrap.cpp
@@ -68,11 +68,16 @@
 #include "llvm/Target/TargetInstrInfo.h"
 // To access TargetInstrInfo.
 #include "llvm/Target/TargetSubtargetInfo.h"
+#include "llvm/Support/CommandLine.h"
 
 #define DEBUG_TYPE "shrink-wrap"
 
 using namespace llvm;
 
+static cl::opt<cl::boolOrDefault>
+    EnableShrinkWrapOpt("enable-shrink-wrap", cl::Hidden,
+                        cl::desc("enable the shrink-wrapping pass"));
+
 STATISTIC(NumFunc, "Number of functions");
 STATISTIC(NumCandidates, "Number of shrink-wrapping candidates");
 STATISTIC(NumCandidatesDropped,
@@ -154,6 +159,11 @@ public:
   ShrinkWrap() : MachineFunctionPass(ID) {
     initializeShrinkWrapPass(*PassRegistry::getPassRegistry());
   }
+  
+  ShrinkWrap(std::function<bool(const MachineFunction &)> Ftor) :
+      MachineFunctionPass(ID), PredicateFtor(Ftor) {
+    initializeShrinkWrapPass(*PassRegistry::getPassRegistry());
+  }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
@@ -171,6 +181,15 @@ public:
   /// \brief Perform the shrink-wrapping analysis and update
   /// the MachineFrameInfo attached to \p MF with the results.
   bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  /// \brief Predicate function to determine if shrink wrapping should run.
+  ///
+  /// This function will be run at the beginning of shrink wrapping and
+  /// determine whether shrink wrapping should run on the given MachineFunction.
+  /// \param[in] MF The MachineFunction to run shrink wrapping on.
+  /// \return true if shrink wrapping should be run, false otherwise.
+  std::function<bool(const MachineFunction &MF)> PredicateFtor;
 };
 } // End anonymous namespace.
 
@@ -301,8 +320,12 @@ void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB) {
 }
 
 bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
-  if (MF.empty())
+  if (PredicateFtor && !PredicateFtor(MF)) 
     return false;
+  
+  if (MF.empty() || skipOptnoneFunction(*MF.getFunction()))
+    return false;
+
   DEBUG(dbgs() << "**** Analysing " << MF.getName() << '\n');
 
   init(MF);
@@ -386,3 +409,26 @@ bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
   ++NumCandidates;
   return false;
 }
+
+/// If EnableShrinkWrap is set run shrink wrapping on the given Machine
+/// Function. Otherwise, shrink wrapping is disabled.
+/// This function can be overridden in each target-specific TargetPassConfig
+/// class to allow different predicate logic for each target. 
+bool TargetPassConfig::runShrinkWrap(const MachineFunction &Fn) const {
+  switch (EnableShrinkWrapOpt) {
+  case cl::BOU_TRUE:
+    return true;
+  case cl::BOU_UNSET:
+  case cl::BOU_FALSE:
+    return false;
+  }
+  llvm_unreachable("Invalid shrink-wrapping state");
+}
+
+/// Create a ShrinkWrap FunctionPass using the runShrinkWrap predicate
+/// function.
+FunctionPass *TargetPassConfig::createShrinkWrapPass() {
+  std::function<bool(const MachineFunction &Fn)> Ftor =
+    std::bind(&TargetPassConfig::runShrinkWrap, this, std::placeholders::_1);
+  return new ShrinkWrap(Ftor);
+}