Home
last modified time | relevance | path

Searched refs:RendezvousKey (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcollective_permute_thunk.cc68 struct RendezvousKey { struct
78 friend H AbslHashValue(H h, const RendezvousKey& k) { in AbslHashValue()
81 friend bool operator==(const RendezvousKey& a, const RendezvousKey& b) { in operator ==()
84 friend bool operator!=(const RendezvousKey& a, const RendezvousKey& b) { in operator !=()
111 explicit Rendezvous(const RendezvousKey& key) : key_(key) {} in Rendezvous()
123 const RendezvousKey key_;
214 RefcountingHashMap<RendezvousKey, Rendezvous>& GlobalRendezvousMap() { in GlobalRendezvousMap()
215 static auto& m = *new RefcountingHashMap<RendezvousKey, Rendezvous>(); in GlobalRendezvousMap()
236 RendezvousKey key{params.run_id, params.device_assn->replica_count()}; in ExecuteOnStream()
237 auto rendezvous_factory = [](const RendezvousKey& key) { in ExecuteOnStream()
Dnccl_collective_thunk.h72 RendezvousKey::CollectiveOpKind collective_op_kind;
92 config.collective_op_kind = RendezvousKey::kCrossModule; in GetNcclCollectiveConfigForMlir()
95 config.collective_op_kind = RendezvousKey::kCrossReplica; in GetNcclCollectiveConfigForMlir()
Dnccl_utils.cc218 NcclCliqueRendezvous(const RendezvousKey& rendezvous_key, in NcclCliqueRendezvous()
330 const RendezvousKey& rendezvous_key, int local_device_ordinal, in AcquireNcclClique()
338 *new RefcountingHashMap<RendezvousKey, NcclCliqueRendezvous>(); in AcquireNcclClique()
346 rendezvous_key, [&](const RendezvousKey& rendezvous_key) { in AcquireNcclClique()
Dnccl_collective_thunk.cc89 RendezvousKey rendezvous_key(params.run_id, std::move(participants), in ExecuteOnStream()
Dnccl_utils.h151 const RendezvousKey& rendezvous_key, int local_device_ordinal,
/external/tensorflow/tensorflow/compiler/xla/service/
Dcollective_ops_utils.h75 struct RendezvousKey { struct
81 explicit RendezvousKey(const RunId& run_id, in RendezvousKey() argument
92 friend H AbslHashValue(H h, const RendezvousKey& k) { in AbslHashValue() argument
97 friend bool operator==(const RendezvousKey& a, const RendezvousKey& b) {
103 friend bool operator!=(const RendezvousKey& a, const RendezvousKey& b) {
142 ParticipantData(const RendezvousKey& rendezvous_key, int64 device_ordinal, in ParticipantData() argument
150 RendezvousKey rendezvous_key;
159 AllReduceParticipantData(const RendezvousKey& rendezvous_key_p, in AllReduceParticipantData()
212 explicit Rendezvous(const RendezvousKey& k) : key_(k) {} in Rendezvous()
301 const RendezvousKey key_;
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_runtime.cc138 CollectivePermuteParticipantData(const xla::RendezvousKey& rendezvous_key_p, in CollectivePermuteParticipantData()
160 AllToAllParticipantData(const xla::RendezvousKey& rendezvous_key_p, in AllToAllParticipantData()
322 explicit CpuAllToAllRendezvous(const xla::RendezvousKey& k) in CpuAllToAllRendezvous()
383 explicit CpuCollectivePermuteRendezvous(const xla::RendezvousKey& k) in CpuCollectivePermuteRendezvous()
425 explicit CpuAllReduceRendezvous(const xla::RendezvousKey& k) in CpuAllReduceRendezvous()
566 xla::RefcountingHashMap<xla::RendezvousKey, CpuAllReduceRendezvous>&
569 *new xla::RefcountingHashMap<xla::RendezvousKey, CpuAllReduceRendezvous>; in GlobalAllReduceRendezvousMap()
573 xla::RefcountingHashMap<xla::RendezvousKey, CpuCollectivePermuteRendezvous>&
575 static auto& m = *new xla::RefcountingHashMap<xla::RendezvousKey, in GlobalCollectivePermuteRendezvousMap()
580 xla::RefcountingHashMap<xla::RendezvousKey, CpuAllToAllRendezvous>&
[all …]