[llvm] Make RAUW skip instances where operand is the replacement node (PR #212724)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-selectiondag

Author: Soham Karandikar (skadewdl3)

<details>
<summary>Changes</summary>

Fixes #<!-- -->211496.

I used `llvm-reduce` to reduce the reproducer from the linked issue to this:
```llvm
; reduced.ll
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

define i32 @<!-- -->f18() #<!-- -->0 {
entry:
  %0 = load <15 x i64>, ptr null, align 128
  %shuffle = shufflevector <15 x i64> %0, <15 x i64> zeroinitializer, <16 x i32> <i32 7, i32 6, i32 4, i32 12, i32 9, i32 5, i32 3, i32 5, i32 12, i32 13, i32 0, i32 14, i32 7, i32 8, i32 4, i32 2>
  %shuffle1 = shufflevector <16 x i64> %shuffle, <16 x i64> zeroinitializer, <16 x i32> <i32 1, i32 5, i32 4, i32 2, i32 7, i32 13, i32 3, i32 2, i32 5, i32 10, i32 8, i32 4, i32 0, i32 8, i32 7, i32 6>
  store <16 x i64> %shuffle1, ptr null, align 128
  ret i32 0
}

attributes #<!-- -->0 = { "target-features"="+avx2" }
```
Running `llc reduced.ll -o reduced` through a LLDB to remove the token factor node `t233` from the CSE map.

---


Error call stack: `[a bunch of selection dag machinery] -> combineBROADCAST_LOAD -> makeEquivalentMemoryOrdering -> ReplaceAllUsesOfValueWith -> RemoveNodeFromCSEMaps`.

- `combineBROADCAST_LOAD` calls `makeEquivalentMemoryOrdering` with `t261:1` and `t193:1` as the old and new chains respectively.
- `makeEquivalentMemoryOrdering` creates a token factor calls `ReplaceAllUsesOfValueWith` to replace `t261:1` with the token factor `t233`.
- `ReplaceAllUsesOfValueWith` does `Use.set(To)`, which replaces the uses of `t261:1` with `t233`.
  - Before `combineBROADCAST_LOAD` runs, lots of combine and RAUW operations had run which results in this node: `t194 = TokenFactor t261:1, t193:1`. Since `t194` is also a user f `t261:1`, it too got rewritten to `t194 = TokenFactor t233, t193:1`.
  - Since `t233` is a user of `t261` too, a cycle is introduced in the DAG: `t233 = TokenFactor t233, t193:1`.
  - CSE then attempted to merge `t233` into `t194`. Recursive RAUW then encountered t233 as its own user and attempted to remove it from the CSE map after it had already been removed, producing the error from the issue: `Node is not in map!`.

`makeEquivalentMemoryOrdering` does call `UpdateNodeOperands` afterward to restore the `TokenFactor`'s original operands. In this case, however, the CSE merge and the error occur inside `ReplaceAllUsesOfValueWith`, so execution never reaches that repair.

This PR adds a check to `ReplaceAllUsesOfValueWith` to the skip replacement of uses if the user and the replacement node are the same.

I'm quite new to x86 backend stuff and SelectionDAG in general, so I'd appreciate any guidance on whether this fix is appropriate, whether the test is in the right place, etc.

---
Full diff: https://github.com/llvm/llvm-project/pull/212724.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+12) 
- (added) llvm/test/CodeGen/X86/dagcombine-tokenfactor-cse-crash.ll (+18) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 679fa3fe36e27..cb79b69d9f1f3 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -13301,6 +13301,18 @@ void SelectionDAG::ReplaceAllUsesOfValueWith(SDValue From, SDValue To){
   RAUWUpdateListener Listener(*this, UI, UE);
   while (UI != UE) {
     SDNode *User = UI->getUser();
+    
+    // Don't replace uses of the To node if is is self-referential.
+    // If CSE tries to merge the Use with another node post-replacement
+    // it'll try to remove the same node from the CSE map twice, causing a crash.
+    if (User == To.getNode()) {
+      do {
+        ++UI;
+      } while (UI != UE && UI->getUser() == User);
+      continue;
+    }
+
+
     bool UserRemovedFromCSEMaps = false;
 
     // A user can appear in a use list multiple times, and when this
diff --git a/llvm/test/CodeGen/X86/dagcombine-tokenfactor-cse-crash.ll b/llvm/test/CodeGen/X86/dagcombine-tokenfactor-cse-crash.ll
new file mode 100644
index 0000000000000..5e9eee133ad6a
--- /dev/null
+++ b/llvm/test/CodeGen/X86/dagcombine-tokenfactor-cse-crash.ll
@@ -0,0 +1,18 @@
+; RUN: llc < %s -mtriple=x86_64 -mattr=+avx2 -o /dev/null
+;
+; Verify that CSE-reusing a TokenFactor during broadcast-load combining does not
+; introduce a self-reference while replacing the old chain.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @f18() #0 {
+entry:
+  %0 = load <15 x i64>, ptr null, align 128
+  %shuffle = shufflevector <15 x i64> %0, <15 x i64> zeroinitializer, <16 x i32> <i32 7, i32 6, i32 4, i32 12, i32 9, i32 5, i32 3, i32 5, i32 12, i32 13, i32 0, i32 14, i32 7, i32 8, i32 4, i32 2>
+  %shuffle1 = shufflevector <16 x i64> %shuffle, <16 x i64> zeroinitializer, <16 x i32> <i32 1, i32 5, i32 4, i32 2, i32 7, i32 13, i32 3, i32 2, i32 5, i32 10, i32 8, i32 4, i32 0, i32 8, i32 7, i32 6>
+  store <16 x i64> %shuffle1, ptr null, align 128
+  ret i32 0
+}
+
+attributes #0 = { "target-features"="+avx2" }

``````````

</details>


https://github.com/llvm/llvm-project/pull/212724