[DAG] Strip poison generating flags in freeze(op()) -> op(freeze()) fold

[RKSimon]

Summary

This patch follows the InstCombine approach of stripping poison generating flags (nsw/nuw from add/sub etc.) to allow us to push a freeze() through the op. Unlike InstCombine it doesn't retain any flags, but we have plenty of DAG folds that do the same thing already. We assert that the newly generated op isGuaranteedNotToBeUndefOrPoison.

Similar to the ValueTracking approach, isGuaranteedNotToBeUndefOrPoison has been updated to confirm that if an op can't create undef/poison and its operands are guaranteed not to be undef/poison - then its not undef/poison. This is just for the generic opcodes - target specific opcodes will need to do this manually just in case they have some special cases.

[craig.topper]

It's coaser than InstCombine

coaser?

[RKSimon]

In D132333#3738272, @craig.topper wrote:

It's coaser than InstCombine

coaser?

this is what happens when spellcheck fails :)

[RKSimon]

ping?

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitFREEZE(SDNode *N) {
 
   // Fold freeze(op(x,y,z)) -> op(freeze(x),y,z).
   // Try to push freeze through instructions that propagate but don't produce
   // poison as far as possible. If an operand of freeze follows three
   // conditions 1) one-use, 2) does not produce poison, and 3) has all but one
   // guaranteed-non-poison operands then push the freeze through to the one
   // operand that is not guaranteed non-poison.
   if (!DAG.canCreateUndefOrPoison(N0, /*PoisonOnly*/ false,
-                                  /*ConsiderFlags*/ true) &&
+                                  /*ConsiderFlags*/ false) &&
       N0->getNumValues() == 1 && N0->hasOneUse()) {
     SDValue MaybePoisonOperand;
     for (SDValue Op : N0->ops()) {
       if (DAG.isGuaranteedNotToBeUndefOrPoison(Op, /*PoisonOnly*/ false))
         continue;
       if ((!MaybePoisonOperand && N0->isOnlyUserOf(Op.getNode())) ||
           MaybePoisonOperand == Op) {
         MaybePoisonOperand = Op;
         continue;
       }
       // Multiple maybe-poison ops - bail out.
       MaybePoisonOperand = SDValue();
       break;
     }
     if (MaybePoisonOperand) {
       // Recreate the node with the frozen maybe-poison operand.
-      // TODO: Disable ConsiderFlags and just strip poison generating flags?
       // TODO: Drop the isOnlyUserOf constraint and replace all users of
       // MaybePoisonOperand with FrozenMaybePoisonOperand
       // to match pushFreezeToPreventPoisonFromPropagating behavior.
       SDValue FrozenMaybePoisonOperand = DAG.getFreeze(MaybePoisonOperand);
       SmallVector<SDValue> Ops(N0->op_begin(), N0->op_end());
       for (SDValue &Op : Ops)
         if (Op == MaybePoisonOperand)
           Op = FrozenMaybePoisonOperand;
-      return DAG.getNode(N0.getOpcode(), SDLoc(N0), N0->getVTList(), Ops);
+      // TODO: Just strip poison generating flags?
+      SDValue R = DAG.getNode(N0.getOpcode(), SDLoc(N0), N0->getVTList(), Ops);
+      assert(DAG.isGuaranteedNotToBeUndefOrPoison(R, /*PoisonOnly*/ false) &&
+             "Can't create node that may be undef/poison!");
+      return R;
     }
   }
 
   return SDValue();
 }
 
 /// We know that BV is a build_vector node with Constant, ConstantFP or Undef
 /// operands. DstEltVT indicates the destination element value type.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

for (unsigned i = 0, e = Op.getNumOperands(); i < e; ++i) {
       if (!DemandedElts[i])
         continue;
       if (!isGuaranteedNotToBeUndefOrPoison(Op.getOperand(i), PoisonOnly,
                                             Depth + 1))
         return false;
     }
     return true;
 
     // TODO: Search for noundef attributes from library functions.
 
     // TODO: Pointers dereferenced by ISD::LOAD/STORE ops are noundef.
 
   default:
     // Allow the target to implement this method for its nodes.
     if (Opcode >= ISD::BUILTIN_OP_END || Opcode == ISD::INTRINSIC_WO_CHAIN ||
         Opcode == ISD::INTRINSIC_W_CHAIN || Opcode == ISD::INTRINSIC_VOID)
       return TLI->isGuaranteedNotToBeUndefOrPoisonForTargetNode(
           Op, DemandedElts, *this, PoisonOnly, Depth);
     break;
   }
 
-  return false;
+  // If Op can't create undef/poison and none of its operands are undef/poison
+  // then Op is never undef/poison.
+  // NOTE: TargetNodes should handle this in themselves in
+  // isGuaranteedNotToBeUndefOrPoisonForTargetNode.
+  return !canCreateUndefOrPoison(Op, PoisonOnly, /*ConsiderFlags*/ true,
+                                 Depth) &&
+         all_of(Op->ops(), [&](SDValue V) {
+           return isGuaranteedNotToBeUndefOrPoison(V, PoisonOnly, Depth + 1);
+         });
 }
 
 bool SelectionDAG::canCreateUndefOrPoison(SDValue Op, bool PoisonOnly,
                                           bool ConsiderFlags,
                                           unsigned Depth) const {
   // TODO: Assume we don't know anything for now.
   EVT VT = Op.getValueType();
   if (VT.isScalableVector())

llvm/test/CodeGen/X86/freeze-binary.ll

; X64-NEXT:    retq
   %z = add i32 %y, 1
   ret i32 %z
 }
 
 define i32 @freeze_add_nsw(i32 %a0) nounwind {
 ; X86-LABEL: freeze_add_nsw:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    incl %eax
-; X86-NEXT:    incl %eax
+; X86-NEXT:    addl $2, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_add_nsw:
 ; X64:       # %bb.0:
 ; X64-NEXT:    # kill: def $edi killed $edi def $rdi
-; X64-NEXT:    leal 1(%rdi), %eax
-; X64-NEXT:    incl %eax
+; X64-NEXT:    leal 2(%rdi), %eax
 ; X64-NEXT:    retq
   %x = add nsw i32 %a0, 1
   %y = freeze i32 %x
   %z = add i32 %y, 1
   ret i32 %z
 }
 
 define <4 x i32> @freeze_add_vec(<4 x i32> %a0) nounwind {
; X64-NEXT:    retq
   %z = mul i32 %y, 2
   ret i32 %z
 }
 
 define i32 @freeze_mul_nsw(i32 %a0) nounwind {
 ; X86-LABEL: freeze_mul_nsw:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    leal (%eax,%eax,2), %eax
 ; X86-NEXT:    leal (%eax,%eax,4), %eax
+; X86-NEXT:    leal (%eax,%eax,2), %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_mul_nsw:
 ; X64:       # %bb.0:
 ; X64-NEXT:    # kill: def $edi killed $edi def $rdi
-; X64-NEXT:    leal (%rdi,%rdi,2), %eax
-; X64-NEXT:    leal (%rax,%rax,4), %eax
+; X64-NEXT:    leal (%rdi,%rdi,4), %eax
+; X64-NEXT:    leal (%rax,%rax,2), %eax
 ; X64-NEXT:    retq
   %x = mul nsw i32 %a0, 3
   %y = freeze i32 %x
   %z = mul i32 %y, 5
   ret i32 %z
 }
 
 define <8 x i16> @freeze_mul_vec(<8 x i16> %a0) nounwind {
; X64-NEXT:    retq
   %z = shl i32 %y, 2
   ret i32 %z
 }
 
 define i32 @freeze_shl_nsw(i32 %a0) nounwind {
 ; X86-LABEL: freeze_shl_nsw:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    shll $3, %eax
-; X86-NEXT:    shll $5, %eax
+; X86-NEXT:    shll $8, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_shl_nsw:
 ; X64:       # %bb.0:
-; X64-NEXT:    # kill: def $edi killed $edi def $rdi
-; X64-NEXT:    leal (,%rdi,8), %eax
-; X64-NEXT:    shll $5, %eax
+; X64-NEXT:    movl %edi, %eax
+; X64-NEXT:    shll $8, %eax
 ; X64-NEXT:    retq
   %x = shl nsw i32 %a0, 3
   %y = freeze i32 %x
   %z = shl i32 %y, 5
   ret i32 %z
 }
 
 define i32 @freeze_shl_outofrange(i32 %a0) nounwind {
; X64-NEXT:    retq
   %y = freeze <2 x i64> %x
   %z = shl <2 x i64> %y, <i64 2, i64 1>
   ret <2 x i64> %z
 }
 
 define <2 x i64> @freeze_shl_vec_outofrange(<2 x i64> %a0) nounwind {
 ; X86-LABEL: freeze_shl_vec_outofrange:
 ; X86:       # %bb.0:
-; X86-NEXT:    paddq %xmm0, %xmm0
-; X86-NEXT:    psllq $2, %xmm0
+; X86-NEXT:    psllq $3, %xmm0
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: freeze_shl_vec_outofrange:
 ; X64:       # %bb.0:
 ; X64-NEXT:    vpsllvq {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
 ; X64-NEXT:    vpsllq $2, %xmm0, %xmm0
 ; X64-NEXT:    retq
   %x = shl <2 x i64> %a0, <i64 1, i64 64>