[llvm] 8623da5 - [DAG] visitFREEZE - generalize freeze(op()) -> op(freeze()) to any number of operands

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2022-08-10T13:12:46+01:00
New Revision: 8623da5f748bc36fc78fa148fd34332d386533ea

URL: https://github.com/llvm/llvm-project/commit/8623da5f748bc36fc78fa148fd34332d386533ea
DIFF: https://github.com/llvm/llvm-project/commit/8623da5f748bc36fc78fa148fd34332d386533ea.diff

LOG: [DAG] visitFREEZE - generalize freeze(op()) -> op(freeze()) to any number of operands

canCreateUndefOrPoison currently only handles unary ops, but we intend to change that soon - this more closely matches the pushFreezeToPreventPoisonFromPropagating behaviour where the freeze is pushed up to a single operand value, as long as all others are guaranteed not to be poison/undef.

However, pushFreezeToPreventPoisonFromPropagating would freeze all uses of the value - whilst this variant requires the frozen value to be only used in the op - we can look at generalize multiple uses later if the need arises.

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 7ba494033262..392b24af3716 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -13868,13 +13868,40 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
   if (DAG.isGuaranteedNotToBeUndefOrPoison(N0, /*PoisonOnly*/ false))
     return N0;
 
-  // Fold freeze(unaryop(x)) -> unaryop(freeze(x)).
-  // TODO: Replace with pushFreezeToPreventPoisonFromPropagating fold and
-  // support getNumOperands() >= 1.
-  if (N0.getNumOperands() == 1 &&
-      !DAG.canCreateUndefOrPoison(N0, /*PoisonOnly*/ false) && N0->hasOneUse())
-    return DAG.getNode(N0.getOpcode(), SDLoc(N0), N->getValueType(0),
-                       DAG.getFreeze(N0.getOperand(0)));
+  // 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) &&
+      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: 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);
+    }
+  }
 
   return SDValue();
 }