[SelectionDAG] Properly copy ExtraInfo on RAUW

[melver]

Summary

During SelectionDAG legalization SDNodes with associated extra info may
be replaced with a new SDNode. Preserve associated extra info on
ReplaceAllUsesWith and remove entries in DeallocateNode.

Depends on D130880

[melver]

Rebase.

[vitalybuka]

Would you be able to reproduce the issue as a new test in TEST_F(SelectionDAGAddressAnalysisTest ?

[melver]

Add unit test for RAUW.
Note: Since SelectionDAG construction needs a target, the AArch64 test is used here (as done for several other test cases that are arch-independent).

[melver]

In D130881#3703690, @vitalybuka wrote:

Would you be able to reproduce the issue as a new test in TEST_F(SelectionDAGAddressAnalysisTest ?

That test seems only related to BaseIndexOffset::computeAliasing(). Added to a different test that seems to be used for that (needs a target, so there's an AArch64 one).

[vitalybuka]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

11819

this assert is inconsistent with the rest of the code

[melver]

• Assert message.

[melver]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

11819

It's similar to the one in transferDbgValues(). Only that copyExtraInfo already takes the unwrapped pointers as it makes usage a bit easier where we don't have the source SDValue. I'd like to keep it to prevent easily avoidable mistakes.

[melver]

Rebase.

llvm/include/llvm/CodeGen/SelectionDAG.h

if (NoMerge)
       SDEI[Node].NoMerge = NoMerge;
   }
   /// Return NoMerge info associated with Node.
   bool getNoMergeSiteInfo(const SDNode *Node) const {
     auto I = SDEI.find(Node);
     return I != SDEI.end() ? I->second.NoMerge : false;
   }
 
+  /// Copy extra info associated with one node to another.
+  void copyExtraInfo(SDNode *From, SDNode *To);
+
   /// Return the current function's default denormal handling kind for the given
   /// floating point type.
   DenormalMode getDenormalMode(EVT VT) const {
     return MF->getDenormalMode(EVTToAPFloatSemantics(VT));
   }
 
   bool shouldOptForSize() const;
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

void SelectionDAG::DeallocateNode(SDNode *N) {
   // FIXME: There are places in SDag that have grown a dependency on the opcode
   // value in the released node.
   __asan_unpoison_memory_region(&N->NodeType, sizeof(N->NodeType));
   N->NodeType = ISD::DELETED_NODE;
 
   // If any of the SDDbgValue nodes refer to this SDNode, invalidate
   // them and forget about that node.
   DbgInfo->erase(N);
+
+  // Invalidate extra info.
+  SDEI.erase(N);
 }
 
 #ifndef NDEBUG
 /// VerifySDNode - Check the given SDNode.  Aborts if it is invalid.
 static void VerifySDNode(SDNode *N) {
   switch (N->getOpcode()) {
   default:
     break;

 void SelectionDAG::ReplaceAllUsesWith(SDValue FromN, SDValue To) {
   SDNode *From = FromN.getNode();
   assert(From->getNumValues() == 1 && FromN.getResNo() == 0 &&
          "Cannot replace with this method!");
   assert(From != To.getNode() && "Cannot replace uses of with self");
 
   // Preserve Debug Values
   transferDbgValues(FromN, To);
+  // Preserve extra info.
+  copyExtraInfo(From, To.getNode());
 
   // Iterate over all the existing uses of From. New uses will be added
   // to the beginning of the use list, which we avoid visiting.
   // This specifically avoids visiting uses of From that arise while the
   // replacement is happening, because any such uses would be the result
   // of CSE: If an existing node looks like From after one of its operands
   // is replaced by To, we don't want to replace of all its users with To
   // too. See PR3018 for more info.
if (From == To)
     return;
 
   // Preserve Debug Info. Only do this if there's a use.
   for (unsigned i = 0, e = From->getNumValues(); i != e; ++i)
     if (From->hasAnyUseOfValue(i)) {
       assert((i < To->getNumValues()) && "Invalid To location");
       transferDbgValues(SDValue(From, i), SDValue(To, i));
     }
+  // Preserve extra info.
+  copyExtraInfo(From, To);
 
   // Iterate over just the existing users of From. See the comments in
   // the ReplaceAllUsesWith above.
   SDNode::use_iterator UI = From->use_begin(), UE = From->use_end();
   RAUWUpdateListener Listener(*this, UI, UE);
   while (UI != UE) {
     SDNode *User = *UI;
 

 /// This can cause recursive merging of nodes in the DAG.
 ///
 /// This version can replace From with any result values.  To must match the
 /// number and types of values returned by From.
 void SelectionDAG::ReplaceAllUsesWith(SDNode *From, const SDValue *To) {
   if (From->getNumValues() == 1)  // Handle the simple case efficiently.
     return ReplaceAllUsesWith(SDValue(From, 0), To[0]);
 
+  for (unsigned i = 0, e = From->getNumValues(); i != e; ++i) {
     // Preserve Debug Info.
-  for (unsigned i = 0, e = From->getNumValues(); i != e; ++i)
     transferDbgValues(SDValue(From, i), To[i]);
+    // Preserve extra info.
+    copyExtraInfo(From, To[i].getNode());
+  }
 
   // Iterate over just the existing users of From. See the comments in
   // the ReplaceAllUsesWith above.
   SDNode::use_iterator UI = From->use_begin(), UE = From->use_end();
   RAUWUpdateListener Listener(*this, UI, UE);
   while (UI != UE) {
     SDNode *User = *UI;
 
void SelectionDAG::ReplaceAllUsesOfValueWith(SDValue From, SDValue To){
   // Handle the simple, trivial, case efficiently.
   if (From.getNode()->getNumValues() == 1) {
     ReplaceAllUsesWith(From, To);
     return;
   }
 
   // Preserve Debug Info.
   transferDbgValues(From, To);
+  copyExtraInfo(From.getNode(), To.getNode());
 
   // Iterate over just the existing users of From. See the comments in
   // the ReplaceAllUsesWith above.
   SDNode::use_iterator UI = From.getNode()->use_begin(),
                        UE = From.getNode()->use_end();
   RAUWUpdateListener Listener(*this, UI, UE);
   while (UI != UE) {
     SDNode *User = *UI;

 void SelectionDAG::ReplaceAllUsesOfValuesWith(const SDValue *From,
                                               const SDValue *To,
                                               unsigned Num){
   // Handle the simple, trivial case efficiently.
   if (Num == 1)
     return ReplaceAllUsesOfValueWith(*From, *To);
 
   transferDbgValues(*From, *To);
+  copyExtraInfo(From->getNode(), To->getNode());
 
   // Read up all the uses and make records of them. This helps
   // processing new uses that are introduced during the
   // replacement process.
   SmallVector<UseMemo, 4> Uses;
   for (unsigned i = 0; i != Num; ++i) {
     unsigned FromResNo = From[i].getResNo();
     SDNode *FromNode = From[i].getNode();
case ISD::FMAXNUM: {
     if (Opcode == ISD::FMAXNUM)
       NeutralAF.changeSign();
 
     return getConstantFP(NeutralAF, DL, VT);
   }
   }
 }
 
+void SelectionDAG::copyExtraInfo(SDNode *From, SDNode *To) {
+  assert(From && To && "Invalid SDNode; empty source SDValue?");
+  auto I = SDEI.find(From);
+  if (I == SDEI.end())
+    return;
+  SDEI[To] = I->second;
+}
+
 #ifndef NDEBUG
 static void checkForCyclesHelper(const SDNode *N,
                                  SmallPtrSetImpl<const SDNode*> &Visited,
                                  SmallPtrSetImpl<const SDNode*> &Checked,
                                  const llvm::SelectionDAG *DAG) {
   // If this node has already been checked, don't check it again.
   if (Checked.count(N))
     return;

llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp

TEST_F(AArch64SelectionDAGTest, TestFold_STEP_VECTOR) {
   auto VecVT = EVT::getVectorVT(Context, MVT::i8, 16, true);
 
   // Should create SPLAT_VECTOR
   SDValue Zero = DAG->getConstant(0, Loc, IntVT);
   SDValue Op = DAG->getNode(ISD::STEP_VECTOR, Loc, VecVT, Zero);
   EXPECT_EQ(Op.getOpcode(), ISD::SPLAT_VECTOR);
 }
 
+TEST_F(AArch64SelectionDAGTest, ReplaceAllUsesWith) {
+  SDLoc Loc;
+  EVT IntVT = EVT::getIntegerVT(Context, 8);
+
+  SDValue N0 = DAG->getConstant(0x42, Loc, IntVT);
+  SDValue N1 = DAG->getRegister(0, IntVT);
+  // Construct node to fill arbitrary ExtraInfo.
+  SDValue N2 = DAG->getNode(ISD::SUB, Loc, IntVT, N0, N1);
+  EXPECT_FALSE(DAG->getHeapAllocSite(N2.getNode()));
+  EXPECT_FALSE(DAG->getNoMergeSiteInfo(N2.getNode()));
+  MDNode *MD = MDNode::get(Context, None);
+  DAG->addHeapAllocSite(N2.getNode(), MD);
+  DAG->addNoMergeSiteInfo(N2.getNode(), true);
+  EXPECT_EQ(DAG->getHeapAllocSite(N2.getNode()), MD);
+  EXPECT_TRUE(DAG->getNoMergeSiteInfo(N2.getNode()));
+
+  SDValue Root = DAG->getNode(ISD::ADD, Loc, IntVT, N2, N2);
+  EXPECT_EQ(Root->getOperand(0)->getOpcode(), ISD::SUB);
+  // Create new node and check that ExtraInfo is propagated on RAUW.
+  SDValue New = DAG->getNode(ISD::ADD, Loc, IntVT, N1, N1);
+  EXPECT_FALSE(DAG->getHeapAllocSite(New.getNode()));
+  EXPECT_FALSE(DAG->getNoMergeSiteInfo(New.getNode()));
+
+  DAG->ReplaceAllUsesWith(N2, New);
+  EXPECT_EQ(Root->getOperand(0), New);
+  EXPECT_EQ(DAG->getHeapAllocSite(New.getNode()), MD);
+  EXPECT_TRUE(DAG->getNoMergeSiteInfo(New.getNode()));
+}
+
 } // end namespace llvm