experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA/GPU] Fix a bug in CanShareOperandBufferWithUser().

Browse Source 
ElementsIn() supports only array shape; avoid it. We don't need to check the
sharing between operands and users for horizontal fusion for non-array shape
anyway.
This commit is contained in:
Trent Lo 2021-02-24 15:30:10 -08:00
parent b8c9a576c9
commit 941e049f0b
1 changed files with 2 additions and 1 deletions

3
tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
View File

@ -1445,7 +1445,8 @@ bool HloDataflowAnalysis::CanShareOperandBufferWithUser(
user->fused_parameter(user->operand_index(operand));
// We don't require the same dimensions but only the same number of elements
// and type (to make sure the same buffer size).
return ShapeUtil::ElementsIn(operand_subshape) ==
return operand_subshape.IsArray() && user_subshape.IsArray() &&
ShapeUtil::ElementsIn(operand_subshape) ==
ShapeUtil::ElementsIn(user_subshape) &&
ShapeUtil::SameElementType(operand_subshape, user_subshape) &&
AreTransitiveUsesEffectivelyElementwise(