[CGP] Avoid segmentation fault when doing PHI node simplifications

[bjope]

Summary

Made PHI node simplifiations more robust in several ways:

• Minor refactoring to let the SimplificationTracker own the

sets with new PHI/Select nodes that are introduced. This is
maybe not mapping to the original intention with the
SimplificationTracker, but IMHO it encapsulates the logic behind
those sets a little bit better.

• MatchPhiNode can sometimes populate the Matched set with

several entries, where it maps one PHI node to different candidates
for replacement. The Matched set is changed into a SmallSetVector
to make sure we get a deterministic iteration when doing
the replacements.

• As described above we may get several different replacements

for a single PHI node. The loop in MatchPhiSet that is doing
the replacements could end up calling eraseFromParent several
times for the same PHI node, resulting in segmentation faults.
This problem was supposed to be fixed in rL327250, but due to
the non-determinism(?) it only appeared to be fixed (I still
got crashes sometime when turning on/off -print-after-all etc
to get different iteration order in the DenseSets).
With this patch we follow the deterministic ordering in the
Matched set when replacing the PHI nodes. If we find a new
replacement for an already replaced PHI node we replace the
new replacement by the old replacement instead. This is quite
similar to what happened in the rl327250 patch, but here we
also recursively verify that the old replacement hasn't been
replaced already.

• It was really hard to track down the fault described above

(segementation fault due to doing eraseFromParent multiple
times for the same instruction). The fault was intermittent and
small changes in the code, or simply turning on -print-after-all
etc could make the problem go away. This was basically due to
the iteration over PhiNodesToMatch in MatchPhiSet no being
deterministic. Therefore I've changed the data structure for
the SimplificationTracker::AllPhiNodes into an SmallSetVector.
This gives a deterministic behavior.

[bjope]

I had not seen https://reviews.llvm.org/D43758 so this is my attempt at solving the same problem. Although, this patch also takes care of some problems with non-deterministic behavior.

I'm also a little bit afraid that the solution in D43758 has some other limitations. See the XXX at line 2983.
The new loop that was added in D43758 (to avoid multiple eraseFromParent() in the second loop) is doing eraseFromParent() itself. Without removing the erased PHI-nodes from the PhiNodesToMatch, and without making sure that the erased PHI nodes aren't already in the MatchedPHINodeMapping.

[skatkov]

lib/CodeGen/CodeGenPrepare.cpp

2642

It can be an assert now (taking into account two loops in the place of invocation.

2986

Make sense however I guess that in general MV.second will never be in the PhiNodesToMatch.

The idea here that PhiNodesToMatch contains a Phi nodes created on previous steps. And they are matched against already existed Phi nodes before this algorithm started to work. So in general PhiNodesToMatch should not intersect with MV.second.

I guess assert will be ok in this case.

2992

you do not need {} now.

[bjope]

Basically moved the solution from rL327250 into SimplificationTracker::ReplacePhi. This avoids the need for the DenseSet (with non-deterministic iteration) that was added in rL327250.

I also made sure that we recursively check if the new replacement already have been replaced when doing the simplifications for "already matched" nodes.

lib/CodeGen/CodeGenPrepare.cpp

 };
 
 /// \brief Keep track of simplification of Phi nodes.
 /// Accept the set of all phi nodes and erase phi node from this set
 /// if it is simplified.
 class SimplificationTracker {
   DenseMap<Value *, Value *> Storage;
   const SimplifyQuery &SQ;
-  SmallPtrSetImpl<PHINode *> &AllPhiNodes;
-  SmallPtrSetImpl<SelectInst *> &AllSelectNodes;
+  // Tracks newly created Phi nodes. We use a SetVector to get deterministic
+  // order when iterating over the set in MatchPhiSet.
+  SmallSetVector<PHINode *, 32> AllPhiNodes;
+  // Tracks newly created Select nodes.
+  SmallPtrSet<SelectInst *, 32> AllSelectNodes;
 
 public:
-  SimplificationTracker(const SimplifyQuery &sq,
-                        SmallPtrSetImpl<PHINode *> &APN,
-                        SmallPtrSetImpl<SelectInst *> &ASN)
-      : SQ(sq), AllPhiNodes(APN), AllSelectNodes(ASN) {}
+  SimplificationTracker(const SimplifyQuery &sq)
+      : SQ(sq) {}
 
   Value *Get(Value *V) {
     do {
       auto SV = Storage.find(V);
       if (SV == Storage.end())
         return V;
       V = SV->second;
     } while (true);
while (!WorkList.empty()) {
         continue;
       if (auto *PI = dyn_cast<Instruction>(P))
         if (Value *V = SimplifyInstruction(cast<Instruction>(PI), SQ)) {
           for (auto *U : PI->users())
             WorkList.push_back(cast<Value>(U));
           Put(PI, V);
           PI->replaceAllUsesWith(V);
           if (auto *PHI = dyn_cast<PHINode>(PI))
-            AllPhiNodes.erase(PHI);
+            AllPhiNodes.remove(PHI);
           if (auto *Select = dyn_cast<SelectInst>(PI))
             AllSelectNodes.erase(Select);
           PI->eraseFromParent();
         }
     }
     return Get(Val);
   }
 
   void Put(Value *From, Value *To) {
     Storage.insert({ From, To });
   }
+
+  void ReplacePhi(PHINode *From, PHINode *To) {
+    if (!AllPhiNodes.count(From)) {
+      // The From pointer is probably dangling, so return here to make sure we
+      // do not dereference or try to erase the node from parent a second time.
+      DEBUG(dbgs() << "CGP: Trying to replace an already replaced PHI node.");
+      return;
+    }
+    Put(From, To);
+    From->replaceAllUsesWith(To);
+    AllPhiNodes.remove(From);
+    From->eraseFromParent();
+  }
+
+  SmallSetVector<PHINode *, 32>& newPhiNodes() { return AllPhiNodes; }
+
+  void insertNewPhi(PHINode *PN) { AllPhiNodes.insert(PN); }
+
+  void insertNewSelect(SelectInst *SI) { AllSelectNodes.insert(SI); }
+
+  unsigned countNewPhiNodes() const { return AllPhiNodes.size(); }
+
+  unsigned countNewSelectNodes() const { return AllSelectNodes.size(); }
+
+  void destroyNewNodes(Type *CommonType) {
+    // For safe erasing, replace the uses with dummy value first.
+    auto Dummy = UndefValue::get(CommonType);
+    for (auto I : AllPhiNodes) {
+      I->replaceAllUsesWith(Dummy);
+      I->eraseFromParent();
+    }
+    AllPhiNodes.clear();
+    for (auto I : AllSelectNodes) {
+      I->replaceAllUsesWith(Dummy);
+      I->eraseFromParent();
+    }
+    AllSelectNodes.clear();
+  }
 };
 
 /// \brief A helper class for combining addressing modes.
 class AddressingModeCombiner {
   typedef std::pair<Value *, BasicBlock *> ValueInBB;
   typedef DenseMap<ValueInBB, Value *> FoldAddrToValueMapping;
   typedef std::pair<PHINode *, PHINode *> PHIPair;
 
private:
   //   v = load p
   // Map is
   //   <p1, BB1> -> b1
   //   <p2, BB2> -> b2
   // Request is
   //   <p, BB3> -> ?
   // The function tries to find or build phi [b1, BB1], [b2, BB2] in BB3
   Value *findCommon(FoldAddrToValueMapping &Map) {
-    // Tracks newly created Phi nodes.
-    SmallPtrSet<PHINode *, 32> NewPhiNodes;
-    // Tracks newly created Select nodes.
-    SmallPtrSet<SelectInst *, 32> NewSelectNodes;
     // Tracks the simplification of newly created phi nodes. The reason we use
     // this mapping is because we will add new created Phi nodes in AddrToBase.
     // Simplification of Phi nodes is recursive, so some Phi node may
     // be simplified after we added it to AddrToBase.
     // Using this mapping we can find the current value in AddrToBase.
-    SimplificationTracker ST(SQ, NewPhiNodes, NewSelectNodes);
+    SimplificationTracker ST(SQ);
 
     // First step, DFS to create PHI nodes for all intermediate blocks.
     // Also fill traverse order for the second step.
     SmallVector<ValueInBB, 32> TraverseOrder;
-    InsertPlaceholders(Map, TraverseOrder, NewPhiNodes, NewSelectNodes);
+    InsertPlaceholders(Map, TraverseOrder, ST);
 
     // Second Step, fill new nodes by merged values and simplify if possible.
     FillPlaceholders(Map, TraverseOrder, ST);
 
-    if (!AddrSinkNewSelects && NewSelectNodes.size() > 0) {
-      DestroyNodes(NewPhiNodes);
-      DestroyNodes(NewSelectNodes);
+    if (!AddrSinkNewSelects && ST.countNewSelectNodes() > 0) {
+      ST.destroyNewNodes(CommonType);
       return nullptr;
     }
 
     // Now we'd like to match New Phi nodes to existed ones.
     unsigned PhiNotMatchedCount = 0;
-    if (!MatchPhiSet(NewPhiNodes, ST, AddrSinkNewPhis, PhiNotMatchedCount)) {
-      DestroyNodes(NewPhiNodes);
-      DestroyNodes(NewSelectNodes);
+    if (!MatchPhiSet(ST, AddrSinkNewPhis, PhiNotMatchedCount)) {
+      ST.destroyNewNodes(CommonType);
       return nullptr;
     }
 
     auto *Result = ST.Get(Map.find(Original)->second);
     if (Result) {
-      NumMemoryInstsPhiCreated += NewPhiNodes.size() + PhiNotMatchedCount;
-      NumMemoryInstsSelectCreated += NewSelectNodes.size();
+      NumMemoryInstsPhiCreated += ST.countNewPhiNodes() + PhiNotMatchedCount;
+      NumMemoryInstsSelectCreated += ST.countNewSelectNodes();
     }
     return Result;
   }
 
-  /// \brief Destroy nodes from a set.
-  template <typename T> void DestroyNodes(SmallPtrSetImpl<T *> &Instructions) {
-    // For safe erasing, replace the Phi with dummy value first.
-    auto Dummy = UndefValue::get(CommonType);
-    for (auto I : Instructions) {
-      I->replaceAllUsesWith(Dummy);
-      I->eraseFromParent();
-    }
-  }
-
   /// \brief Try to match PHI node to Candidate.
   /// Matcher tracks the matched Phi nodes.
   bool MatchPhiNode(PHINode *PHI, PHINode *Candidate,
-                    DenseSet<PHIPair> &Matcher,
-                    SmallPtrSetImpl<PHINode *> &PhiNodesToMatch) {
+                    SmallSetVector<PHIPair, 8> &Matcher,
+                    SmallSetVector<PHINode *, 32> &PhiNodesToMatch) {
     SmallVector<PHIPair, 8> WorkList;
     Matcher.insert({ PHI, Candidate });
     WorkList.push_back({ PHI, Candidate });
     SmallSet<PHIPair, 8> Visited;
     while (!WorkList.empty()) {
       auto Item = WorkList.pop_back_val();
       if (!Visited.insert(Item).second)
         continue;
while (!WorkList.empty()) {
         Matcher.insert({ FirstPhi, SecondPhi });
         // But me must check it.
         WorkList.push_back({ FirstPhi, SecondPhi });
       }
     }
     return true;
   }
 
-  /// \brief For the given set of PHI nodes try to find their equivalents.
+  /// \brief For the given set of PHI nodes (in the SimplificationTracker) try
+  /// to find their equivalents.
   /// Returns false if this matching fails and creation of new Phi is disabled.
-  bool MatchPhiSet(SmallPtrSetImpl<PHINode *> &PhiNodesToMatch,
-                   SimplificationTracker &ST, bool AllowNewPhiNodes,
+  bool MatchPhiSet(SimplificationTracker &ST, bool AllowNewPhiNodes,
                    unsigned &PhiNotMatchedCount) {
-    DenseSet<PHIPair> Matched;
+    // Use a SetVector for Matched to make sure we do replacements (ReplacePhi)
+    // in a deterministic order below.
+    SmallSetVector<PHIPair, 8> Matched;
     SmallPtrSet<PHINode *, 8> WillNotMatch;
+    SmallSetVector<PHINode *, 32> &PhiNodesToMatch = ST.newPhiNodes();
     while (PhiNodesToMatch.size()) {
       PHINode *PHI = *PhiNodesToMatch.begin();
 
       // Add us, if no Phi nodes in the basic block we do not match.
       WillNotMatch.clear();
       WillNotMatch.insert(PHI);
 
       // Traverse all Phis until we found equivalent or fail to do that.
while (PhiNodesToMatch.size()) {
         // If we matched phi node to different but identical phis then
         // make a simplification here.
         DenseMap<PHINode *, PHINode *> MatchedPHINodeMapping;
         for (auto MV : Matched) {
           auto AlreadyMatched = MatchedPHINodeMapping.find(MV.first);
           if (AlreadyMatched != MatchedPHINodeMapping.end()) {
             MV.second->replaceAllUsesWith(AlreadyMatched->second);
             ST.Put(MV.second, AlreadyMatched->second);
+            // XXX: This looks dangerous. What if MV.second already have been
+            // removed this way? Or if MV.second also is in the Matched set (as
+            // a "first" node)? Shouldn't we also make sure to remove MV.second
+            // from PhiNodesToMatch now when it is erased?
             MV.second->eraseFromParent();
           } else
             MatchedPHINodeMapping.insert({ MV.first, MV.second });
         }
         // Replace all matched values and erase them.
         for (auto MV : MatchedPHINodeMapping) {
-          MV.first->replaceAllUsesWith(MV.second);
-          PhiNodesToMatch.erase(MV.first);
-          ST.Put(MV.first, MV.second);
-          MV.first->eraseFromParent();
+          ST.ReplacePhi(MV.first, MV.second);
         }
         Matched.clear();
         continue;
       }
       // If we are not allowed to create new nodes then bail out.
       if (!AllowNewPhiNodes)
         return false;
       // Just remove all seen values in matcher. They will not match anything.
       PhiNotMatchedCount += WillNotMatch.size();
       for (auto *P : WillNotMatch)
-        PhiNodesToMatch.erase(P);
+        PhiNodesToMatch.remove(P);
     }
     return true;
   }
   /// \brief Fill the placeholder with values from predecessors and simplify it.
   void FillPlaceholders(FoldAddrToValueMapping &Map,
                         SmallVectorImpl<ValueInBB> &TraverseOrder,
                         SimplificationTracker &ST) {
     while (!TraverseOrder.empty()) {
private:
 
   /// Starting from value recursively iterates over predecessors up to known
   /// ending values represented in a map. For each traversed block inserts
   /// a placeholder Phi or Select.
   /// Reports all new created Phi/Select nodes by adding them to set.
   /// Also reports and order in what basic blocks have been traversed.
   void InsertPlaceholders(FoldAddrToValueMapping &Map,
                           SmallVectorImpl<ValueInBB> &TraverseOrder,
-                          SmallPtrSetImpl<PHINode *> &NewPhiNodes,
-                          SmallPtrSetImpl<SelectInst *> &NewSelectNodes) {
+                          SimplificationTracker &ST) {
     SmallVector<ValueInBB, 32> Worklist;
     assert((isa<PHINode>(Original.first) || isa<SelectInst>(Original.first)) &&
            "Address must be a Phi or Select node");
     auto *Dummy = UndefValue::get(CommonType);
     Worklist.push_back(Original);
     while (!Worklist.empty()) {
       auto Current = Worklist.pop_back_val();
       // If value is not an instruction it is something global, constant,
while (!Worklist.empty()) {
           std::distance(pred_begin(CurrentBlock), pred_end(CurrentBlock));
       // if Current Value is not defined in this basic block we are interested
       // in values in predecessors.
       if (!IsDefinedInThisBB) {
         assert(PredCount && "Unreachable block?!");
         PHINode *PHI = PHINode::Create(CommonType, PredCount, "sunk_phi",
                                        &CurrentBlock->front());
         Map[Current] = PHI;
-        NewPhiNodes.insert(PHI);
+        ST.insertNewPhi(PHI);
         // Add all predecessors in work list.
         for (auto B : predecessors(CurrentBlock))
           Worklist.push_back({ CurrentValue, B });
         continue;
       }
       // Value is defined in this basic block.
       if (SelectInst *OrigSelect = dyn_cast<SelectInst>(CurrentI)) {
         // Is it OK to get metadata from OrigSelect?!
         // Create a Select placeholder with dummy value.
         SelectInst *Select =
             SelectInst::Create(OrigSelect->getCondition(), Dummy, Dummy,
                                OrigSelect->getName(), OrigSelect, OrigSelect);
         Map[Current] = Select;
-        NewSelectNodes.insert(Select);
+        ST.insertNewSelect(Select);
         // We are interested in True and False value in this basic block.
         Worklist.push_back({ OrigSelect->getTrueValue(), CurrentBlock });
         Worklist.push_back({ OrigSelect->getFalseValue(), CurrentBlock });
       } else {
         // It must be a Phi node then.
         auto *CurrentPhi = cast<PHINode>(CurrentI);
         // Create new Phi node for merge of bases.
         assert(PredCount && "Unreachable block?!");
         PHINode *PHI = PHINode::Create(CommonType, PredCount, "sunk_phi",
                                        &CurrentBlock->front());
         Map[Current] = PHI;
-        NewPhiNodes.insert(PHI);
+        ST.insertNewPhi(PHI);
 
         // Add all predecessors in work list.
         for (auto B : predecessors(CurrentBlock))
           Worklist.push_back({ CurrentPhi->getIncomingValueForBlock(B), B });
       }
     }
   }