Searched refs:NcclAllReduceThunk (Results 1 – 3 of 3) sorted by relevance
31 /* static */ bool NcclAllReduceThunk::NcclIsEnabled() { in NcclIsEnabled()62 NcclAllReduceThunk* originator;314 Status NcclAllReduceThunk::ExecuteOnStream( in ExecuteOnStream()334 Status NcclAllReduceThunk::ExecuteOnStream( in ExecuteOnStream()344 NcclAllReduceThunk::NcclAllReduceThunk( in NcclAllReduceThunk() function in xla::gpu::NcclAllReduceThunk
31 class NcclAllReduceThunk : public Thunk {43 NcclAllReduceThunk(int64 replica_count, int64 element_count,
1392 << "; NCCL is enabled: " << NcclAllReduceThunk::NcclIsEnabled(); in HandleAllReduce()1407 AddThunkToThunkSequence(absl::make_unique<NcclAllReduceThunk>( in HandleAllReduce()1421 crs->IsCrossReplicaAllReduce(), NcclAllReduceThunk::NcclIsEnabled()); in HandleAllReduce()