| 39 | } // namespace |
| 40 | |
| 41 | void NcclReducer::Run(StatusCallback done) { |
| 42 | ncclRedOp_t reduction_op; |
| 43 | Status s = ReductionOp(col_params_->merge_op->type_string(), &reduction_op); |
| 44 | if (!s.ok()) { |
| 45 | done(s); |
| 46 | return; |
| 47 | } |
| 48 | |
| 49 | Tensor group_size; |
| 50 | Notification group_size_ready; |
| 51 | Status group_size_status; |
| 52 | if (col_params_->final_op) { |
| 53 | // Create an on-device scalar value from group_size_. |
| 54 | // TODO(ayushd, tucker): avoid this copy by either reusing across |
| 55 | // invocations or providing the scalar to the kernel in host memory. |
| 56 | Tensor group_size_val(col_ctx_->output->dtype(), TensorShape({})); |
| 57 | switch (col_ctx_->output->dtype()) { |
| 58 | case DT_FLOAT: |
| 59 | group_size_val.scalar<float>()() = col_params_->group.group_size; |
| 60 | break; |
| 61 | case DT_DOUBLE: |
| 62 | group_size_val.scalar<double>()() = col_params_->group.group_size; |
| 63 | break; |
| 64 | case DT_INT32: |
| 65 | group_size_val.scalar<int32>()() = col_params_->group.group_size; |
| 66 | break; |
| 67 | case DT_INT64: |
| 68 | group_size_val.scalar<int64>()() = col_params_->group.group_size; |
| 69 | break; |
| 70 | default: |
| 71 | done(errors::Internal("Unsupported type ", col_ctx_->output->dtype())); |
| 72 | return; |
| 73 | } |
| 74 | group_size = Tensor( |
| 75 | col_ctx_->device->GetAllocator(col_ctx_->op_ctx->input_alloc_attr(0)), |
| 76 | col_ctx_->output->dtype(), TensorShape({})); |
| 77 | DeviceContext* op_dev_ctx = col_ctx_->op_ctx->op_device_context(); |
| 78 | // Enqueue copy on gpu stream. |
| 79 | op_dev_ctx->CopyCPUTensorToDevice( |
| 80 | &group_size_val, col_ctx_->device, &group_size, |
| 81 | [&group_size_ready, &group_size_status](const Status& s) { |
| 82 | group_size_status = s; |
| 83 | group_size_ready.Notify(); |
| 84 | }); |
| 85 | } else { |
| 86 | group_size_ready.Notify(); |
| 87 | } |
| 88 | |
| 89 | Notification nccl_done; |
| 90 | Status nccl_status; |
| 91 | auto* compute_stream = col_ctx_->op_ctx->op_device_context()->stream(); |
| 92 | auto* gpu_info = col_ctx_->op_ctx->device()->tensorflow_gpu_device_info(); |
| 93 | // `AddToAllReduce` performs consistency checks for the NCCL call and enqueues |
| 94 | // the `Participant` struct locally. When all local participants with this |
| 95 | // `nccl_collective_key` have called `AddToAllReduce` and |
| 96 | // `SignalMultiNodeReady`, all devices at this worker are ready to process |
| 97 | // this NCCL op. |
| 98 | // |
nothing calls this directly
no test coverage detected