Searched refs:RendezvousKey (Results 1 – 7 of 7) sorted by relevance
/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | collective_permute_thunk.cc | 68 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()
|
D | nccl_collective_thunk.h | 72 RendezvousKey::CollectiveOpKind collective_op_kind; 92 config.collective_op_kind = RendezvousKey::kCrossModule; in GetNcclCollectiveConfigForMlir() 95 config.collective_op_kind = RendezvousKey::kCrossReplica; in GetNcclCollectiveConfigForMlir()
|
D | nccl_utils.cc | 218 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()
|
D | nccl_collective_thunk.cc | 89 RendezvousKey rendezvous_key(params.run_id, std::move(participants), in ExecuteOnStream()
|
D | nccl_utils.h | 151 const RendezvousKey& rendezvous_key, int local_device_ordinal,
|
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | collective_ops_utils.h | 75 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/ |
D | cpu_runtime.cc | 138 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 …]
|