[CodeExtractor] Remove stale llvm.assume calls from extracted region

During extraction, stale llvm.assume handles may be retained in the
original function. The setup is:

1) CodeExtractor unregisters assumptions in the blocks that are to be
   extracted.

2) Extraction happens. There are now two functions: f1 and f1.extracted.

3) Leftover assumptions in f1 (/not/ removed as they were not in the set of
   blocks to be extracted) now have affected-value llvm.assume handles in
   f1.extracted.

When assumptions for a value used in f1 are looked up, ValueTracking can assert
as some of the handles are in the wrong function. To fix this, simply erase the
llvm.assume calls in the extracted function.

Alternatives include flushing the assumption cache in the original function, or
walking all values used in the original function to prune stale affected-value
handles. Both seem more expensive.

Testing: check-llvm, LNT run with -mllvm -hot-cold-split enabled

rdar://58460728

include/llvm/Transforms/Utils/CodeExtractor.h

@@ -140,9 +140,11 @@ public:
     Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC);
 
     /// Verify that assumption cache isn't stale after a region is extracted.
-    /// Returns false when verifier finds errors. AssumptionCache is passed as
+    /// Returns true when verifier finds errors. AssumptionCache is passed as
     /// parameter to make this function stateless.
-    static bool verifyAssumptionCache(const Function& F, AssumptionCache *AC);
+    static bool verifyAssumptionCache(const Function &OldFunc,
+                                      const Function &NewFunc,
+                                      AssumptionCache *AC);
 
     /// Test whether this code extractor is eligible.
     ///

lib/Transforms/Utils/CodeExtractor.cpp

@@ -1531,13 +1531,19 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC) {
     }
   }
 
-  if (AC) {
-    // Remove @llvm.assume calls that were moved to the new function from the
+  // Remove @llvm.assume calls that will be 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));
+  for (BasicBlock *Block : Blocks) {
+    for (auto It = Block->begin(), End = Block->end(); It != End;) {
+      Instruction *I = &*It;
+      ++It;
+
+      if (match(I, m_Intrinsic<Intrinsic::assume>())) {
+        if (AC)
+          AC->unregisterAssumption(cast<CallInst>(I));
+        I->eraseFromParent();
+      }
+    }
   }
 
   // If we have any return instructions in the region, split those blocks so
@@ -1711,17 +1717,36 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC) {
   });
   LLVM_DEBUG(if (verifyFunction(*oldFunction))
              report_fatal_error("verification of oldFunction failed!"));
-  LLVM_DEBUG(if (AC && verifyAssumptionCache(*oldFunction, AC))
+  LLVM_DEBUG(if (AC && verifyAssumptionCache(*oldFunction, *newFunction, AC))
                  report_fatal_error("Stale Asumption cache for old Function!"));
   return newFunction;
 }
 
-bool CodeExtractor::verifyAssumptionCache(const Function& F,
+bool CodeExtractor::verifyAssumptionCache(const Function &OldFunc,
+                                          const Function &NewFunc,
                                           AssumptionCache *AC) {
   for (auto AssumeVH : AC->assumptions()) {
-    CallInst *I = cast<CallInst>(AssumeVH);
-    if (I->getFunction() != &F)
+    CallInst *I = dyn_cast_or_null<CallInst>(AssumeVH);
+    if (!I)
+      continue;
+
+    // There shouldn't be any llvm.assume intrinsics in the new function.
+    if (I->getFunction() != &OldFunc)
       return true;
+
+    // There shouldn't be any stale affected values in the assumption cache
+    // that were previously in the old function, but that have now been moved
+    // to the new function.
+    for (auto AffectedValVH : AC->assumptionsFor(I->getOperand(0))) {
+      CallInst *AffectedCI = dyn_cast_or_null<CallInst>(AffectedValVH);
+      if (!AffectedCI)
+        continue;
+      if (AffectedCI->getFunction() != &OldFunc)
+        return true;
+      auto *AssumedInst = dyn_cast<Instruction>(AffectedCI->getOperand(0));
+      if (AssumedInst->getFunction() != &OldFunc)
+        return true;
+    }
   }
   return false;
 }

test/Transforms/CodeExtractor/extract-assume.ll

@@ -5,7 +5,7 @@
 
 ; CHECK: Cached assumptions for function: fun
 ; CHECK-NEXT: Cached assumptions for function: fun.cold
-; CHECK-NEXT:   %cmp = icmp uge i32 %x, 64
+; CHECK-NOT: icmp uge
 
 declare void @fun2(i32) #0
 

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

@@ -16,7 +16,7 @@ target triple = "aarch64"
 ; CHECK: define {{.*}}@f.cold.1(i64 %0)
 ; CHECK-LABEL: newFuncRoot:
 ; CHECK: %1 = icmp eq i64 %0, 0
-; CHECK: call void @llvm.assume(i1 %1)
+; CHECK-NOT: call void @llvm.assume
 
 define void @f() {
 entry:

test/Transforms/HotColdSplit/stale-assume-in-original-func.ll

@@ -0,0 +1,38 @@
+; RUN: opt -S -hotcoldsplit -hotcoldsplit-threshold=-1 < %s 2>&1 | FileCheck %s
+
+; CHECK-LABEL: define {{.*}} @foo(
+; CHECK-NOT: llvm.assume
+; CHECK: call void @foo.cold.1()
+; CHECK: llvm.assume
+; CHECK-NEXT: ret void
+
+; CHECK-LABEL: define {{.*}} @foo.cold.1(
+; CHECK-NOT: llvm.assume
+; CHECK: call void @cold()
+; CHECK-NOT: llvm.assume
+; CHECK: }
+
+define void @foo(i1 %cond) {
+entry:
+  br i1 %cond, label %cold, label %cont
+
+cold:
+  call void @llvm.assume(i1 %cond)
+  call void @cold()
+  br label %cont
+
+cont:
+  %cmp = icmp eq i1 %cond, true
+  br i1 %cmp, label %exit1, label %exit2
+
+exit1:
+  call void @llvm.assume(i1 %cond)
+  ret void
+
+exit2:
+  ret void
+}
+
+declare void @llvm.assume(i1)
+
+declare void @cold() cold

unittests/Transforms/Utils/CodeExtractorTest.cpp

@@ -280,6 +280,6 @@ TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
-  EXPECT_FALSE(CE.verifyAssumptionCache(*Func, &AC));
+  EXPECT_FALSE(CE.verifyAssumptionCache(*Func, *Outlined, &AC));
 }
 } // end anonymous namespace