Invalidate assumption cache before outlining.

[hiraditya]

Reviewed here: https://reviews.llvm.org/D67941

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
                   std::string Suffix = "");
 
     /// Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
     Function *extractCodeRegion();
 
+    /// Verify that assumption cache isn't stale after a region is extracted.
+    /// Returns false when verifier finds errors. AssumptionCache is passed as
+    /// parameter to make this function stateless.
+    static bool verifyAssumptionCache(const Function& F, AssumptionCache *AC);
+
     /// Test whether this code extractor is eligible.
     ///
     /// Based on the blocks used when constructing the code extractor,
     /// determine whether it is eligible for extraction.
     /// 
     /// Checks that varargs handling (with vastart and vaend) is only done in
     /// the outlined blocks.
     bool isEligible() const;

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

do {
                                              OutlinedFunctionID);
       if (Outlined) {
         ++OutlinedFunctionID;
         Changed = true;
       }
     } while (!Region.empty());
   }
 
-  // We need to explicitly clear the assumption cache since the value tracking
-  // may now be invalid as part of the function has changed.
-  if (Changed)
-    if (AssumptionCache *AC = LookupAC(F))
-      AC->clear();
   return Changed;
 }
 
 bool HotColdSplitting::run(Module &M) {
   bool Changed = false;
   bool HasProfileSummary = (M.getProfileSummary(/* IsCS */ false) != nullptr);
   for (auto It = M.begin(), End = M.end(); It != End; ++It) {
     Function &F = *It;

llvm/lib/Transforms/Utils/CodeExtractor.cpp

void CodeExtractor::moveCodeToFunction(Function *newFunction) {
   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
 
   for (BasicBlock *Block : Blocks) {
     // Delete the basic block from the old function, and the list of blocks
     oldBlocks.remove(Block);
 
     // Insert this basic block into the new function
     newBlocks.push_back(Block);
-
-    // Remove @llvm.assume calls that were moved to the new function from the
-    // old function's assumption cache.
-    if (AC)
-      for (auto &I : *Block)
-        if (match(&I, m_Intrinsic<Intrinsic::assume>()))
-          AC->unregisterAssumption(cast<CallInst>(&I));
   }
 }
 
 void CodeExtractor::calculateNewCallTerminatorWeights(
     BasicBlock *CodeReplacer,
     DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
     BranchProbabilityInfo *BPI) {
   using Distribution = BlockFrequencyInfoImplBase::Distribution;
if (BFI) {
     for (BasicBlock *Pred : predecessors(header)) {
       if (Blocks.count(Pred))
         continue;
       EntryFreq +=
           BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, header);
     }
   }
 
+  if (AC) {
+    // Remove @llvm.assume calls that were moved to the new function from the
+    // old function's assumption cache.
+    for (BasicBlock *Block : Blocks)
+      for (auto &I : *Block)
+        if (match(&I, m_Intrinsic<Intrinsic::assume>()))
+          AC->unregisterAssumption(cast<CallInst>(&I));
+  }
+
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
   // Calculate the exit blocks for the extracted region and the total exit
   // weights for each of those blocks.
   DenseMap<BasicBlock *, BlockFrequency> ExitWeights;
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
if (doesNotReturn)
     newFunction->setDoesNotReturn();
 
   LLVM_DEBUG(if (verifyFunction(*newFunction, &errs())) {
     newFunction->dump();
     report_fatal_error("verification of newFunction failed!");
   });
   LLVM_DEBUG(if (verifyFunction(*oldFunction))
              report_fatal_error("verification of oldFunction failed!"));
+  LLVM_DEBUG(if (AC && verifyAssumptionCache(*oldFunction, AC))
+             report_fatal_error("Stale Asumption cache for old Function!"));
   return newFunction;
 }
+
+bool CodeExtractor::verifyAssumptionCache(const Function& F,
+                                          AssumptionCache *AC) {
+  for (auto AssumeVH : AC->assumptions()) {
+    CallInst *I = cast<CallInst>(AssumeVH);
+    if (I->getFunction() != &F)
+      return true;
+  }
+  return false;
+}

llvm/test/Transforms/HotColdSplit/assumption-cache-invalidation.ll

+; REQUIRES: asserts
+; RUN: opt -S -instsimplify -hotcoldsplit -debug < %s 2>&1 | FileCheck %s
 ; RUN: opt -instcombine -hotcoldsplit -instsimplify %s -o /dev/null
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64"
 
 %a = type { i64, i64 }
 %b = type { i64 }
 
+; CHECK: @f
+; CHECK-LABEL: codeRepl:
+; CHECK-NOT: @llvm.assume
+; CHECK: }
+; CHECK: @llvm.assume
+; CHECK: @f.cold.1(i64 %0)
+; CHECK-LABEL: newFuncRoot:
+; CHECK: %1 = icmp eq i64 %0, 0
+; CHECK: call void @llvm.assume(i1 %1)
+
 define void @f() {
 entry:
   %0 = getelementptr inbounds %a, %a* null, i64 0, i32 1
   br label %label
 
 label:
   %1 = bitcast i64* %0 to %b**
   %2 = load %b*, %b** %1, align 8