Invalidate assumption cache before outlining.

Subscribers: llvm-commits

Tags: #llvm

Reviewers: compnerd, vsk, sebpop, fhahn, tejohnson

Reviewed by: vsk

Differential Revision: https://reviews.llvm.org/D68478

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@373807 91177308-0d34-0410-b5e6-96231b3b80d8
(cherry picked from commit 83e4df0)

Author: hiraditya

include/llvm/Transforms/Utils/CodeExtractor.h

@@ -106,6 +106,11 @@ class Value;
     /// 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,

lib/Transforms/Utils/CodeExtractor.cpp

@@ -1297,13 +1297,6 @@ void CodeExtractor::moveCodeToFunction(Function *newFunction) {
 
     // 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));
   }
 }
 
@@ -1374,6 +1367,15 @@ Function *CodeExtractor::extractCodeRegion() {
     }
   }
 
+  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();
@@ -1564,5 +1566,17 @@ Function *CodeExtractor::extractCodeRegion() {
   });
   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;
+}

unittests/Transforms/Utils/CodeExtractorTest.cpp

@@ -236,7 +236,7 @@ TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
         %b = type { i64 }
         declare void @g(i8*)
 
-        declare void @llvm.assume(i1) #0
+        declare void @llvm.assume(i1)
 
         define void @test() {
         entry:
@@ -259,8 +259,6 @@ TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
           call void @llvm.assume(i1 %4)
           unreachable
         }
-
-        attributes #0 = { nounwind willreturn }
   )ir",
                                                 Err, Ctx));