[llvm] 23bc8f7 - [DAGCombiner] `visitFREEZE()`: allow, and update, other uses of maybe-poison operand

[Roman Lebedev via llvm-commits]

Author: Roman Lebedev
Date: 2022-12-22T23:23:19+03:00
New Revision: 23bc8f730d99960c93946286f03d97d2fc54e8fe

URL: https://github.com/llvm/llvm-project/commit/23bc8f730d99960c93946286f03d97d2fc54e8fe
DIFF: https://github.com/llvm/llvm-project/commit/23bc8f730d99960c93946286f03d97d2fc54e8fe.diff

LOG: [DAGCombiner] `visitFREEZE()`: allow, and update, other uses of maybe-poison operand

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
    llvm/test/CodeGen/X86/freeze-binary.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 390ee7eec9767..c62de472e66a1 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -14261,8 +14261,7 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
       if (DAG.isGuaranteedNotToBeUndefOrPoison(Op, /*PoisonOnly*/ false,
                                                /*Depth*/ 1))
         continue;
-      if ((!MaybePoisonOperand && N0->isOnlyUserOf(Op.getNode())) ||
-          MaybePoisonOperand == Op) {
+      if (!MaybePoisonOperand || MaybePoisonOperand == Op) {
         MaybePoisonOperand = Op;
         continue;
       }
@@ -14271,15 +14270,17 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
       break;
     }
     if (MaybePoisonOperand) {
-      // Recreate the node with the frozen maybe-poison operand.
-      // TODO: Drop the isOnlyUserOf constraint and replace all users of
-      // MaybePoisonOperand with FrozenMaybePoisonOperand
-      // to match pushFreezeToPreventPoisonFromPropagating behavior.
+      // First, freeze the offending operand.
       SDValue FrozenMaybePoisonOperand = DAG.getFreeze(MaybePoisonOperand);
+      // Then, change all other uses of unfrozen operand to use frozen operand.
+      DAG.ReplaceAllUsesOfValueWith(MaybePoisonOperand,
+                                    FrozenMaybePoisonOperand);
+      // But, that also updated the use in the freeze we just created, thus
+      // creating a cycle in a DAG. Let's undo that by mutating the freeze.
+      DAG.UpdateNodeOperands(FrozenMaybePoisonOperand.getNode(),
+                             MaybePoisonOperand);
+      // Finally, recreate the node with the frozen maybe-poison operand.
       SmallVector<SDValue> Ops(N0->op_begin(), N0->op_end());
-      for (SDValue &Op : Ops)
-        if (Op == MaybePoisonOperand)
-          Op = FrozenMaybePoisonOperand;
       // TODO: Just strip poison generating flags?
       SDValue R = DAG.getNode(N0.getOpcode(), SDLoc(N0), N0->getVTList(), Ops);
       assert(DAG.isGuaranteedNotToBeUndefOrPoison(R, /*PoisonOnly*/ false) &&

diff  --git a/llvm/test/CodeGen/X86/freeze-binary.ll b/llvm/test/CodeGen/X86/freeze-binary.ll
index 0f3e2fbe662b8..339f1f9939709 100644
--- a/llvm/test/CodeGen/X86/freeze-binary.ll
+++ b/llvm/test/CodeGen/X86/freeze-binary.ll
@@ -26,7 +26,6 @@ define i32 @freeze_and_extra_use(i32 %a0, ptr %escape) nounwind {
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    movl %eax, (%ecx)
-; X86-NEXT:    andl $15, %eax
 ; X86-NEXT:    andl $7, %eax
 ; X86-NEXT:    retl
 ;
@@ -34,7 +33,6 @@ define i32 @freeze_and_extra_use(i32 %a0, ptr %escape) nounwind {
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
 ; X64-NEXT:    movl %edi, (%rsi)
-; X64-NEXT:    andl $15, %eax
 ; X64-NEXT:    andl $7, %eax
 ; X64-NEXT:    retq
   store i32 %a0, ptr %escape
@@ -46,17 +44,14 @@ define i32 @freeze_and_extra_use(i32 %a0, ptr %escape) nounwind {
 define i32 @freeze_and_extra_use2(i32 %a0, ptr %escape) nounwind {
 ; X86-LABEL: freeze_and_extra_use2:
 ; X86:       # %bb.0:
-; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; X86-NEXT:    movl %ecx, %eax
+; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    andl $15, %eax
-; X86-NEXT:    andl %ecx, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_and_extra_use2:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl %edi, %eax
 ; X64-NEXT:    andl $15, %eax
-; X64-NEXT:    andl %edi, %eax
 ; X64-NEXT:    retq
   %x = and i32 %a0, 15
   %y = freeze i32 %x