[SCCP] Add API for updating the state of the Solver.

[labrinea]

Summary

• Notify users when replacing a value.
• Invalidate a deleted instruction.

Discussion: https://discourse.llvm.org/t/rfc-should-we-enable-function-specialization/61518

[labrinea]

When invalidating remove the predicate info from additional users of all instructions, not only CmpInst (the sanitizer found a use after free for a PHINode).

[labrinea]

Changes from last revision:

• allow removing lattice state of struct-typed values
• maybe it possible to conditionally notify users when replacing a value

[labrinea]

ping

[labrinea]

Abandoning in favor of D126455.

llvm/include/llvm/Transforms/Utils/SCCPSolver.h

public:
   void markArgInFuncSpecialization(Function *F,
                                    const SmallVectorImpl<ArgInfo> &Args);
 
   /// Mark all of the blocks in function \p F non-executable. Clients can used
   /// this method to erase a function from the module (e.g., if it has been
   /// completely specialized and is no longer needed).
   void markFunctionUnreachable(Function *F);
 
+  /// Replace all uses and visit their users (optionally remove from parent).
+  void replaceUsesOfWith(Value *Old, Value *New, bool NotifyUsers=true,
+                                                 bool ForceDeletion=false);
+
+  /// Notify the visitor that this instruction has been deleted.
+  void invalidate(Instruction *I);
+
   void visit(Instruction *I);
   void visitCall(CallInst &I);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_SCCPSOLVER_H

llvm/lib/Transforms/Utils/SCCPSolver.cpp

 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/ValueLattice.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include <cassert>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "sccp"
 

 namespace llvm {
 
 /// Helper class for SCCPSolver. This implements the instruction visitor and
 /// holds all the state.
 class SCCPInstVisitor : public InstVisitor<SCCPInstVisitor> {
   const DataLayout &DL;
   std::function<const TargetLibraryInfo &(Function &)> GetTLI;
   SmallPtrSet<BasicBlock *, 8> BBExecutable; // The BBs that are executable.
+  DenseSet<Value *> InvalidatedInsts; // Instructions which have been deleted.
   DenseMap<Value *, ValueLatticeElement>
       ValueState; // The state each value is in.
 
   /// StructValueState - This maintains ValueState for values that have
   /// StructType, for example for formal arguments, calls, insertelement, etc.
   DenseMap<std::pair<Value *, unsigned>, ValueLatticeElement> StructValueState;
 
   /// GlobalValue - If we are tracking any values for the contents of a global
private:
   /// markEdgeExecutable - Mark a basic block as executable, adding it to the BB
   /// work list if it is not already executable.
   bool markEdgeExecutable(BasicBlock *Source, BasicBlock *Dest);
 
   // getFeasibleSuccessors - Return a vector of booleans to indicate which
   // successors are reachable from a given terminator instruction.
   void getFeasibleSuccessors(Instruction &TI, SmallVectorImpl<bool> &Succs);
 
-  // OperandChangedState - This method is invoked on all of the users of an
-  // instruction that was just changed state somehow.  Based on this
-  // information, we need to update the specified user of this instruction.
-  void operandChangedState(Instruction *I) {
-    if (BBExecutable.count(I->getParent())) // Inst is executable?
-      visit(*I);
-  }
-
   // Add U as additional user of V.
   void addAdditionalUser(Value *V, User *U) {
     AdditionalUsers[V].insert(U);
     AdditionalUserOf[U].insert(V);
   }
 
   // Mark I's users as changed, including AdditionalUsers.
   void markUsersAsChanged(Value *I) {
public:
   void findAdditionalUsersFor(Value *V, SmallVectorImpl<Instruction *> &Users) {
     auto Iter = AdditionalUsers.find(V);
     if (Iter != AdditionalUsers.end())
       for (User *U : Iter->second)
         if (auto *I = dyn_cast<Instruction>(U))
           Users.push_back(I);
   }
 
+  void invalidate(Instruction *I) {
+    auto Iter = AdditionalUsers.find(I);
+    if (Iter != AdditionalUsers.end())
+      for (User *U : Iter->second)
+        if (auto *I = dyn_cast<Instruction>(U))
+          removePredicateInfoFor(I);
+    removeFromAdditionalUsers(I);
+    removeFromAdditionalUserOf(I);
+    AdditionalUsers.erase(I);
+    AdditionalUserOf.erase(I);
+    removeLatticeValueFor(I);
+    InvalidatedInsts.insert(I);
+  }
+
+  // OperandChangedState - This method is invoked on all of the users of an
+  // instruction that was just changed state somehow.  Based on this
+  // information, we need to update the specified user of this instruction.
+  void operandChangedState(Instruction *I) {
+    if (BBExecutable.count(I->getParent())) // Inst is executable?
+      visit(*I);
+  }
+
   DomTreeUpdater getDTU(Function &F) {
     auto A = AnalysisResults.find(&F);
     assert(A != AnalysisResults.end() && "Need analysis results for function.");
     return {A->second.DT, A->second.PDT, DomTreeUpdater::UpdateStrategy::Lazy};
   }
 
   SCCPInstVisitor(const DataLayout &DL,
                   std::function<const TargetLibraryInfo &(Function &)> GetTLI,
std::vector<ValueLatticeElement> getStructLatticeValueFor(Value *V) const {
     for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i) {
       auto I = StructValueState.find(std::make_pair(V, i));
       assert(I != StructValueState.end() && "Value not in valuemap!");
       StructValues.push_back(I->second);
     }
     return StructValues;
   }
 
-  void removeLatticeValueFor(Value *V) { ValueState.erase(V); }
+  void removeLatticeValueFor(Value *V) {
+    if (auto *STy = dyn_cast<StructType>(V->getType()))
+      for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i)
+        StructValueState.erase(std::make_pair(V, i));
+    else
+      ValueState.erase(V);
+  }
 
   const ValueLatticeElement &getLatticeValueFor(Value *V) const {
     assert(!V->getType()->isStructTy() &&
            "Should use getStructLatticeValueFor");
     DenseMap<Value *, ValueLatticeElement>::const_iterator I =
         ValueState.find(V);
     assert(I != ValueState.end() &&
            "V not found in ValueState nor Paramstate map!");
void SCCPSolver::markArgInFuncSpecialization(
     Function *F, const SmallVectorImpl<ArgInfo> &Args) {
   Visitor->markArgInFuncSpecialization(F, Args);
 }
 
 void SCCPSolver::markFunctionUnreachable(Function *F) {
   Visitor->markFunctionUnreachable(F);
 }
 
+static bool canRemoveInstruction(Instruction *I) {
+  if (wouldInstructionBeTriviallyDead(I))
+    return true;
+
+  // Some instructions can be handled but are rejected above. Catch
+  // those cases by falling through to here.
+  // TODO: Mark globals as being constant earlier, so
+  // TODO: wouldInstructionBeTriviallyDead() knows that atomic loads
+  // TODO: are safe to remove.
+  return isa<LoadInst>(I);
+}
+
+void SCCPSolver::replaceUsesOfWith(Value *Old, Value *New, bool NotifyUsers,
+                                                           bool ForceDeletion) {
+  // Record uses of Old to avoid visiting irrelevant uses of New later.
+  SmallVector<Instruction *> ToNotify;
+
+  if (NotifyUsers) {
+    Visitor->findAdditionalUsersFor(Old, ToNotify);
+    for (User *U : Old->users())
+      if (auto *I = dyn_cast<Instruction>(U))
+        if (I != Old)
+          ToNotify.push_back(I);
+  }
+
+  Old->replaceAllUsesWith(New);
+
+  // Remove the instruction from Block and Solver.
+  if (auto *I = dyn_cast<Instruction>(Old)) {
+    if (canRemoveInstruction(I) || ForceDeletion) {
+      Visitor->invalidate(I);
+      I->eraseFromParent();
+    }
+  }
+
+  for (Instruction *I : ToNotify)
+    Visitor->operandChangedState(I);
+}
+
+void SCCPSolver::invalidate(Instruction *I) { Visitor->invalidate(I); }
+
 void SCCPSolver::visit(Instruction *I) { Visitor->visit(I); }
 
 void SCCPSolver::visitCall(CallInst &I) { Visitor->visitCall(I); }