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 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(); }