Note that this function is blocking and must not run in any thread which cannot be blocked.
| 113 | // Note that this function is blocking and must not run in any thread |
| 114 | // which cannot be blocked. |
| 115 | void RingReducer::ContinueAfterInputCopy() { |
| 116 | AllocatorAttributes attr = col_ctx_->op_ctx->output_alloc_attr(0); |
| 117 | ca_.reset(MakeCollectiveAdapter(col_ctx_->output, group_size_ * num_subdivs_, |
| 118 | col_ctx_->device->GetAllocator(attr))); |
| 119 | |
| 120 | if (col_params_->final_op) { |
| 121 | // Create an on-device scalar value from group_size_ that may be needed |
| 122 | // later. |
| 123 | // TODO(tucker): Cache and reuse across invocations? Or maybe the scalar |
| 124 | // can be provided to the kernel in host memory? |
| 125 | Tensor group_size_val = ca_->Scalar(group_size_); |
| 126 | if (col_params_->group.device_type != "CPU") { |
| 127 | uint64 safe_alloc_frontier = col_ctx_->device->SafeAllocFrontier(0); |
| 128 | AllocationAttributes aa; |
| 129 | std::function<uint64()> freed_by_func = [this, &safe_alloc_frontier]() { |
| 130 | safe_alloc_frontier = |
| 131 | col_ctx_->device->SafeAllocFrontier(safe_alloc_frontier); |
| 132 | return safe_alloc_frontier; |
| 133 | }; |
| 134 | if (safe_alloc_frontier > 0) { |
| 135 | aa.freed_by_func = &freed_by_func; |
| 136 | } |
| 137 | group_size_tensor_ = ca_->Scalar( |
| 138 | col_ctx_->device->GetAllocator(col_ctx_->op_ctx->input_alloc_attr(0)), |
| 139 | aa); |
| 140 | DeviceContext* op_dev_ctx = col_ctx_->op_ctx->op_device_context(); |
| 141 | op_dev_ctx->CopyCPUTensorToDevice( |
| 142 | &group_size_val, col_ctx_->device, &group_size_tensor_, |
| 143 | [this](const Status& s) { |
| 144 | if (!s.ok()) { |
| 145 | StartAbort(s); |
| 146 | } |
| 147 | group_size_tensor_ready_.Notify(); |
| 148 | }, |
| 149 | (safe_alloc_frontier == 0)); |
| 150 | } else { |
| 151 | group_size_tensor_ = group_size_val; |
| 152 | group_size_tensor_ready_.Notify(); |
| 153 | } |
| 154 | } else { |
| 155 | // Value won't be used, so no need to initialize. |
| 156 | group_size_tensor_ready_.Notify(); |
| 157 | } |
| 158 | Finish(RunAsyncParts()); |
| 159 | } |
| 160 | |
| 161 | void RingReducer::InitRingField(RingField* rf, int chunk_idx, int subdiv_idx, |
| 162 | int field_idx) { |
nothing calls this directly
no test coverage detected