Index: /Users/rriddle/Desktop/llvm/llvm/include/llvm/Transforms/IPO/PartialInlining.h
===================================================================
--- /Users/rriddle/Desktop/llvm/llvm/include/llvm/Transforms/IPO/PartialInlining.h
+++ /Users/rriddle/Desktop/llvm/llvm/include/llvm/Transforms/IPO/PartialInlining.h
@@ -24,9 +24,6 @@
 class PartialInlinerPass : public PassInfoMixin<PartialInlinerPass> {
 public:
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
-
-private:
-  Function *unswitchFunction(Function *F);
 };
 }
 #endif // LLVM_TRANSFORMS_IPO_PARTIALINLINING_H
Index: /Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp
===================================================================
--- /Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp
+++ /Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/PartialInlining.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
@@ -29,40 +30,71 @@
 STATISTIC(NumPartialInlined, "Number of functions partially inlined");
 
 namespace {
-struct PartialInlinerLegacyPass : public ModulePass {
-  static char ID; // Pass identification, replacement for typeid
-  PartialInlinerLegacyPass() : ModulePass(ID) {
-    initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
-  }
-
-  bool runOnModule(Module &M) override {
-    if (skipModule(M))
-      return false;
-    ModuleAnalysisManager DummyMAM;
-    auto PA = Impl.run(M, DummyMAM);
-    return !PA.areAllPreserved();
-  }
+  struct PartialInlinerImpl{
+    PartialInlinerImpl(InlineFunctionInfo IFI) : IFI(IFI){}
+    bool run(Module &M);
+    Function *unswitchFunction(Function *F);
+  private:
+    InlineFunctionInfo IFI;
+  };
+  struct PartialInlinerLegacyPass : public ModulePass {
+    static char ID; // Pass identification, replacement for typeid
+    PartialInlinerLegacyPass() : ModulePass(ID) {
+      initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
+    }
 
-private:
-  PartialInlinerPass Impl;
+    void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<AssumptionCacheTracker>();
+    }
+    bool runOnModule(Module &M) override {
+      if (skipModule(M))
+        return false;
+      InlineFunctionInfo IFI;
+      IFI.ACT = &getAnalysis<AssumptionCacheTracker>();
+      return PartialInlinerImpl(IFI).run(M);
+    }
   };
 }
 
-char PartialInlinerLegacyPass::ID = 0;
-INITIALIZE_PASS(PartialInlinerLegacyPass, "partial-inliner", "Partial Inliner",
-                false, false)
+bool PartialInlinerImpl::run(Module &M) {
+  std::vector<Function*> worklist;
+  worklist.reserve(M.size());
+  for (Function &F : M)
+    if (!F.use_empty() && !F.isDeclaration())
+      worklist.push_back(&F);
 
-ModulePass *llvm::createPartialInliningPass() {
-  return new PartialInlinerLegacyPass();
+  bool changed = false;
+  while (!worklist.empty()) {
+    Function* currFunc = worklist.back();
+    worklist.pop_back();
+
+    if (currFunc->use_empty()) continue;
+
+    bool recursive = false;
+    for (User *U : currFunc->users())
+      if (Instruction* I = dyn_cast<Instruction>(U))
+      if (I->getParent()->getParent() == currFunc) {
+        recursive = true;
+        break;
+      }
+    if (recursive) continue;
+
+    if (Function* newFunc = unswitchFunction(currFunc)) {
+      worklist.push_back(newFunc);
+      changed = true;
+    }
+
+  }
+  return changed;
 }
 
-Function *PartialInlinerPass::unswitchFunction(Function *F) {
+Function *PartialInlinerImpl::unswitchFunction(Function *F) {
   // First, verify that this function is an unswitching candidate...
   BasicBlock *entryBlock = &F->front();
   BranchInst *BR = dyn_cast<BranchInst>(entryBlock->getTerminator());
   if (!BR || BR->isUnconditional())
     return nullptr;
-  
+
   BasicBlock* returnBlock = nullptr;
   BasicBlock* nonReturnBlock = nullptr;
   unsigned returnCount = 0;
@@ -73,10 +105,10 @@
     } else
       nonReturnBlock = BB;
   }
-  
+
   if (returnCount != 1)
     return nullptr;
-  
+
   // Clone the function, so that we can hack away on it.
   ValueToValueMapTy VMap;
   Function* duplicateFunction = CloneFunction(F, VMap);
@@ -84,11 +116,11 @@
   BasicBlock* newEntryBlock = cast<BasicBlock>(VMap[entryBlock]);
   BasicBlock* newReturnBlock = cast<BasicBlock>(VMap[returnBlock]);
   BasicBlock* newNonReturnBlock = cast<BasicBlock>(VMap[nonReturnBlock]);
-  
+
   // Go ahead and update all uses to the duplicate, so that we can just
   // use the inliner functionality when we're done hacking.
   F->replaceAllUsesWith(duplicateFunction);
-  
+
   // Special hackery is needed with PHI nodes that have inputs from more than
   // one extracted block.  For simplicity, just split the PHIs into a two-level
   // sequence of PHIs, some of which will go in the extracted region, and some
@@ -110,11 +142,11 @@
     retPhi->addIncoming(OldPhi->getIncomingValueForBlock(newEntryBlock),
                         newEntryBlock);
     OldPhi->removeIncomingValue(newEntryBlock);
-    
+
     ++I;
   }
   newEntryBlock->getTerminator()->replaceUsesOfWith(preReturn, newReturnBlock);
-  
+
   // Gather up the blocks that we're going to extract.
   std::vector<BasicBlock*> toExtract;
   toExtract.push_back(newNonReturnBlock);
@@ -129,10 +161,8 @@
 
   // Extract the body of the if.
   Function* extractedFunction
-    = CodeExtractor(toExtract, &DT).extractCodeRegion();
-  
-  InlineFunctionInfo IFI;
-  
+      = CodeExtractor(toExtract, &DT).extractCodeRegion();
+
   // Inline the top-level if test into all callers.
   std::vector<User *> Users(duplicateFunction->user_begin(),
                             duplicateFunction->user_end());
@@ -141,49 +171,32 @@
       InlineFunction(CI, IFI);
     else if (InvokeInst *II = dyn_cast<InvokeInst>(User))
       InlineFunction(II, IFI);
-  
+
   // Ditch the duplicate, since we're done with it, and rewrite all remaining
   // users (function pointers, etc.) back to the original function.
   duplicateFunction->replaceAllUsesWith(F);
   duplicateFunction->eraseFromParent();
-  
+
   ++NumPartialInlined;
-  
+
   return extractedFunction;
 }
 
-PreservedAnalyses PartialInlinerPass::run(Module &M, ModuleAnalysisManager &) {
-  std::vector<Function*> worklist;
-  worklist.reserve(M.size());
-  for (Function &F : M)
-    if (!F.use_empty() && !F.isDeclaration())
-      worklist.push_back(&F);
 
-  bool changed = false;
-  while (!worklist.empty()) {
-    Function* currFunc = worklist.back();
-    worklist.pop_back();
-  
-    if (currFunc->use_empty()) continue;
-    
-    bool recursive = false;
-    for (User *U : currFunc->users())
-      if (Instruction* I = dyn_cast<Instruction>(U))
-        if (I->getParent()->getParent() == currFunc) {
-          recursive = true;
-          break;
-        }
-    if (recursive) continue;
-          
-    
-    if (Function* newFunc = unswitchFunction(currFunc)) {
-      worklist.push_back(newFunc);
-      changed = true;
-    }
-    
-  }
+char PartialInlinerLegacyPass::ID = 0;
+INITIALIZE_PASS_BEGIN(PartialInlinerLegacyPass,  "partial-inliner", "Partial Inliner",
+                      false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_END(PartialInlinerLegacyPass,  "partial-inliner", "Partial Inliner",
+                    false, false)
+
+ModulePass *llvm::createPartialInliningPass() {
+  return new PartialInlinerLegacyPass();
+}
 
-  if (changed)
+PreservedAnalyses PartialInlinerPass::run(Module &M, ModuleAnalysisManager &AM) {
+  InlineFunctionInfo IFI;
+  if (PartialInlinerImpl(IFI).run(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
-}
+}
\ No newline at end of file