Home
last modified time | relevance | path

Searched refs:col_ctx_ (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/core/common_runtime/
Dring_reducer.cc54 CHECK(col_ctx_); in Run()
77 VLOG(1) << "RingReducer::Run for device " << col_ctx_->device_name in Run()
84 if ((col_ctx_->input != col_ctx_->output) && in Run()
85 (DMAHelper::base(col_ctx_->input) != DMAHelper::base(col_ctx_->output))) { in Run()
91 col_ctx_->op_ctx->input_device_context(0), in Run()
92 col_ctx_->op_ctx->op_device_context(), col_ctx_->device, in Run()
93 col_ctx_->device, col_ctx_->op_ctx->input_alloc_attr(0), in Run()
94 col_ctx_->op_ctx->output_alloc_attr(0), col_ctx_->input, in Run()
95 col_ctx_->output, 0 /*dev_to_dev_stream_index*/, in Run()
112 AllocatorAttributes attr = col_ctx_->op_ctx->output_alloc_attr(0); in ContinueAfterInputCopy()
[all …]
Dhierarchical_tree_broadcaster.cc56 : col_ctx_(nullptr), in HierarchicalTreeBroadcaster()
190 col_ctx_ = col_ctx; in InitializeCollectiveContext()
198 CHECK(col_ctx_); in Run()
307 VLOG(1) << "Running Broadcast tree device=" << col_ctx_->device_name in RunTree()
320 DispatchRecv(si, recv_from_rank, my_rank, col_ctx_->output, in RunTree()
340 (is_source_ ? col_ctx_->input : col_ctx_->output), in RunTree()
360 VLOG(2) << "copying input to output for device=" << col_ctx_->device_name in RunTree()
362 if (col_ctx_->input != col_ctx_->output && in RunTree()
363 (DMAHelper::base(col_ctx_->input) != in RunTree()
364 DMAHelper::base(col_ctx_->output))) { in RunTree()
[all …]
Dring_gatherer.cc62 DCHECK(col_ctx_); in Run()
85 VLOG(1) << "RingGatherer::Run for device " << col_ctx_->device_name in Run()
91 AllocatorAttributes attr = col_ctx_->op_ctx->output_alloc_attr(0); in Run()
92 ca_.reset(MakeCollectiveAdapter(col_ctx_->output, group_size_ * num_subdivs_, in Run()
93 col_ctx_->device->GetAllocator(attr), in Run()
103 col_ctx_->op_ctx->input_device_context(0), in Run()
104 col_ctx_->op_ctx->op_device_context(), col_ctx_->device, col_ctx_->device, in Run()
105 col_ctx_->op_ctx->input_alloc_attr(0), in Run()
106 col_ctx_->op_ctx->output_alloc_attr(0), col_ctx_->input, &alias_chunk, in Run()
136 col_ctx_->device->tensorflow_gpu_device_info(); in RunAsyncParts()
[all …]
Dring_alg.cc101 col_ctx_(nullptr), in RingAlg()
245 col_ctx_ = col_ctx; in InitializeCollectiveContext()
254 col_ctx_->op_ctx->device()->tensorflow_gpu_device_info(); in TensorDebugString()
259 &tensor, "" /*tensor_name*/, col_ctx_->device, &cpu_tensor, in TensorDebugString()
289 col_ctx_->col_exec->StartAbort(s); in StartAbort()
296 ca_->ConsumeFinalValue(col_ctx_->output); in Finish()
382 string send_buf_key = RingAlgBufKey(name_, col_ctx_->exec_key, in DispatchSend()
390 col_ctx_->col_exec->PostToPeer( in DispatchSend()
393 col_ctx_->device, col_ctx_->op_ctx->op_device_context(), in DispatchSend()
394 col_ctx_->op_ctx->output_alloc_attr(0), &rf->chunk, in DispatchSend()
[all …]
Dhierarchical_tree_broadcaster.h82 CollectiveContext* col_ctx_; // Not owned variable
Dring_alg.h110 CollectiveContext* col_ctx_; // Not owned variable
/external/tensorflow/tensorflow/core/kernels/
Dcollective_nccl_reducer.cc29 NcclReducer::NcclReducer() : col_ctx_(nullptr), col_params_(nullptr) {} in NcclReducer()
45 col_ctx_ = col_ctx; in InitializeCollectiveContext()
89 Tensor group_size_val(col_ctx_->output->dtype(), TensorShape({})); in Run()
90 switch (col_ctx_->output->dtype()) { in Run()
104 done(errors::Internal("Unsupported type ", col_ctx_->output->dtype())); in Run()
108 col_ctx_->device->GetAllocator(col_ctx_->op_ctx->input_alloc_attr(0)), in Run()
109 col_ctx_->output->dtype(), TensorShape({})); in Run()
110 DeviceContext* op_dev_ctx = col_ctx_->op_ctx->op_device_context(); in Run()
113 &group_size_val, col_ctx_->device, &group_size, in Run()
124 auto* compute_stream = col_ctx_->op_ctx->op_device_context()->stream(); in Run()
[all …]
Dcollective_nccl_reducer.h42 CollectiveContext* col_ctx_; // Not owned