| 54 | } |
| 55 | |
| 56 | void RingReducer::Run(StatusCallback done) { |
| 57 | CHECK(col_ctx_); |
| 58 | CHECK(col_params_); |
| 59 | done_ = std::move(done); |
| 60 | group_size_ = col_params_->group.group_size; |
| 61 | num_subdivs_ = static_cast<int>( |
| 62 | col_params_->instance.impl_details.subdiv_permutations.size()); |
| 63 | CHECK_GT(num_subdivs_, 0); |
| 64 | |
| 65 | if (VLOG_IS_ON(1)) { |
| 66 | string buf; |
| 67 | for (int r = 0; r < col_params_->instance.device_names.size(); ++r) { |
| 68 | strings::StrAppend(&buf, "dev ", r, " : ", |
| 69 | col_params_->instance.device_names[r], "\n"); |
| 70 | } |
| 71 | for (int sd = 0; |
| 72 | sd < col_params_->instance.impl_details.subdiv_permutations.size(); |
| 73 | ++sd) { |
| 74 | strings::StrAppend(&buf, "\nsubdiv ", sd, " perm: "); |
| 75 | for (auto x : |
| 76 | col_params_->instance.impl_details.subdiv_permutations[sd]) { |
| 77 | strings::StrAppend(&buf, x, ", "); |
| 78 | } |
| 79 | } |
| 80 | VLOG(1) << "RingReducer::Run for device " << col_ctx_->device_name |
| 81 | << " default_rank " << col_params_->default_rank << "\n" |
| 82 | << buf; |
| 83 | } |
| 84 | |
| 85 | // Start by copying input to output if they're not already the same, i.e. if |
| 86 | // we're not computing in-place on the input tensor. |
| 87 | if ((col_ctx_->input != col_ctx_->output) && |
| 88 | (DMAHelper::base(col_ctx_->input) != DMAHelper::base(col_ctx_->output))) { |
| 89 | // We are running in a blockable thread and the callback can't block so |
| 90 | // just wait here on the copy. |
| 91 | Notification note; |
| 92 | Status status; |
| 93 | profiler::TraceMe activity("MemCpyAsync", profiler::TraceMeLevel::kInfo); |
| 94 | CollectiveRemoteAccessLocal::MemCpyAsync( |
| 95 | col_ctx_->op_ctx->op_device_context(), |
| 96 | col_ctx_->op_ctx->op_device_context(), col_ctx_->device, |
| 97 | col_ctx_->device, col_ctx_->op_ctx->input_alloc_attr(0), |
| 98 | col_ctx_->op_ctx->output_alloc_attr(0), col_ctx_->input, |
| 99 | col_ctx_->output, 0 /*dev_to_dev_stream_index*/, |
| 100 | [¬e, &status](const Status& s) { |
| 101 | status.Update(s); |
| 102 | note.Notify(); |
| 103 | }); |
| 104 | note.WaitForNotification(); |
| 105 | if (!status.ok()) { |
| 106 | done_(status); |
| 107 | return; |
| 108 | } |
| 109 | } |
| 110 | ContinueAfterInputCopy(); |
| 111 | } |
| 112 | |
| 113 | // Note that this function is blocking and must not run in any thread |
nothing calls this directly
no test coverage detected