Invalidate assumption cache before outlining.

[hiraditya]

Summary

Assumption cache pointers become stale after function splitting and dependencies from old function can't be removed as such.
This patch:

1. un-registers the assumption cache before basic blocks are outlined
2. verifies the AssumptionCache of the original function to make sure no stale entries remain
3. Fixed a bug in assumption cache where the AssumeHandle wasn't removing an assume instrinsic while unregistering an Assumption

See D67936 for the quickfix to this bug.

[hfinkel]

Please post patches with full context (the full source file should be here). I can't see where moveCodeToFunction is called.

[vsk]

Is there a test case for the assertion failure? (I'm assuming that the test attached to D67936 does _not_, is that right?)

[hiraditya]

The test case in D67936 does assert, I mean without the fix from that diff.

[vsk]

Ah, could you add a test for this that asserts with D67936 applied?

[hiraditya]

With D67936, it wont assert anymore (with hot-cold-split). That was a quick fix but the real bug is in the CodeExtractor. let me add a revert of that patch in here as well. Because with this patch we wont need D67936.

[hiraditya]

Reverting r372667

[vsk]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1377

Is the findAffectedValues() operation inside of AssumptionCache not finding values because the llvm.assume is moved to another function? I.e. the search stops at an Argument? I think adding a unit test to CodeExtractorTest.cpp would really help.

[hiraditya]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1377

Basically the assertion: ValueTracking.cpp: 594 assert(I->getParent()->getParent() == Q.CxtI->getParent()->getParent() &&... triggers because one of the values is not unregistered even after it was moved to a new function. The value which are moved to new function get successfully unregistered without this patch but the values which remained in the original function didnt. This is because pointers to values got updated and Assumption cache had stale pointers.

Let me try to add a test case to CodeExtractorTest.cpp, seems slightly tricky from the first glance.

[hiraditya]

Added AssumptionCache verifier to verify the original function.

[hiraditya]

@vsk I tried adding test to CodeExtractorTest (D68095), for some reason the verifier didn't fail. I may be missing something and can use some help.
We can follow up on the test in the new diff if that is reasonable.

[hiraditya]

Is there any feedback for improvement, I'd like to push this if reasonable.

[fhahn]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1569

Can we just use AC->verifyAnalysis()?

[hiraditya]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1569

It seems, verifyAnalysis only checks if the AssumptionCache has the entry for a llvm.assume instruction or not.

AssumptionCache.cpp:268
        if (match(&II, m_Intrinsic<Intrinsic::assume>()) &&
            !AssumptionSet.count(cast<CallInst>(&II)))
          report_fatal_error("Assumption in scanned function not in cache");

[hiraditya]

@vsk Fixed the unittest in: D68095, it triggers the failure when AssumptionCache is stale.

[vsk]

I see, this looks good to me with a small test update.

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

3

Pass '-implicit-check-not=llvm.assume' to FileCheck, so this tests that the assumption in '@f' was removed?

[hiraditya]

added flag to the testcase.

[hiraditya]

Fixed a bug in assumption cache where the AssumeVH wasn't getting cleared.

[hiraditya]

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

3

the llvm.assume will be there in the new function after outlining.

[vsk]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1377

Could you explain the reason for the movement of logic that unregisters assumptions from moveCodeToFunction to here, before the new function is created? I don't understand why it's necessary.

1568

The if (AC) check needs to be wrapped in LLVM_DEBUG, otherwise the function will end without returning anything in Release mode.

1577

I->getFunction(), please.

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

3

Yes, I see that you've added a check for the llvm.assume in the new (split) function. Please add the -implicit-check-not as well, to test that no llvm.assume is left in the original function.

[hiraditya]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1377

I think we should fix assumption cache before moving the blocks to new function. Once we move the block to new function the instructions technically belong to the new function and assumption cache still treats them as if belonging to the old function. Even if clearing the AC after splitting works, it doesn't appear to be logically the right place to put it. I'm happy to put this where it was as I dont have a preference.

1568

It shouldn't because of ';', but i see the point.

1577

sure.

[hiraditya]

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

3

I'm not sure how to make -implicit-check-not work when llvm.assume is present in the function. I can use some help, can you share an example? For now I used CHECK to make sure llvm.assume isn't present in the original function.

[vsk]

Thanks, lgtm with some minor test updates.

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1377

Thanks for explaining, it makes sense to move the un-registration before splitting occurs.

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

3

Hm, my theory is that the -implicit-check-not is probably causing a test failure because it's matching the declaration of @llvm.assume.

The CHECK-NOT you've added in the original function's codeRepl block probably makes the implicit-check-not redundant, though.

15

Should this be "declare {{.*}}@llvm.assume"?

16

Should this be "define {{.*}}@f.cold.1("?

[hiraditya]

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

3

Hm, my theory is that the -implicit-check-not is probably causing a test failure because it's matching the declaration of @llvm.assume.

possible.

The CHECK-NOT you've added in the original function's codeRepl block probably makes the implicit-check-not redundant, though.

Thanks.

15

I'll make these changes before committing.

[hiraditya]

Closed in: https://reviews.llvm.org/rL373807

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.
     bool isEligible() const { return !Blocks.empty(); }
 
     /// Compute the set of input values and output values for the code.
     ///

llvm/lib/Analysis/AssumptionCache.cpp

void AssumptionCache::unregisterAssumption(CallInst *CI) {
   SmallVector<Value *, 16> Affected;
   findAffectedValues(CI, Affected);
 
   for (auto &AV : Affected) {
     auto AVI = AffectedValues.find_as(AV);
     if (AVI != AffectedValues.end())
       AffectedValues.erase(AVI);
   }
-  remove_if(AssumeHandles, [CI](WeakTrackingVH &VH) { return CI == VH; });
+  AssumeHandles.erase(remove_if(AssumeHandles,
+                                [CI](WeakTrackingVH &VH) { return CI == VH; }),
+                      AssumeHandles.end());
 }
 
 void AssumptionCache::AffectedValueCallbackVH::deleted() {
   auto AVI = AC->AffectedValues.find(getValPtr());
   if (AVI != AC->AffectedValues.end())
     AC->AffectedValues.erase(AVI);
   // 'this' now dangles!
 }

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