[SelectionDAG] Transitively copy NodeExtraInfo on RAUW

During legalization of the SelectionDAG, some nodes are replaced with
arch-specific nodes. These may be complex nodes, where the root node no
longer corresponds to the node that should carry the extra info.

Fix the issue by copying extra info to the new node and all its new
transitive operands during RAUW. See code comments for more details.

This fixes the remaining pcsections-atomics.ll tests on X86.

v2: Optimize copyExtraInfo() deep copy. For now we assume that only
NodeExtraInfo that have PCSections set require deep copy. Furthermore,
limit the depth of graph search while pre-populating the visited set,
assuming the to-be-replaced subgraph 'From' has limited complexity. An
assertion catches if the maximum depth needs to be increased.

Reviewed By: dvyukov

Differential Revision: https://reviews.llvm.org/D144677

Author: melver

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -17,6 +17,7 @@
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -12223,7 +12224,56 @@ void SelectionDAG::copyExtraInfo(SDNode *From, SDNode *To) {
   // Use of operator[] on the DenseMap may cause an insertion, which invalidates
   // the iterator, hence the need to make a copy to prevent a use-after-free.
   NodeExtraInfo Copy = I->second;
-  SDEI[To] = std::move(Copy);
+  if (LLVM_LIKELY(!Copy.PCSections)) {
+    // No deep copy required for the types of extra info set.
+    SDEI[To] = std::move(Copy);
+    return;
+  }
+
+  // We need to copy NodeExtraInfo to all _new_ nodes that are being introduced
+  // through the replacement of From with To. Otherwise, replacements of a node
+  // (From) with more complex nodes (To and its operands) may result in lost
+  // extra info where the root node (To) is insignificant in further propagating
+  // and using extra info when further lowering to MIR.
+  //
+  // In the first step pre-populate the visited set with the nodes reachable
+  // from the old From node. This avoids copying NodeExtraInfo to parts of the
+  // DAG that is not new and should be left untouched.
+  DenseSet<const SDNode *> Visited;
+  constexpr int MaxDepth = 16;
+  auto VisitFrom = [&Visited](auto &&Self, SDNode *N, int Depth) {
+    if (Depth >= MaxDepth)
+      return;
+    if (!Visited.insert(N).second)
+      return;
+    for (const SDValue &Op : N->op_values())
+      Self(Self, Op.getNode(), Depth + 1);
+  };
+  VisitFrom(VisitFrom, From, 0);
+
+  // Copy extra info to To and all its transitive operands (that are new).
+  auto DeepCopyTo = [this, &Copy, &Visited](auto &&Self, SDNode *To) {
+    if (!Visited.insert(To).second)
+      return true;
+    if (getEntryNode().getNode() == To) {
+      // This should not happen - and if it did, that means From has a depth
+      // greater or equal to MaxDepth, and VisitFrom() could not visit all
+      // common operands. As a result, we're able to reach the entry node.
+      assert(false && "Too complex 'From' node - increase MaxDepth?");
+      return false;
+    }
+    for (const SDValue &Op : To->op_values()) {
+      if (!Self(Self, Op.getNode()))
+        return false;
+    }
+    SDEI[To] = Copy;
+    return true;
+  };
+
+  if (LLVM_UNLIKELY(!DeepCopyTo(DeepCopyTo, To))) {
+    // Fallback - see assert above.
+    SDEI[To] = std::move(Copy);
+  }
 }
 
 #ifndef NDEBUG

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

@@ -851,6 +851,12 @@ void SDNode::print_details(raw_ostream &OS, const SelectionDAG *G) const {
           Dbg->print(OS);
     } else if (getHasDebugValue())
       OS << " [NoOfDbgValues>0]";
+
+    if (const auto *MD = G ? G->getPCSections(this) : nullptr) {
+      OS << " [pcsections ";
+      MD->printAsOperand(OS, G->getMachineFunction().getFunction().getParent());
+      OS << ']';
+    }
   }
 }