diff options
author | 2018-07-17 18:01:52 -0700 | |
---|---|---|
committer | 2018-07-17 18:07:00 -0700 | |
commit | 07cc6474b219ee3ad9f55860e621f61b34bb6bd1 (patch) | |
tree | c09bbe69d49db3b6a6acd6f38545915ae2369035 /tensorflow/compiler/xla/service/hlo_module_group_metadata.cc | |
parent | 1662a105497e60d002e101161987cbbd48ba06c6 (diff) |
Add single-sided host send and receive operations.
Adds a bit on kSend/kReceive instructions and their Done variants indicated whether the operations communicates with the host or another device (the default). Host send/recv operations are single-sided without a complementary recv/send instruction in another module.
Host send/recv operations are exposed in the XLA builder API as SendToHost and RecvFromHost.
PiperOrigin-RevId: 205008138
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_module_group_metadata.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_module_group_metadata.cc | 10 |
1 files changed, 8 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index 3ffac2f413..10bf9ffd6c 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -20,6 +20,8 @@ limitations under the License. #include <utility> #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/hlo_casting_utils.h" +#include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/util.h" @@ -175,8 +177,12 @@ bool HloModuleGroupMetadata::IsChannelInstruction( case HloOpcode::kSend: case HloOpcode::kRecv: case HloOpcode::kSendDone: - case HloOpcode::kRecvDone: - return true; + case HloOpcode::kRecvDone: { + const HloSendRecvInstruction* send_recv_instr = + DynCast<HloSendRecvInstruction>(instruction); + CHECK(send_recv_instr != nullptr); + return !send_recv_instr->is_host_transfer(); + } default: return false; } |