diff options
| author |  HyoukJoong Lee <hyouklee@google.com> | 2018-04-30 12:49:33 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-30 12:53:12 -0700 |
| commit | 5cdcb47361e9923c418c16fee6510a472a928427 (patch) | |
| tree | 3954662bd0c2ee98b5e2a8584ecc9590d6bafcb4 /tensorflow/compiler/xla/service/hlo_module_group_metadata.cc | |
| parent | 6e9d8abcdc44552a53475405f6cf0fdbffb40613 (diff) | |
Fix device assignment in xla/service/service.cc to build the assignment based on
the provided device handles rather than using the default assignment.
PiperOrigin-RevId: 194829761
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_module_group_metadata.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_module_group_metadata.cc | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index 54c34ce116..3367d76ded 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -194,6 +194,13 @@ int64 HloModuleGroupMetadata::GetModuleId(const HloModule* module) const { LOG(FATAL) << "unknown module"; } +int64 HloModuleGroupMetadata::GetDeviceModulesCount() const { + return std::count_if(modules_.begin(), modules_.end(), + [](const HloModule* module) { + return !module->config().is_host_module(); + }); +} + Status HloModuleGroupMetadata::RecordInstructions() { const auto visitor = [this](HloInstruction* hlo) -> Status { if (hlo->opcode() == HloOpcode::kWhile) { |