| 62 | } |
| 63 | |
| 64 | void RingGatherer::Run(StatusCallback done) { |
| 65 | DCHECK(col_ctx_); |
| 66 | DCHECK(col_params_); |
| 67 | done_ = std::move(done); |
| 68 | group_size_ = col_params_->group.group_size; |
| 69 | num_subdivs_ = static_cast<int>( |
| 70 | col_params_->instance.impl_details.subdiv_permutations.size()); |
| 71 | DCHECK_GT(num_subdivs_, 0); |
| 72 | |
| 73 | if (VLOG_IS_ON(1)) { |
| 74 | string buf; |
| 75 | for (int r = 0; r < col_params_->instance.device_names.size(); ++r) { |
| 76 | strings::StrAppend(&buf, "dev ", r, " : ", |
| 77 | col_params_->instance.device_names[r], "\n"); |
| 78 | } |
| 79 | for (int sd = 0; |
| 80 | sd < col_params_->instance.impl_details.subdiv_permutations.size(); |
| 81 | ++sd) { |
| 82 | strings::StrAppend(&buf, "\nsubdiv ", sd, " perm: "); |
| 83 | for (auto x : |
| 84 | col_params_->instance.impl_details.subdiv_permutations[sd]) { |
| 85 | strings::StrAppend(&buf, x, ", "); |
| 86 | } |
| 87 | } |
| 88 | VLOG(1) << "RingGatherer::Run for device " << col_ctx_->device_name |
| 89 | << " default_rank " << col_params_->default_rank << "\n" |
| 90 | << buf; |
| 91 | } |
| 92 | |
| 93 | // Prepare to alias fields within the output. |
| 94 | AllocatorAttributes attr = col_ctx_->op_ctx->output_alloc_attr(0); |
| 95 | ca_.reset(MakeCollectiveAdapter(col_ctx_->output, group_size_ * num_subdivs_, |
| 96 | col_ctx_->device->GetAllocator(attr), |
| 97 | false /*align_chunks*/)); |
| 98 | |
| 99 | // Start by copying input to the rank-specific offset of output. |
| 100 | // We are running in a blockable thread and the callback can't block so |
| 101 | // just wait here on the copy. |
| 102 | { |
| 103 | profiler::TraceMe activity("MemCpyAsync", profiler::TraceMeLevel::kInfo); |
| 104 | Notification note; |
| 105 | Status status; |
| 106 | Tensor alias_chunk(ca_->ChunkAlias(col_params_->subdiv_rank[0])); |
| 107 | CollectiveRemoteAccessLocal::MemCpyAsync( |
| 108 | col_ctx_->op_ctx->op_device_context(), |
| 109 | col_ctx_->op_ctx->op_device_context(), col_ctx_->device, |
| 110 | col_ctx_->device, col_ctx_->op_ctx->input_alloc_attr(0), |
| 111 | col_ctx_->op_ctx->output_alloc_attr(0), col_ctx_->input, &alias_chunk, |
| 112 | 0 /*dev_to_dev_stream_index*/, [¬e, &status](const Status& s) { |
| 113 | status.Update(s); |
| 114 | note.Notify(); |
| 115 | }); |
| 116 | note.WaitForNotification(); |
| 117 | if (!status.ok()) { |
| 118 | done_(status); |
| 119 | return; |
| 120 | } |
| 121 | } |
nothing calls this directly
no test coverage detected