[GVNSink][Fix] Unregister assumptions that are replaced (PR49043)

[jdoerfert]

Summary

If we replace an llvm.assume with another one and we have an active
assumption cache, we might end up with an llvm.assume in the cache
twice, breaking an invariant. Avoid this by unregistering the assumption
when necessary.

[nikic]

Can you explain why this requires manual handling, and is not handled by the ValueHandle RAUW logic in AC?

[jdoerfert]

In D96106#2544339, @nikic wrote:

Can you explain why this requires manual handling, and is not handled by the ValueHandle RAUW logic in AC?

Good point. The current value handle RAUW logic is only dealing with affected values, the idea is still valid.
I could replace the

struct ResultElem {                                                                                                                                                                                                                                        
  WeakTrackingVH Assume;

with a CallbackVH to deal with this.

I'll try it out, thought's welcome.

[jdoerfert]

@nikic Using CallbackVH for the llvm.assume in the Cache is not as simple as I thought. I will do it but it changes the data structures a bit, unsure how that will impact compile time.

[jdoerfert]

llvm/lib/Transforms/Scalar/GVNSink.cpp

884

@nikic This is where one assume is replaced with another, e.g., in the attached test.

[nikic]

llvm/lib/Transforms/Scalar/GVNSink.cpp

884

Okay, I think now I get it. The problem is that we use a WeakTrackingVH, which means that it will follow RAUW, and we do multiple replacements to the same assume.

This makes me wonder though: Can't we just replace the WeakTrackingVH with a WeakVH? Doing RAUW on assumes is not really meaningful because the instruction has no return value. I believe it's fine to just ignore RAUW by using WeakVH, and then in this case, the old assumes would get nulled out by the eraseFromParent below.

[jdoerfert]

llvm/lib/Transforms/Scalar/GVNSink.cpp

884

Trying that now. Avoiding dangling references is still kinda nice but I'm OK with a minimal fix as well. Either way.

llvm/lib/Transforms/Scalar/GVNSink.cpp

 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/ArrayRecycler.h"
 #include "llvm/Support/AtomicOrdering.h"
uint32_t getMemoryUseOrder(Instruction *Inst) {
     return 0;
   }
 };
 
 //===----------------------------------------------------------------------===//
 
 class GVNSink {
 public:
-  GVNSink() = default;
+  GVNSink(AssumptionCache *AC) : AC(AC) {}
 
   bool run(Function &F) {
     LLVM_DEBUG(dbgs() << "GVNSink: running on function @" << F.getName()
                       << "\n");
 
     unsigned NumSunk = 0;
     ReversePostOrderTraversal<Function*> RPOT(&F);
     for (auto *N : RPOT)
       NumSunk += sinkBB(N);
 
     return NumSunk > 0;
   }
 
 private:
+  AssumptionCache *AC;
   ValueTable VN;
 
   bool shouldAvoidSinkingInstruction(Instruction *I) {
     // These instructions may change or break semantics if moved.
     if (isa<PHINode>(I) || I->isEHPad() || isa<AllocaInst>(I) ||
         I->getType()->isTokenTy())
       return true;
     return false;
for (unsigned O = 0, E = I0->getNumOperands(); O != E; ++O)
     I0->getOperandUse(O).set(NewOperands[O]);
   I0->moveBefore(&*BBEnd->getFirstInsertionPt());
 
   // Update metadata and IR flags.
   for (auto *I : Insts)
     if (I != I0) {
       combineMetadataForCSE(I0, I, true);
       I0->andIRFlags(I);
+      // If assumptions are tracked we need to unregister any we will replace.
+      if (AC && isa<CallInst>(I) &&
+          match(cast<CallInst>(I),
+                llvm::PatternMatch::m_Intrinsic<Intrinsic::assume>()))
+        AC->unregisterAssumption(cast<CallInst>(I));
     }
 
   for (auto *I : Insts)
     if (I != I0)
       I->replaceAllUsesWith(I0);
   foldPointlessPHINodes(BBEnd);
 
   // Finally nuke all instructions apart from the common instruction.
   for (auto *I : Insts)
     if (I != I0)
       I->eraseFromParent();
 
   NumRemoved += Insts.size() - 1;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Pass machinery / boilerplate
 
 class GVNSinkLegacyPass : public FunctionPass {
 public:
   static char ID;
 
   GVNSinkLegacyPass() : FunctionPass(ID) {
     initializeGVNSinkLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
-    GVNSink G;
+    AssumptionCache *AC = nullptr;
+    if (auto *ACT = getAnalysisIfAvailable<AssumptionCacheTracker>())
+      AC = &ACT->getAssumptionCache(F);
+    GVNSink G(AC);
     return G.run(F);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addPreserved<GlobalsAAWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 PreservedAnalyses GVNSinkPass::run(Function &F, FunctionAnalysisManager &AM) {
-  GVNSink G;
+  AssumptionCache *AC = AM.getCachedResult<AssumptionAnalysis>(F);
+  GVNSink G(AC);
   if (!G.run(F))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserve<GlobalsAA>();
   return PA;
 }
 

llvm/test/Transforms/GVNSink/assumption.ll

+; RUN: opt < %s -S -passes="print<assumptions>,gvn-sink,loop-unroll" -unroll-count=3 | FileCheck %s
+;
+; This crashed because the cached assumption was replaced and the replacement
+; was then in the cache twice.
+;
+; PR49043
+
+@g = external global i32
+
+define void @main() {
+bb:
+  %i1.i = load volatile i32, i32* @g
+  %i32.i = icmp eq i32 %i1.i, 0
+  call void @llvm.assume(i1 %i32.i) #3
+  br label %bb4.i
+
+bb4.i:                                            ; preds = %bb4.i, %bb
+  %i.i = load volatile i32, i32* @g
+  %i3.i = icmp eq i32 %i.i, 0
+  call void @llvm.assume(i1 %i3.i) #3
+  br label %bb4.i
+
+func_1.exit:                                      ; No predecessors!
+  unreachable
+}
+
+declare void @llvm.assume(i1)
+
+; CHECK:  call void @llvm.assume(
+; CHECK:  call void @llvm.assume(
+; CHECK:  call void @llvm.assume(
+