Home
last modified time | relevance | path

Searched refs:async_comms_stream (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dnccl_all_reduce_thunk.cc286 se::Stream& async_comms_stream = *params.async_comms_stream; in RunNcclCollective() local
288 async_comms_stream.ThenWaitFor(params.stream); in RunNcclCollective()
295 async_comms_stream, comm)); in RunNcclCollective()
298 se::Event done_event(async_comms_stream.parent()); in RunNcclCollective()
300 async_comms_stream.ThenRecordEvent(&done_event); in RunNcclCollective()
302 int device_ordinal = async_comms_stream.parent()->device_ordinal(); in RunNcclCollective()
H A Dthunk.cc24 se::Stream* async_comms_stream) in ExecuteParams() argument
27 async_comms_stream(async_comms_stream), in ExecuteParams()
H A Dgpu_executable.cc326 StatusOr<StreamPool::Ptr> async_comms_stream = in ExecuteThunks() local
341 TF_RET_CHECK(async_comms_stream.ok() || !NeedsAsyncCommsStream(*thunk)) in ExecuteThunks()
346 async_comms_stream.ok() ? async_comms_stream->get() : nullptr}; in ExecuteThunks()
638 StatusOr<StreamPool::Ptr> async_comms_stream = in ExecuteJitRt() local
645 async_comms_stream.ok() ? async_comms_stream->get() : nullptr); in ExecuteJitRt()
H A Dthunk.h109 se::Stream* stream, se::Stream* async_comms_stream);
113 se::Stream* async_comms_stream; member