Home
last modified time | relevance | path

Searched refs:col_params_ (Results 1 – 17 of 17) sorted by relevance

/external/tensorflow/tensorflow/core/common_runtime/
Dpermuter.cc40 : col_ctx_(nullptr), col_params_(nullptr), done_(nullptr), counter_(0) {} in Permuter()
57 col_params_ = col_ctx->col_params; in InitializeCollectiveContext()
64 if (col_params_->instance.permutation.size() != in Run()
65 col_params_->instance.devices.size()) { in Run()
69 DispatchSend(col_params_->default_rank, in Run()
70 col_params_->instance.permutation[col_params_->default_rank], in Run()
72 for (int i = 0; i < col_params_->instance.permutation.size(); ++i) { in Run()
73 if (col_params_->default_rank == col_params_->instance.permutation[i]) { in Run()
74 DispatchRecv(i, col_params_->instance.permutation[i], col_ctx_->output, in Run()
86 << col_params_->instance.devices[target_rank] in DispatchSend()
[all …]
Dpermuter_test.cc111 : device_type_(DEVICE_CPU), col_exec_(nullptr), col_params_(nullptr) {} in PermuterTest()
117 if (col_params_) col_params_->Unref(); in ~PermuterTest()
175 col_params_ = new CollectiveParams(); in Init()
176 col_params_->name = "test_collective"; in Init()
177 col_params_->instance.data_type = dtype; in Init()
179 col_params_->instance.instance_key = kInstanceKey; in Init()
180 col_params_->group.device_type = device_type; in Init()
181 col_params_->instance.type = PERMUTE_COLLECTIVE; in Init()
193 col_params_->group.device_names.push_back(dev_name); in Init()
194 col_params_->instance.devices.push_back(dev_name); in Init()
[all …]
Dhierarchical_tree_broadcaster_test.cc201 : device_type_(DEVICE_CPU), col_exec_(nullptr), col_params_(nullptr) {} in HierarchicalTreeBroadcasterTest()
207 if (col_params_) col_params_->Unref(); in ~HierarchicalTreeBroadcasterTest()
268 col_params_ = new CollectiveParams(); in Init()
269 col_params_->name = "test_collective"; in Init()
270 col_params_->instance.data_type = dtype; in Init()
272 col_params_->group.group_key = kGroupKey; in Init()
274 col_params_->instance.instance_key = kInstanceKey; in Init()
275 col_params_->group.device_type = device_type; in Init()
276 col_params_->group.group_size = num_workers * num_devices_per_worker; in Init()
277 col_params_->instance.impl_details.subdiv_offsets.clear(); in Init()
[all …]
Dring_reducer_test.cc142 : device_type_(DEVICE_CPU), col_exec_(nullptr), col_params_(nullptr) {} in RingReducerTest()
159 if (col_params_) col_params_->Unref(); in ~RingReducerTest()
208 col_params_ = new CollectiveParams(); in Init()
209 col_params_->name = "test_collective"; in Init()
211 col_params_->group.group_key = kGroupKey; in Init()
212 col_params_->group.device_type = device_type; in Init()
213 col_params_->group.group_size = num_workers * num_devices; in Init()
215 col_params_->instance.instance_key = kInstanceKey; in Init()
216 col_params_->instance.impl_details.subdiv_offsets.clear(); in Init()
217 col_params_->instance.type = REDUCTION_COLLECTIVE; in Init()
[all …]
Dring_gatherer_test.cc119 : device_type_(DEVICE_CPU), col_exec_(nullptr), col_params_(nullptr) {} in RingGathererTest()
136 if (col_params_) col_params_->Unref(); in ~RingGathererTest()
185 col_params_ = new CollectiveParams(); in Init()
186 col_params_->name = "test_collective"; in Init()
188 col_params_->group.group_key = kGroupKey; in Init()
189 col_params_->group.device_type = device_type; in Init()
190 col_params_->group.group_size = num_workers * num_devices; in Init()
192 col_params_->instance.instance_key = kInstanceKey; in Init()
193 col_params_->instance.impl_details.subdiv_offsets.clear(); in Init()
194 col_params_->instance.type = GATHER_COLLECTIVE; in Init()
[all …]
Dhierarchical_tree_broadcaster.cc58 col_params_(nullptr), in HierarchicalTreeBroadcaster()
191 col_params_ = col_ctx->col_params; in InitializeCollectiveContext()
199 CHECK(col_params_); in Run()
201 is_source_ = col_params_->is_source; in Run()
291 int num_subdivs = static_cast<int>(col_params_->subdiv_rank.size()); in RunTree()
298 int my_rank = col_params_->subdiv_rank[si]; in RunTree()
301 int source_rank = col_params_->instance.impl_details.subdiv_source_rank[si]; in RunTree()
304 for (int r : col_params_->instance.impl_details.subdiv_permutations[si]) { in RunTree()
321 int recv_from_rank = TreeRecvFrom(*col_params_, si); in RunTree()
339 TreeSendTo(*col_params_, si, &send_to_ranks); in RunTree()
[all …]
Dring_reducer.cc57 CHECK(col_params_); in Run()
60 col_ctx_->col_exec->UnblockDependencies(*col_params_); in Run()
63 group_size_ = col_params_->group.group_size; in Run()
65 col_params_->instance.impl_details.subdiv_permutations.size()); in Run()
70 for (int r = 0; r < col_params_->group.device_names.size(); ++r) { in Run()
72 col_params_->group.device_names[r], "\n"); in Run()
75 sd < col_params_->instance.impl_details.subdiv_permutations.size(); in Run()
79 col_params_->instance.impl_details.subdiv_permutations[sd]) { in Run()
84 << " default_rank " << col_params_->default_rank << "\n" in Run()
123 if (col_params_->final_op) { in ContinueAfterInputCopy()
[all …]
Dring_alg.cc103 col_params_(nullptr), in RingAlg()
248 col_params_ = col_ctx->col_params; in InitializeCollectiveContext()
324 rf->rank = col_params_->subdiv_rank[subdiv_idx]; in InitRingField()
330 rf->recv_dev_idx = col_params_->instance.impl_details in InitRingField()
332 int send_dev_idx = col_params_->instance.impl_details in InitRingField()
334 rf->recv_is_remote = !col_params_->task.is_local[rf->recv_dev_idx]; in InitRingField()
335 rf->send_is_remote = !col_params_->task.is_local[send_dev_idx]; in InitRingField()
386 VLOG(3) << "DispatchSend rank=" << col_params_->default_rank << " send key " in DispatchSend()
390 int send_to_dev_idx = col_params_->instance.impl_details in DispatchSend()
393 col_params_->group.device_names[send_to_dev_idx], in DispatchSend()
[all …]
Dring_gatherer.cc65 DCHECK(col_params_); in Run()
67 group_size_ = col_params_->group.group_size; in Run()
69 col_params_->instance.impl_details.subdiv_permutations.size()); in Run()
74 for (int r = 0; r < col_params_->group.device_names.size(); ++r) { in Run()
76 col_params_->group.device_names[r], "\n"); in Run()
79 sd < col_params_->instance.impl_details.subdiv_permutations.size(); in Run()
83 col_params_->instance.impl_details.subdiv_permutations[sd]) { in Run()
88 << " default_rank " << col_params_->default_rank << "\n" in Run()
105 Tensor alias_chunk(ca_->ChunkAlias(col_params_->subdiv_rank[0])); in Run()
Dpermuter.h68 const CollectiveParams* col_params_; // Not owned variable
Dhierarchical_tree_broadcaster.h85 const CollectiveParams* col_params_; // Not owned variable
Dring_alg.h113 const CollectiveParams* col_params_; // Not owned variable
/external/tensorflow/tensorflow/core/kernels/
Dcollective_ops.cc56 : AsyncOpKernel(c), name_(name()), col_params_(new CollectiveParams()) {} in CollectiveOpV1Kernel()
58 ~CollectiveOpV1Kernel() override { col_params_->Unref(); } in ~CollectiveOpV1Kernel()
93 return CollectiveKey(c, col_params_->group.group_key, in GetCollectiveKey()
94 col_params_->instance.instance_key); in GetCollectiveKey()
101 if (col_params_->group.group_size > in CanProceedWithCompute()
102 col_params_->group.device_names.size()) { in CanProceedWithCompute()
108 << col_params_->name << " device " << c->device()->name() in CanProceedWithCompute()
109 << " group " << col_params_->group.group_key << " instance " in CanProceedWithCompute()
110 << col_params_->instance.instance_key; in CanProceedWithCompute()
112 c->device()->attributes(), col_params_, c->cancellation_manager(), in CanProceedWithCompute()
[all …]
Dcollective_nccl_test.cc92 col_params_(nullptr) {} in NcclTestBase()
96 if (col_params_) col_params_->Unref(); in ~NcclTestBase()
131 col_params_ = new CollectiveParams(); in Init()
132 col_params_->name = "test_nccl_collective_op"; in Init()
134 col_params_->group.group_key = group_key; in Init()
135 col_params_->group.device_type = DEVICE_GPU; in Init()
136 col_params_->group.group_size = num_ranks; in Init()
137 col_params_->instance.instance_key = instance_key; in Init()
138 col_params_->instance.type = collective_type_; in Init()
139 col_params_->instance.data_type = DT_FLOAT; in Init()
[all …]
Dcollective_nccl_reducer.cc31 if (col_params_->final_op) { in Run()
40 Tensor(static_cast<Eigen::half>(col_params_->group.group_size)); in Run()
44 Tensor(static_cast<float>(col_params_->group.group_size)); in Run()
48 Tensor(static_cast<double>(col_params_->group.group_size)); in Run()
52 Tensor(static_cast<int32>(col_params_->group.group_size)); in Run()
56 Tensor(static_cast<int64>(col_params_->group.group_size)); in Run()
82 if (col_params_->final_op) { in Run()
92 col_params_->Ref(); in Run()
93 core::ScopedUnref unref(col_params_); in Run()
97 if (!col_params_->final_op) { in Run()
[all …]
Dcollective_nccl.cc27 : type_(type), name_(name), col_ctx_(nullptr), col_params_(nullptr) {} in NcclBase()
64 col_params_ = col_ctx->col_params; in InitializeCollectiveContext()
Dcollective_nccl.h43 const CollectiveParams* col_params_; // Not owned variable