Skip to content
Snippets Groups Projects
Commit 85dbdae5 authored by Aden Grue's avatar Aden Grue Committed by TensorFlower Gardener
Browse files

Don't allow multi-output fusion when the "producer" op is in-place

Fusion in this case is dangerous. Without in-depth analysis of the consumer we can't know whether or not it will treat the producer's in-place operands safely.

PiperOrigin-RevId: 432523155
parent 74d537e3
No related merge requests found
......@@ -254,6 +254,21 @@ bool IsProducerConsumerMultiOutputFusible(const HloInstruction& producer,
return false;
}
// If the producer is an in-place operation (or a fusion with an in-place
// operation at its root), fusion is dangerous. Only certain data access
// patterns in the consumer are valid, otherwise there is danger of the
// consumer accidentally reading parts of the in-place operand that have
// already been modified.
// For simplicitly, we simply ban fusion altogether when the producer is
// in-place. (It would be possible to make this more lenient in the future by
// allowing fusion in certain cases where the consumer can be proven to be
// 'safe').
if (HloDataflowAnalysis::IsInPlaceOperation(producer.opcode()) ||
(producer.opcode() == HloOpcode::kFusion &&
HloDataflowAnalysis::IsInPlaceOperation(
producer.fused_expression_root()->opcode()))) {
return false;
}
if (!IsLoopFusible(producer) || !IsFusibleAsMultiOutputFusionRoot(consumer)) {
return false;
}
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment