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

Change _Send/_Recv attrs in VirtualScheduler

Browse Source 
Change src_device_ to send_device and dst_device_ to recv_device. This complies
with tensorflow naming, so that VirtualScheduler can handle graphs generated
on inspectz with _Send/_Recv nodes and AutoGrappler does not need to remove
them.

PiperOrigin-RevId: 221660625
This commit is contained in:
Andiry Xu 2018-11-15 11:37:50 -08:00 committed by TensorFlower Gardener
parent dcf390ede4
commit 4ea51ae99c
2 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/grappler/costs/virtual_scheduler.cc
View File

@ -639,6 +639,8 @@ std::pair<const NodeDef*, const NodeDef*> VirtualScheduler::CreateSendRecv(
send->set_device(ChannelDeviceName(from, to));
auto& send_attr = *(send->mutable_attr());
send_attr[kAttrInputSrc].set_s(input_name);
// Use input_name as tensor_name.
send_attr[kAttrTensorName].set_s(input_name);
send_attr[kAttrSrcDevice].set_s(DeviceName(from));
send_attr[kAttrDstDevice].set_s(DeviceName(to));
@ -650,6 +652,8 @@ std::pair<const NodeDef*, const NodeDef*> VirtualScheduler::CreateSendRecv(
recv->set_device(DeviceName(to));
auto& recv_attr = *(recv->mutable_attr());
recv_attr[kAttrInputSrc].set_s(input_name);
// Use input_name as tensor_name.
recv_attr[kAttrTensorName].set_s(input_name);
// NodeState for _Send op.
auto& send_node_state = GetNodeStateOrCreateIt(send);

5
tensorflow/core/grappler/costs/virtual_scheduler.h
View File

@ -308,8 +308,9 @@ class VirtualScheduler {
private:
// Constants.
const string kAttrInputSrc = "input_source_";
const string kAttrSrcDevice = "src_device_";
const string kAttrDstDevice = "dst_device_";
const string kAttrSrcDevice = "send_device";
const string kAttrDstDevice = "recv_device";
const string kAttrTensorName = "tensor_name";
const string kChannelDevice = "Channel";
// Methods called from Init(). Fails if initialize_ is set.