return device;
}
+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) {
tensorflow::gtl::optional<int64> GetInstructionDevice(
const HloInstruction& instruction) const;
+ // Returns the number of modules for devices (excluding the host module).
+ int64 GetDeviceModulesCount() const;
+
// Returns the companion instructions for the given instruction.
//
// Precondition: IsCompanionWhile(instruction) is true.
// profiled.
std::map<int64, se::Stream*> index_to_profiled_streams;
- TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment,
- backend->computation_placer()->AssignDevices(
- options_.number_of_replicas(), executables.size()));
+ // Build DeviceAssignment for all cores based on the provided device handles.
+ DeviceAssignment device_assignment(options_.number_of_replicas(),
+ executables.size());
+ for (int64 i = 0; i < executables.size(); i++) {
+ TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handles[i]));
+ CHECK_EQ(replicas.size(), arguments[i].size());
+ for (int64 replica = 0; replica < replicas.size(); ++replica) {
+ device_assignment(replica, i) = replicas[replica]->device_ordinal();
+ }
+ }
for (int64 i = 0; i < executables.size(); i++) {
// Stream executors for the replicas of the current computation.