[llvm-branch-commits] [llvm] release/19.x: [DAGCombiner] Fix ReplaceAllUsesOfValueWith mutation bug in visitFREEZE (#104924) (PR #105627)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-selectiondag

Author: None (llvmbot)

<details>
<summary>Changes</summary>

Backport 278fc8efdf004a1959a31bb4c208df5ee733d5c8

Requested by: @<!-- -->nikic

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+18-4) 
- (added) llvm/test/CodeGen/AArch64/dag-combine-freeze.ll (+31) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index aa9032ea2574c4..71cdec91e5f67a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -15680,13 +15680,16 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
     }
   }
 
-  SmallSetVector<SDValue, 8> MaybePoisonOperands;
-  for (SDValue Op : N0->ops()) {
+  SmallSet<SDValue, 8> MaybePoisonOperands;
+  SmallVector<unsigned, 8> MaybePoisonOperandNumbers;
+  for (auto [OpNo, Op] : enumerate(N0->ops())) {
     if (DAG.isGuaranteedNotToBeUndefOrPoison(Op, /*PoisonOnly*/ false,
                                              /*Depth*/ 1))
       continue;
     bool HadMaybePoisonOperands = !MaybePoisonOperands.empty();
-    bool IsNewMaybePoisonOperand = MaybePoisonOperands.insert(Op);
+    bool IsNewMaybePoisonOperand = MaybePoisonOperands.insert(Op).second;
+    if (IsNewMaybePoisonOperand)
+      MaybePoisonOperandNumbers.push_back(OpNo);
     if (!HadMaybePoisonOperands)
       continue;
     if (IsNewMaybePoisonOperand && !AllowMultipleMaybePoisonOperands) {
@@ -15698,7 +15701,18 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
   // it could create undef or poison due to it's poison-generating flags.
   // So not finding any maybe-poison operands is fine.
 
-  for (SDValue MaybePoisonOperand : MaybePoisonOperands) {
+  for (unsigned OpNo : MaybePoisonOperandNumbers) {
+    // N0 can mutate during iteration, so make sure to refetch the maybe poison
+    // operands via the operand numbers. The typical scenario is that we have
+    // something like this
+    //   t262: i32 = freeze t181
+    //   t150: i32 = ctlz_zero_undef t262
+    //   t184: i32 = ctlz_zero_undef t181
+    //   t268: i32 = select_cc t181, Constant:i32<0>, t184, t186, setne:ch
+    // When freezing the t181 operand we get t262 back, and then the
+    // ReplaceAllUsesOfValueWith call will not only replace t181 by t262, but
+    // also recursively replace t184 by t150.
+    SDValue MaybePoisonOperand = N->getOperand(0).getOperand(OpNo);
     // Don't replace every single UNDEF everywhere with frozen UNDEF, though.
     if (MaybePoisonOperand.getOpcode() == ISD::UNDEF)
       continue;
diff --git a/llvm/test/CodeGen/AArch64/dag-combine-freeze.ll b/llvm/test/CodeGen/AArch64/dag-combine-freeze.ll
new file mode 100644
index 00000000000000..4f0c3d0ce18006
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/dag-combine-freeze.ll
@@ -0,0 +1,31 @@
+; RUN: llc -mtriple aarch64 -o /dev/null %s
+
+; This used to fail with:
+;    Assertion `N1.getOpcode() != ISD::DELETED_NODE &&
+;               "Operand is DELETED_NODE!"' failed.
+; Just make sure we do not crash here.
+define void @test_fold_freeze_over_select_cc(i15 %a, ptr %p1, ptr %p2) {
+entry:
+  %a2 = add nsw i15 %a, 1
+  %sext = sext i15 %a2 to i32
+  %ashr = ashr i32 %sext, 31
+  %lshr = lshr i32 %ashr, 7
+  ; Setup an already frozen input to ctlz.
+  %freeze = freeze i32 %lshr
+  %ctlz = call i32 @llvm.ctlz.i32(i32 %freeze, i1 true)
+  store i32 %ctlz, ptr %p1, align 1
+  ; Here is another ctlz, which is used by a frozen select.
+  ; DAGCombiner::visitFREEZE will to try to fold the freeze over a SELECT_CC,
+  ; and when dealing with the condition operand the other SELECT_CC operands
+  ; will be replaced/simplified as well. So the SELECT_CC is mutated while
+  ; freezing the "maybe poison operands". This needs to be handled by
+  ; DAGCombiner::visitFREEZE, as it can't store the list of SDValues that
+  ; should be frozen in a separate data structure that isn't updated when the
+  ; SELECT_CC is mutated.
+  %ctlz1 = call i32 @llvm.ctlz.i32(i32 %lshr, i1 true)
+  %icmp = icmp ne i32 %lshr, 0
+  %select = select i1 %icmp, i32 %ctlz1, i32 0
+  %freeze1 = freeze i32 %select
+  store i32 %freeze1, ptr %p2, align 1
+  ret void
+}

``````````

</details>


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