| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | computation_placer.cc | 42 proto->set_replica_count(replica_count()); in Serialize()
47 for (int replica = 0; replica < replica_count(); ++replica) { in Serialize()
57 if (proto.replica_count() <= 0 || proto.computation_count() <= 0) { in Deserialize()
61 proto.replica_count(), proto.computation_count()); in Deserialize()
64 proto.replica_count(), proto.computation_count()); in Deserialize()
69 proto.replica_count()); in Deserialize()
70 for (int replica = 0; replica < proto.replica_count(); ++replica) { in Deserialize()
80 " Replicas: ", replica_count(), "\n"); in ToString()
83 for (int replica = 0; replica < replica_count(); ++replica) { in ToString()
92 int replica_count, in DeviceId() argument
[all …]
|
| D | computation_placer.h | 41 DeviceAssignment(int replica_count, int computation_count) in DeviceAssignment() argument
42 : Array2D<int>(replica_count, computation_count, -1) { in DeviceAssignment()
43 CHECK_GT(replica_count, 0); in DeviceAssignment()
47 int replica_count() const { return height(); } in replica_count() function
73 int replica_count, int computation_count);
77 virtual StatusOr<DeviceAssignment> AssignDevices(int replica_count,
|
| D | hlo_module_config.h | 80 void set_replica_count(int64 replica_count) { in set_replica_count() argument
81 replica_count_ = replica_count; in set_replica_count()
83 int64 replica_count() const { return replica_count_; } in replica_count() function
|
| D | hlo_module_config.cc | 56 if (replica_count() != 1) { in compilation_cache_key()
57 StrAppend(&key, "::replica_count=", replica_count()); in compilation_cache_key()
|
| D | service.cc | 762 const int64 replica_count = options_.number_of_replicas(); in GetDeviceHandles() local
763 if (replica_count <= 0) { in GetDeviceHandles()
766 if (available_device_count < arg->device_count() * replica_count) { in GetDeviceHandles()
770 arg->device_count(), replica_count, available_device_count); in GetDeviceHandles()
1011 const int64 replica_count = options_.number_of_replicas(); in TransferToInfeed() local
1012 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferToInfeed()
1017 replica_count, ").")); in TransferToInfeed()
1040 const int64 replica_count = options_.number_of_replicas(); in TransferFromOutfeed() local
1041 if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { in TransferFromOutfeed()
1044 arg->replica_id(), replica_count); in TransferFromOutfeed()
|
| D | elemental_ir_emitter.cc | 2441 if (hlo_module_config_.replica_count() != 1) { in MakeElementGenerator()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | nccl_all_reduce_thunk.cc | 48 int64 replica_count; member
68 replica_count, element_count, device_ordinal, generation_counter, in ToString()
151 return participants_.size() >= participant.replica_count; in SubmitParticipant()
164 (participants_.back().replica_count != participant.replica_count || in SubmitParticipant()
207 done_.emplace(participant.replica_count); in SubmitParticipant()
320 participant.replica_count = replica_count_; in ExecuteOnStream()
345 int64 replica_count, int64 element_count, in NcclAllReduceThunk() argument
350 replica_count_(replica_count), in NcclAllReduceThunk()
|
| D | nccl_all_reduce_thunk.h | 43 NcclAllReduceThunk(int64 replica_count, int64 element_count,
|
| D | ir_emitter_unnested.cc | 1390 VLOG(2) << "AllReduce; replica count: " << hlo_module_config_.replica_count() in HandleAllReduce()
1397 hlo_module_config_.replica_count() > 1 && in HandleAllReduce()
1408 /*replica_count=*/hlo_module_config_.replica_count(), in HandleAllReduce()
1415 if (hlo_module_config_.replica_count() != 1) { in HandleAllReduce()
1420 hlo_module_config_.replica_count(), crs->operand_count(), in HandleAllReduce()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | client_library.cc | 93 int replica_count = options.number_of_replicas(); in GetOrCreateLocalClient() local
108 service_options.set_number_of_replicas(replica_count); in GetOrCreateLocalClient()
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla_client.py | 1828 def initialize_replica_count(replica_count): argument
1842 _default_replica_count = replica_count
|
| D | local_computation_builder.cc | 218 int num_replicas = device_assignment_.replica_count(); in DeviceOrdinals()
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | xla_data.proto | 343 int32 replica_count = 1; field
346 // Each logical computation runs on replica_count physical devices.
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | ir_emitter.cc | 1310 if (hlo_module_config_.replica_count() != 1) { in HandleAllReduce()
|