| /external/tensorflow/tensorflow/core/kernels/ |
| D | nccl_ops.cc | 108 auto* compute_stream = c->op_device_context()->stream(); in ComputeAsync() local
111 compute_stream->parent(), compute_stream, gpu_info, input, output, in ComputeAsync()
139 auto* compute_stream = c->op_device_context()->stream(); in ComputeAsync() local
142 compute_stream->parent(), compute_stream, gpu_info, &c->input(0), in ComputeAsync()
175 auto* compute_stream = c->op_device_context()->stream(); in ComputeAsync() local
178 compute_stream->parent(), compute_stream, gpu_info, input, output, in ComputeAsync()
209 auto* compute_stream = c->op_device_context()->stream(); in ComputeAsync() local
212 compute_stream->parent(), compute_stream, gpu_info, &c->input(0), in ComputeAsync()
245 auto* compute_stream = c->op_device_context()->stream(); in ComputeAsync() local
248 compute_stream->parent(), compute_stream, gpu_info, in ComputeAsync()
|
| /external/tensorflow/tensorflow/core/nccl/ |
| D | collective_communicator.cc | 78 auto* compute_stream = col_ctx->op_ctx->op_device_context()->stream(); in Enqueue() local
81 compute_stream->parent(), compute_stream, gpu_info, col_ctx->input, in Enqueue()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | pjrt_stream_executor_client.cc | 282 if (stream != local_device->compute_stream()) { in StallStreamOnError()
283 local_device->compute_stream()->ThenWaitFor(stream); in StallStreamOnError()
384 copy_stream->ThenWaitFor(local_device->compute_stream()); in AllocateDestinationBuffer()
388 local_device->compute_stream()->parent(), dst_buffer)); in AllocateDestinationBuffer()
404 local_device->compute_stream())); in AllocateDestinationBuffer()
406 std::move(event), local_device->compute_stream()); in AllocateDestinationBuffer()
428 tuple_table_stream->ThenWaitFor(local_device->compute_stream()); in AllocateDestinationBuffer()
431 local_device->compute_stream()->parent(), dst_buffer)); in AllocateDestinationBuffer()
1095 local_device_state->compute_stream()) && in Release()
1323 transfer_stream, dst_local_device->compute_stream(), input_buffer, in CopyToDeviceHelper()
[all …]
|
| D | local_device_state.h | 106 se::Stream* compute_stream() const { return compute_stream_.get(); } in compute_stream() function
|
| D | gpu_device.cc | 138 local_device->compute_stream()); in CreateBFCAllocator()
|
| /external/tensorflow/tensorflow/core/common_runtime/device/ |
| D | device_event_mgr_test.cc | 159 se::Stream* compute_stream() { return gpu_->stream_->compute; } in compute_stream() function in tensorflow::GPUDeviceTestHelper
383 gpu_helper_->h2d_stream()->ThenWaitFor(gpu_helper_->compute_stream()); in DoAddChain()
394 gpu_helper_->compute_stream()->ThenWaitFor(gpu_helper_->h2d_stream()); in DoAddChain()
397 gpu_helper_->compute_stream(), [times, r]() { in DoAddChain()
407 gpu_helper_->event_mgr()->ThenExecute(gpu_helper_->compute_stream(), in DoAddChain()
414 gpu_helper_->compute_stream(), [times, r]() { in DoAddChain()
418 gpu_helper_->d2h_stream()->ThenWaitFor(gpu_helper_->compute_stream()); in DoAddChain()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_device_context.cc | 77 std::shared_ptr<se::Stream> compute_stream, in XlaDeviceContext() argument
84 : stream_(std::move(compute_stream)), in XlaDeviceContext()
|
| D | xla_device_context.h | 55 std::shared_ptr<se::Stream> compute_stream,
|
| /external/tensorflow/tensorflow/core/common_runtime/gpu/ |
| D | gpu_device.h | 202 se::Stream* compute_stream, in GPUKernelTracker() argument
207 stream_(compute_stream), in GPUKernelTracker()
|