Home
last modified time | relevance | path

Searched refs:num_local_devices (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/core/nccl/
Dnccl_manager.cc171 num_local_devices(num_local_devices_in), in Collective()
192 const int num_local_devices; // devices local to this node member
313 for (i = 0; i < collective->num_local_devices; ++i) { in GetCommunicator()
319 if (i == collective->num_local_devices) { in GetCommunicator()
352 std::vector<CommunicatorMember> members(collective->num_local_devices); in GetCommunicator()
353 std::vector<int> devices(collective->num_local_devices); in GetCommunicator()
354 for (int i = 0; i < collective->num_local_devices; ++i) { in GetCommunicator()
390 std::vector<ncclComm_t> nccl_comms(collective->num_local_devices); in GetCommunicator()
403 for (int i = 0; i < collective->num_local_devices; ++i) { in GetCommunicator()
420 nccl_comms.data(), collective->num_local_devices, devices.data())); in GetCommunicator()
[all …]
Dnccl_manager.h134 Context(const string& collective_key, int num_local_devices, in Context()
138 num_local_devices(num_local_devices), in Context()
147 int num_local_devices; member
Dcollective_communicator.cc74 const int num_local_devices = col_params->group.num_devices_per_task.at( in Enqueue() local
107 nccl_collective_key, num_local_devices, num_global_devices, in Enqueue()
113 << " num local devices " << num_local_devices in Enqueue()
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client.cc672 const int num_local_devices = local_devices_.size(); in ExecuteOnLocalDevices() local
674 if (argument_handles.size() != num_local_devices) { in ExecuteOnLocalDevices()
678 argument_handles.size(), num_local_devices, num_replicas(), in ExecuteOnLocalDevices()
684 << " num_local_devices=" << num_local_devices; in ExecuteOnLocalDevices()
691 for (int i = 0; i < num_local_devices; ++i) { in ExecuteOnLocalDevices()
701 std::vector<ExecuteResult> results(num_local_devices); in ExecuteOnLocalDevices()
709 for (int i = 0; i < num_local_devices; ++i) { in ExecuteOnLocalDevices()
724 execute_semaphore.Acquire(num_local_devices); in ExecuteOnLocalDevices()
726 for (int i = 0; i < num_local_devices; ++i) { in ExecuteOnLocalDevices()
742 num_local_devices); in ExecuteOnLocalDevices()
[all …]
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtpu_client.cc105 int num_local_devices = host.Cores(kTensorCore).size(); in GetDefaultDeviceAssignment() local
106 if (num_replicas * num_partitions <= num_local_devices) { in GetDefaultDeviceAssignment()
/external/tensorflow/tensorflow/core/distributed_runtime/
Dmaster.cc451 int num_local_devices = 0; in CreateSession() local
454 if (num_local_devices == 0) { in CreateSession()
458 num_local_devices++; in CreateSession()