| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | collective_ops_utils_test.cc | 59 std::vector<GlobalDeviceId> actual = in TEST()
60 GetParticipatingDevices(GlobalDeviceId(42), device_assignment, in TEST()
63 std::vector<GlobalDeviceId> expected = { in TEST()
64 GlobalDeviceId(42), GlobalDeviceId(43), GlobalDeviceId(44)}; in TEST()
82 std::vector<GlobalDeviceId> actual = in TEST()
83 GetParticipatingDevices(GlobalDeviceId(42), device_assignment, in TEST()
86 std::vector<GlobalDeviceId> expected = {GlobalDeviceId(42), in TEST()
87 GlobalDeviceId(45)}; in TEST()
99 std::vector<GlobalDeviceId> actual = in TEST()
100 GetParticipatingDevices(GlobalDeviceId(44), device_assignment, in TEST()
[all …]
|
| D | global_device_id.h | 31 TF_LIB_GTL_DEFINE_INT_TYPE(GlobalDeviceId, int64);
34 std::string GlobalDeviceIdsToString(absl::Span<GlobalDeviceId const> ids);
|
| D | global_device_id.cc | 22 std::string GlobalDeviceIdsToString(absl::Span<GlobalDeviceId const> ids) { in GlobalDeviceIdsToString()
25 for (GlobalDeviceId id : ids) { in GlobalDeviceIdsToString()
|
| D | collective_ops_utils.h | 48 StatusOr<std::vector<GlobalDeviceId>> GetParticipatingDevices(
49 GlobalDeviceId device_id, const DeviceAssignment& device_assignment,
82 std::vector<GlobalDeviceId> global_devices, in RendezvousKey()
116 std::vector<GlobalDeviceId> global_devices;
180 std::vector<std::pair<GlobalDeviceId, int64>> local_devices;
|
| D | collective_ops_utils.cc | 78 StatusOr<std::vector<GlobalDeviceId>> GetParticipatingDevices( in GetParticipatingDevices()
79 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices()
81 std::vector<GlobalDeviceId> participants; in GetParticipatingDevices()
|
| D | computation_placer.h | 55 GlobalDeviceId device_id) const;
57 StatusOr<int> ReplicaIdForDevice(GlobalDeviceId device_id) const;
|
| D | computation_placer.cc | 46 GlobalDeviceId device_id) const { in LogicalIdsForDevice()
69 GlobalDeviceId device_id) const { in ReplicaIdForDevice()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_executable_run_options.h | 35 explicit NcclCliqueKey(std::vector<GlobalDeviceId> devices);
45 const std::vector<GlobalDeviceId>& devices() const { return devices_; } in devices()
50 std::vector<GlobalDeviceId> devices_;
66 absl::optional<std::vector<GlobalDeviceId>> gpu_global_device_ids);
67 const absl::optional<std::vector<GlobalDeviceId>>& gpu_global_device_ids()
77 absl::optional<std::vector<GlobalDeviceId>> gpu_global_device_ids_;
|
| D | gpu_executable_run_options.cc | 23 NcclCliqueKey::NcclCliqueKey(std::vector<GlobalDeviceId> devices) in NcclCliqueKey()
31 absl::optional<std::vector<GlobalDeviceId>> gpu_global_device_ids) { in set_gpu_global_device_ids()
36 const absl::optional<std::vector<GlobalDeviceId>>&
|
| D | nccl_test_utils.cc | 25 absl::flat_hash_set<GlobalDeviceId> DevicesWithOpenNcclChannels() { in DevicesWithOpenNcclChannels()
26 absl::flat_hash_set<GlobalDeviceId> devices; in DevicesWithOpenNcclChannels()
|
| D | thunk.cc | 21 StatusOr<GlobalDeviceId> Thunk::ExecuteParams::GetGlobalDeviceId() const { in GetGlobalDeviceId()
29 return GlobalDeviceId(local_device_ordinal); in GetGlobalDeviceId()
|
| D | nccl_collective_thunk.cc | 70 TF_ASSIGN_OR_RETURN(GlobalDeviceId global_device_id, in ExecuteOnStream()
74 std::vector<GlobalDeviceId> participants, in ExecuteOnStream()
|
| D | thunk.h | 112 const std::vector<GlobalDeviceId>* gpu_global_device_ids; // may be null
115 StatusOr<GlobalDeviceId> GetGlobalDeviceId() const;
|
| D | nccl_utils.h | 109 const std::vector<GlobalDeviceId>& participants,
110 const std::vector<GlobalDeviceId>* local_devices); // may be null
|
| D | nccl_test_utils_dummy.cc | 21 absl::flat_hash_set<GlobalDeviceId> DevicesWithOpenNcclChannels() { return {}; } in DevicesWithOpenNcclChannels()
|
| D | nccl_test_utils.h | 31 absl::flat_hash_set<GlobalDeviceId> DevicesWithOpenNcclChannels();
|
| D | replica_id_thunk.cc | 27 TF_ASSIGN_OR_RETURN(GlobalDeviceId global_device_id, in ExecuteOnStream()
|
| D | nccl_utils.cc | 259 const std::vector<GlobalDeviceId>& participants, in GetLocalParticipants()
260 const std::vector<GlobalDeviceId>* local_devices) { in GetLocalParticipants()
263 absl::flat_hash_map<GlobalDeviceId, int> device_ranks(participants.size()); in GetLocalParticipants()
|
| D | collective_permute_thunk.cc | 243 TF_ASSIGN_OR_RETURN(GlobalDeviceId global_device_id, in ExecuteOnStream()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | gpu_device.cc | 179 absl::flat_hash_map<GlobalDeviceId, int> device_to_node) in NcclIdStore() argument
189 const absl::flat_hash_map<GlobalDeviceId, int> device_to_node_;
268 std::vector<GlobalDeviceId> gpu_device_ids(local_device_states.size()); in BuildDistributedDevices()
269 absl::flat_hash_map<GlobalDeviceId, int> device_to_node; in BuildDistributedDevices()
272 GlobalDeviceId global_device_id(device_proto.global_device_id()); in BuildDistributedDevices()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_runtime.cc | 605 std::vector<xla::GlobalDeviceId> participating_devices = in GetRendezvousKey()
606 xla::GetParticipatingDevices(xla::GlobalDeviceId(device_ordinal), in GetRendezvousKey()
626 ->ReplicaIdForDevice(xla::GlobalDeviceId(device_ordinal)) in __xla_cpu_runtime_AllToAll()
714 ->ReplicaIdForDevice(xla::GlobalDeviceId(device_ordinal)) in __xla_cpu_runtime_ReplicaId()
730 ->ReplicaIdForDevice(xla::GlobalDeviceId(device_ordinal)) in __xla_cpu_runtime_CollectivePermute()
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_execute_op.cc | 326 std::vector<xla::GlobalDeviceId> gpu_global_ids; in RunExecutable()
330 gpu_global_ids.emplace_back(xla::GlobalDeviceId(gid)); in RunExecutable()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | collective_ops_test.cc | 161 absl::flat_hash_set<GlobalDeviceId> OpenNcclChannels() { in OpenNcclChannels()
|