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

Fix a bug in data flow analysis for asynchronous collective-permute.

Browse Source 
PiperOrigin-RevId: 316540956
Change-Id: Icb1fb9d1d445d5aa3cf7afa580eed06607c4ecb3
This commit is contained in:
Jinliang Wei 2020-06-15 14:08:55 -07:00 committed by TensorFlower Gardener
parent 0e40b3e0c3
commit fe6201c57b
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -698,7 +698,7 @@ bool HloDataflowAnalysis::UpdateCollectivePermuteDoneValueSet(
CHECK_EQ(collective_permute_done->opcode(),
HloOpcode::kCollectivePermuteDone);
bool changed = false;
// CollectivePermuteDone forwards the operand value at {0} to its output.
// CollectivePermuteDone forwards the operand value at {1} to its output.
const HloValueSet& operand_value_set =
GetValueSet(collective_permute_done->operand(0), {1});
HloValueSet& value_set = GetValueSet(collective_permute_done);
@ -945,6 +945,17 @@ Status HloDataflowAnalysis::InitializeInstructionValueSets() {
// CopyDone consumes a tuple produced by CopyStart and produces an
// element. Its output aliases its input tuple element {0}.
break;
case HloOpcode::kCollectivePermuteStart:
// CollectivePermuteStart produces a tuple of
// {aliased operand, destination buffer, U32 context, U32 context}.
define_value_at(/*index=*/{});
define_value_at(/*index=*/{1});
define_value_at(/*index=*/{2});
define_value_at(/*index=*/{3});
break;
case HloOpcode::kCollectivePermuteDone:
// CollectivePermuteDone's output aliases its input tuple element {1}.
break;
case HloOpcode::kRecvDone:
// RecvDone produces a two-element tuple. Element zero aliases its
// input tuple element {0}; element one is a token.