| /external/tensorflow/tensorflow/core/common_runtime/ |
| D | permuter.cc | 80 void Permuter::DispatchSend(int src_rank, int target_rank, const Tensor* tensor, in DispatchSend() argument
83 strings::StrCat(col_ctx_->exec_key, src_rank, target_rank); in DispatchSend()
87 << " target_rank=" << target_rank << " src_rank=" << src_rank; in DispatchSend()
96 void Permuter::DispatchRecv(int src_rank, int target_rank, Tensor* tensor, in DispatchRecv() argument
99 strings::StrCat(col_ctx_->exec_key, src_rank, target_rank); in DispatchRecv()
102 << col_params_->instance.devices[src_rank] in DispatchRecv()
103 << " target_rank=" << target_rank << " src_rank=" << src_rank; in DispatchRecv()
105 col_params_->instance.devices[src_rank], in DispatchRecv()
106 col_params_->group.task_names[src_rank], in DispatchRecv()
107 col_params_->task.is_local[src_rank], recv_buf_key, col_ctx_->device, in DispatchRecv()
|
| D | hierarchical_tree_broadcaster.cc | 44 string BroadcastBufKey(const string& exec_key, int subdiv, int src_rank, in BroadcastBufKey() argument
48 "):src(", src_rank, "):dst(", dst_rank, ")"); in BroadcastBufKey()
51 return strings::StrCat(exec_key, ":", subdiv, ":", src_rank, ":", dst_rank); in BroadcastBufKey()
407 int src_rank, in DispatchSend() argument
414 BroadcastBufKey(col_ctx_->exec_key, subdiv, src_rank, dst_rank); in DispatchSend()
430 void HierarchicalTreeBroadcaster::DispatchRecv(int subdiv, int src_rank, in DispatchRecv() argument
434 BroadcastBufKey(col_ctx_->exec_key, subdiv, src_rank, dst_rank); in DispatchRecv()
436 col_params_->instance.impl_details.subdiv_permutations[subdiv][src_rank]; in DispatchRecv()
440 << " src_rank=" << src_rank << " src_idx=" << src_idx; in DispatchRecv()
|
| D | permuter.h | 74 void DispatchSend(int src_rank, int target_rank, const Tensor* tensor,
77 void DispatchRecv(int src_rank, int target_rank, Tensor* tensor,
|
| D | hierarchical_tree_broadcaster.h | 72 void DispatchSend(int subdiv, int dst_rank, int src_rank,
78 void DispatchRecv(int subdiv, int src_rank, int dst_rank, Tensor* dst_tensor,
|
| /external/tensorflow/tensorflow/core/nccl/ |
| D | nccl_manager_test.cc | 211 TensorShape shape, int src_node, int src_rank, in MakeBroadcastTestCase() argument
221 if (node == src_node && local_rank == src_rank) { in MakeBroadcastTestCase()
624 const int src_rank = static_cast<int>(random::New64() % num_ranks); in TYPED_TEST() local
628 /*src_node=*/0, src_rank, in_place); in TYPED_TEST()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | legalize_to_linalg.cc | 1314 int src_rank = src_type.getRank(); in matchAndRewrite() local
1315 if (!src_rank) { in matchAndRewrite()
1335 rewriter.getContext(), src_rank, reduction_dims)); in matchAndRewrite()
1342 for (int i = 0, e = src_rank - reduction_dims.size(); i < e; ++i) in matchAndRewrite()
1344 indexing_maps.emplace_back(AffineMap::get(src_rank, /*symbolCount=*/0, in matchAndRewrite()
1361 GetParallelAndReductionIterators(src_rank, reduction_dims.size())); in matchAndRewrite()
|