Fix : Partial Inliner requires AssumptionCacheTracker

[rriddle]

Summary

The public InlineFunction utility assumes that the passed in InlineFunctionInfo has a valid AssumptionCacheTracker.

[davide]

You mean, port to the new pass manager?

[rriddle]

It was already ported to the new pass manager, but given that InlineFunction directly accesses the ACT it will crash if run via legacy.

[eraman]

Please split the adding of the dependence on AssumptionCache and the porting of the pass into the new PM into two separate passes. Also, note that you are not populating the ACT field of IFI in the new PM version of the code and that needs to be fixed.

[rriddle]

Sorry, I was under the impression that the new pass managed had a functor for the ACT. I guess that was a pass that hadn't landed yet.

[eraman]

Sorry, I looked at Davide's previous comment and quickly made an incorrect conclusion. So this does only add the dependence to ACT, but you still need to populate ACT with the new PM.

[rriddle]

Also, the pass was already ported to the new pass manager, this just moves the code into a different structure.

[davide]

Oh, sigh, my bad. I was confused by the churn. Ignore my comment.

[rriddle]

Explicit dependency for LegacyPass

[rriddle]

The new pass manager doesn't use AssumptionCacheTracker, providing a replacement instead with AssumptionAnalysis.

[silvas]

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp

52–53

Can you clang-format this class?

205

You aren't passing assumption cache analysis stuff in here. This will require some refactoring of stuff that uses InlineFunctionInfo. I've already done that work as part of https://reviews.llvm.org/D21921. Give me a second and I'll spin that out and commit that. Then it should be easy to pass the assumption cache in here.

[silvas]

Also, can you add a test case?

[rriddle]

Fix the formatting of the implementation and added a test case for the assumption cache.

[silvas]

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp

205

Should be done in r276515. Sorry for the delay.

[rriddle]

Updated to reflect r276515

[silvas]

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp

123

Can you split these formatting changes into a separate patch after this one? Even better, formatting + updating the naming would be a very welcome patch for this pass.

/Users/rriddle/Desktop/llvm/llvm/test/Transforms/Inline/partial-inline-act.ll

4

Can you please put explicit names on the arguments (even just %0 and %1)?

[rriddle]

Updated from Sean's comments

[silvas]

r276609, thanks!

/Users/rriddle/Desktop/llvm/llvm/include/llvm/Transforms/IPO/PartialInlining.h

 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// Pass to remove unused function declarations.
 class PartialInlinerPass : public PassInfoMixin<PartialInlinerPass> {
 public:
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
-
-private:
-  Function *unswitchFunction(Function *F);
 };
 }
 #endif // LLVM_TRANSFORMS_IPO_PARTIALINLINING_H

/Users/rriddle/Desktop/llvm/llvm/lib/Transforms/IPO/PartialInlining.cpp

 //===- PartialInlining.cpp - Inline parts of functions --------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass performs partial inlining, typically by inlining an if statement
 // that surrounds the body of the function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "partialinlining"
 
 STATISTIC(NumPartialInlined, "Number of functions partially inlined");
 
 namespace {
+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());
   }
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
+  }
   bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
-    ModuleAnalysisManager DummyMAM;
-    auto PA = Impl.run(M, DummyMAM);
-    return !PA.areAllPreserved();
-  }
 
-private:
-  PartialInlinerPass Impl;
-  };
+    AssumptionCacheTracker *ACT = &getAnalysis<AssumptionCacheTracker>();
+    std::function<AssumptionCache &(Function &)> GetAssumptionCache =
+        [&ACT](Function &F) -> AssumptionCache & {
+      return ACT->getAssumptionCache(F);
+    };
+    InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
+    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;
 
-Function *PartialInlinerPass::unswitchFunction(Function *F) {
+    if (Function *newFunc = unswitchFunction(currFunc)) {
+      worklist.push_back(newFunc);
+      changed = true;
+    }
+  }
+  return changed;
+}
+
+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;
+  BasicBlock *returnBlock = nullptr;
+  BasicBlock *nonReturnBlock = nullptr;
   unsigned returnCount = 0;
   for (BasicBlock *BB : successors(entryBlock)) {
     if (isa<ReturnInst>(BB->getTerminator())) {
       returnBlock = BB;
       returnCount++;
     } 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);
+  Function *duplicateFunction = CloneFunction(F, VMap);
   duplicateFunction->setLinkage(GlobalValue::InternalLinkage);
-  BasicBlock* newEntryBlock = cast<BasicBlock>(VMap[entryBlock]);
-  BasicBlock* newReturnBlock = cast<BasicBlock>(VMap[returnBlock]);
-  BasicBlock* newNonReturnBlock = cast<BasicBlock>(VMap[nonReturnBlock]);
-  
+  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
   // of which will go outside.
-  BasicBlock* preReturn = newReturnBlock;
+  BasicBlock *preReturn = newReturnBlock;
   newReturnBlock = newReturnBlock->splitBasicBlock(
       newReturnBlock->getFirstNonPHI()->getIterator());
   BasicBlock::iterator I = preReturn->begin();
   Instruction *Ins = &newReturnBlock->front();
   while (I != preReturn->end()) {
-    PHINode* OldPhi = dyn_cast<PHINode>(I);
-    if (!OldPhi) break;
+    PHINode *OldPhi = dyn_cast<PHINode>(I);
+    if (!OldPhi)
+      break;
 
     PHINode *retPhi = PHINode::Create(OldPhi->getType(), 2, "", Ins);
     OldPhi->replaceAllUsesWith(retPhi);
     Ins = newReturnBlock->getFirstNonPHI();
 
     retPhi->addIncoming(&*I, preReturn);
     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;
+  std::vector<BasicBlock *> toExtract;
   toExtract.push_back(newNonReturnBlock);
   for (BasicBlock &BB : *duplicateFunction)
     if (&BB != newEntryBlock && &BB != newReturnBlock &&
         &BB != newNonReturnBlock)
       toExtract.push_back(&BB);
 
   // The CodeExtractor needs a dominator tree.
   DominatorTree DT;
   DT.recalculate(*duplicateFunction);
 
   // Extract the body of the if.
-  Function* extractedFunction
-    = CodeExtractor(toExtract, &DT).extractCodeRegion();
-  
-  InlineFunctionInfo IFI;
-  
+  Function *extractedFunction =
+      CodeExtractor(toExtract, &DT).extractCodeRegion();
   // Inline the top-level if test into all callers.
   std::vector<User *> Users(duplicateFunction->user_begin(),
                             duplicateFunction->user_end());
   for (User *User : Users)
     if (CallInst *CI = dyn_cast<CallInst>(User))
       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) {
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  std::function<AssumptionCache &(Function &)> GetAssumptionCache =
+      [&FAM](Function &F) -> AssumptionCache & {
+    return FAM.getResult<AssumptionAnalysis>(F);
+  };
+  InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
+  if (PartialInlinerImpl(IFI).run(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
-}
+}
 No newline at end of file

/Users/rriddle/Desktop/llvm/llvm/test/Transforms/Inline/partial-inline-act.ll

+; RUN: opt < %s -partial-inliner -disable-output
+; This testcase tests the assumption cache
+
+define internal i32 @inlinedFunc(i1, i32* align 4) {
+entry:
+  br i1 %0, label %if.then, label %return
+if.then:
+  ; Dummy store to have more than 0 uses
+  store i32 10, i32* %1, align 4
+  br label %return
+return:             ; preds = %entry
+  ret i32 0
+}
+
+define internal i32 @dummyCaller(i1, i32* align 2) {
+entry:
+  %val = call i32 @inlinedFunc(i1 %0, i32* %1)
+  ret i32 %val
+}
+