[DAG] Fix in ReplaceAllUsesOfValuesWith

[bjope]

Summary

When doing SelectionDAG::ReplaceAllUsesOfValuesWith a worklist is
prepared containing all users that should be updated. Then we use
the RemoveNodeFromCSEMaps/AddModifiedNodeToCSEMaps helpers to handle
recursive CSE updates while doing the replacements.

This patch aims at solving a problem that could arise if the recursive
CSE updates would result in an SDNode present in the worklist is being
removed as a side-effect of morphing a prio user in the worklist.

To examplify such a scenario, imagine that we have these nodes in
the DAG

t12: i64 = add t8, t11
t13: i64 = add t12, t8
t14: i64 = add t11, t11
t15: i64 = add t14, t8
t16: i64 = sub t13, t15

and that the t8 uses should be replaced by t11. An initial worklist
(listing the users that should be morphed) could be [t12, t13, t15].
When updating t12 we get

t12: i64 = add t11, t11

which results in a CSE update that replaces t14 by t12, so we get

t15: i64 = add t12, t8

which results in a CSE update that replaces t13 by t12, so we get

t16: i64 = sub t12, t15

and then t13 is removed given that it was the last use of t13.

So when being done with the updates triggered by rewriting the use
of t8 in t12 the t13 node no longer exist. And we used to end up
hitting an assertion when continuing with the worklist aiming at
replacing the t8 uses in t13.

The solution is based on using a DAGUpdateListener, making sure that
we prune a user from the worklist if it is removed during the
recursive CSE updates.

The bug was found using an OOT target. I think the problem is quite
old, even if the particular intree target reproducer added in this
patch seem to pass when using LLVM 13.0.0.

[RKSimon]

llvm/test/CodeGen/AArch64/bbi-63928-ReplaceAllUsesOfValuesWith.ll

2 ↗

(On Diff #406267)

What's with the naming for bbi-63928-ReplaceAllUsesOfValuesWith.ll ?

[bjope]

llvm/test/CodeGen/AArch64/bbi-63928-ReplaceAllUsesOfValuesWith.ll

2 ↗

(On Diff #406267)

Oh, I should have adjusted that. The bbi-63928 happens to be our downstream TR reference. I can rename it to something more generic such as dag-ReplaceAllUsesOfValuesWith.ll.

[bjope]

Rename the new test case.

[bjope]

ping!

(let me know if I should try to find alternative reviewers)

[bjope]

Fixed clang-format failures.

[niravd]

Since there doesn't seem to be any other suggestions, I see no reason why we should hold off on this landing.

LGTM.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

void SelectionDAG::ReplaceAllUsesOfValueWith(SDValue From, SDValue To){
 
   // If we just RAUW'd the root, take note.
   if (From == getRoot())
     setRoot(To);
 }
 
 namespace {
 
-  /// UseMemo - This class is used by SelectionDAG::ReplaceAllUsesOfValuesWith
-  /// to record information about a use.
-  struct UseMemo {
-    SDNode *User;
-    unsigned Index;
-    SDUse *Use;
-  };
 
-  /// operator< - Sort Memos by User.
-  bool operator<(const UseMemo &L, const UseMemo &R) {
-    return (intptr_t)L.User < (intptr_t)R.User;
-  }
 
+/// RAUOVWUpdateListener - Helper for ReplaceAllUsesOfValuesWith - When the node
+/// pointed to by a UseMemo is deleted, set the User to nullptr to indicate that
+/// the node already has been taken care of recursively.
+class RAUOVWUpdateListener : public SelectionDAG::DAGUpdateListener {
+  SmallVector<UseMemo, 4> &Uses;
+
+  void NodeDeleted(SDNode *N, SDNode *E) override {
+    for (UseMemo &Memo : Uses)
+      if (Memo.User == N)
+        Memo.User = nullptr;
+  }
+
+public:
+  RAUOVWUpdateListener(SelectionDAG &d, SmallVector<UseMemo, 4> &uses)
+      : SelectionDAG::DAGUpdateListener(d), Uses(uses) {}
+};
+
 } // end anonymous namespace
 
 bool SelectionDAG::calculateDivergence(SDNode *N) {
   if (TLI->isSDNodeAlwaysUniform(N)) {
     assert(!TLI->isSDNodeSourceOfDivergence(N, FLI, DA) &&
            "Conflicting divergence information!");
     return false;
   }
for (SDNode::use_iterator UI = FromNode->use_begin(),
         UseMemo Memo = { *UI, i, &Use };
         Uses.push_back(Memo);
       }
     }
   }
 
   // Sort the uses, so that all the uses from a given User are together.
   llvm::sort(Uses);
+  RAUOVWUpdateListener Listener(*this, Uses);
 
   for (unsigned UseIndex = 0, UseIndexEnd = Uses.size();
        UseIndex != UseIndexEnd; ) {
     // We know that this user uses some value of From.  If it is the right
     // value, update it.
     SDNode *User = Uses[UseIndex].User;
+    // If the node has been deleted by recursive CSE updates when updating
+    // another node, then just skip this entry.
+    if (User == nullptr) {
+      ++UseIndex;
+      continue;
+    }
 
     // This node is about to morph, remove its old self from the CSE maps.
     RemoveNodeFromCSEMaps(User);
 
     // The Uses array is sorted, so all the uses for a given User
     // are next to each other in the list.
     // To help reduce the number of CSE recomputations, process all
     // the uses of this user that we can find this way.

llvm/test/CodeGen/AArch64/dag-ReplaceAllUsesOfValuesWith.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple aarch64-- -start-after codegenprepare -o - %s | FileCheck %s
+
+; REQUIRES: asserts
+
+; This used to hit an assertion like this:
+;
+; llc: ../lib/CodeGen/SelectionDAG/SelectionDAG.cpp:1087: bool llvm::SelectionDAG::RemoveNodeFromCSEMaps(llvm::SDNode*): Assertion `N->getOpcode() != ISD::DELETED_NODE && "DELETED_NODE in CSEMap!"' failed.
+; Stack dump:
+; 0.      Program arguments: llc -mtriple aarch64 -o - reduced.ll -start-after codegenprepare
+; 1.      Running pass 'Function Pass Manager' on module 'reduced.ll'.
+; 2.      Running pass 'AArch64 Instruction Selection' on function '@g'
+;  #0 0x00000000031615b8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int)
+;  #1 0x000000000315effe SignalHandler(int) Signals.cpp:0:0
+;  #2 0x00007f2c746b2630 __restore_rt sigaction.c:0:0
+;  #3 0x00007f2c7200f387 raise (/lib64/libc.so.6+0x36387)
+;  #4 0x00007f2c72010a78 abort (/lib64/libc.so.6+0x37a78)
+;  #5 0x00007f2c720081a6 __assert_fail_base (/lib64/libc.so.6+0x2f1a6)
+;  #6 0x00007f2c72008252 (/lib64/libc.so.6+0x2f252)
+;  #7 0x0000000002f06de9 llvm::SelectionDAG::RemoveNodeFromCSEMaps(llvm::SDNode*)
+;  #8 0x0000000002f0f0b4 llvm::SelectionDAG::ReplaceAllUsesOfValuesWith(llvm::SDValue const*, llvm::SDValue const*, unsigned int)
+;  #9 0x0000000002dc8a4f (anonymous namespace)::DAGCombiner::scalarizeExtractedVectorLoad(llvm::SDNode*, llvm::EVT, llvm::SDValue, llvm::LoadSDNode*) DAGCombiner.cpp:0:0
+; #10 0x0000000002de1a8e (anonymous namespace)::DAGCombiner::visitEXTRACT_VECTOR_ELT(llvm::SDNode*) DAGCombiner.cpp:0:0
+; #11 0x0000000002e12f41 (anonymous namespace)::DAGCombiner::visit(llvm::SDNode*) DAGCombiner.cpp:0:0
+; #12 0x0000000002e14fe5 (anonymous namespace)::DAGCombiner::combine(llvm::SDNode*) DAGCombiner.cpp:0:0
+
+define i64 @g({ i64, i64 }* %p) {
+; CHECK-LABEL: g:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ldr x8, [x0, #8]
+; CHECK-NEXT:    add x9, x8, x8
+; CHECK-NEXT:    add x8, x9, x8
+; CHECK-NEXT:    sub x0, x8, x8
+; CHECK-NEXT:    ret
+  %vecp = bitcast { i64, i64 }* %p to <2 x i64>*
+  %vec = load <2 x i64>, <2 x i64>* %vecp, align 1
+  %elt = extractelement <2 x i64> %vec, i32 1
+  %scalarp = getelementptr inbounds { i64, i64 }, { i64, i64 }* %p, i32 0, i32 1
+  %scalar = load i64, i64* %scalarp, align 1
+  %add.i62 = add i64 %elt, %scalar
+  %add.i66 = add i64 %add.i62, %elt
+  %add.i72 = add i64 %scalar, %scalar
+  %add.i76 = add i64 %add.i72, %elt
+  %add.i80 = add i64 %add.i76, 0
+  %sub.i82 = sub i64 %add.i66, %add.i80
+  ret i64 %sub.i82
+}