[llvm] [SelectionDAG] Fix cycle in ReplacedValues during type legalization (PR #184697)

[Peter Rong via llvm-commits]

https://github.com/DataCorrupted created https://github.com/llvm/llvm-project/pull/184697

During type legalization, `ReplaceValueWith` handles morphed nodes by computing `OldValId` and `NewValId`, calling `ReplaceAllUsesOfValueWith` (RAUW), then adding `ReplacedValues[OldValId] = NewValId`. 

However, the RAUW call can trigger `CSE`, which fires NoteDeletion callbacks that add new  entries to `ReplacedValues`. If NoteDeletion adds `ReplacedValues[A] = B`, and the subsequent assignment adds `ReplacedValues[B] = A`, a cycle is created. RemapId then recurses infinitely following the cycle during  path compression, causing a stack overflow crash.

On AArch64, i16 is not a legal type and requires promotion to i32. When IR uses i16 values as `extractelement` indices that feed into other `extractelement` operations, the interplay between result promotion and operand promotion creates the specific CSE pattern that produces the cycle.

Fix by calling RemapId on both OldValId and NewValId after the  RAUW call to account for new ReplacedValues entries added during the RAUW cascade. In the reproducer, this causes both ids to resolve to the same value, skipping the assignment and preventing the cycle.

Fixes https://github.com/llvm/llvm-project/issues/57251

[Assisted-by](https://t.ly/Dkjjk): [Claude Opus 4.6](https://www.anthropic.com/news/claude-opus-4-6)

>From 40f29e092a193edfc0a013b164a4a7e5e477b70e Mon Sep 17 00:00:00 2001
From: Peter Rong <PeterRong at meta.com>
Date: Wed, 4 Mar 2026 14:39:32 -0800
Subject: [PATCH] [SelectionDAG] Fix cycle in ReplacedValues during type
 legalization

---
 .../lib/CodeGen/SelectionDAG/LegalizeTypes.cpp |  6 ++++++
 llvm/test/CodeGen/AArch64/Issue57251.ll        | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/Issue57251.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
index 9634e5afdc738..74b81572b6692 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
@@ -687,6 +687,12 @@ void DAGTypeLegalizer::ReplaceValueWith(SDValue From, SDValue To) {
           auto OldValId = getTableId(OldVal);
           auto NewValId = getTableId(NewVal);
           DAG.ReplaceAllUsesOfValueWith(OldVal, NewVal);
+          // Re-remap ids after RAUW, since the call above may have caused
+          // nodes to be deleted (via CSE), triggering NoteDeletion callbacks
+          // that added new entries to ReplacedValues. Without re-remapping,
+          // we could create a cycle like A -> B -> A.
+          RemapId(OldValId);
+          RemapId(NewValId);
           if (OldValId != NewValId)
             ReplacedValues[OldValId] = NewValId;
         }
diff --git a/llvm/test/CodeGen/AArch64/Issue57251.ll b/llvm/test/CodeGen/AArch64/Issue57251.ll
new file mode 100644
index 0000000000000..285f66b4856f9
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/Issue57251.ll
@@ -0,0 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 < %s | FileCheck %s
+
+define i16 @f(<4 x i16> %L) {
+; CHECK-LABEL: f:
+; CHECK:       // %bb.0: // %BB
+; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    ret
+BB:
+  %B3 = sub <4 x i16> %L, %L
+  %B1 = srem <4 x i16> %L, %L
+  %E = extractelement <4 x i16> %B3, i8 0
+  %I = insertelement <4 x i16> %B1, i16 %E, i16 0
+  %E1 = extractelement <4 x i16> %I, i16 %E
+  %E2 = extractelement <4 x i16> zeroinitializer, i16 %E1
+  %E3 = extractelement <4 x i16> %I, i16 %E2
+  ret i16 %E3
+}