Method ComputeAsync

tensorflow/core/kernels/nccl_ops.cc:97–121

      : NcclReduceOpBase(c) {}

  void ComputeAsync(OpKernelContext* c, DoneCallback done) override {
    const Tensor* input = &c->input(0);
    Tensor* output;
    OP_REQUIRES_OK_ASYNC(
        c, c->forward_input_or_allocate_output({0}, 0, input->shape(), &output),
        done);
    auto actual_done = [c, done](Status s) {
      OP_REQUIRES_OK_ASYNC(c, s, done);
      done();
    };

    auto* compute_stream = c->op_device_context()->stream();
    auto* gpu_info = c->device()->tensorflow_gpu_device_info();
    auto participant = absl::make_unique<NcclManager::Participant>(
        compute_stream->parent(), compute_stream, gpu_info,
        input, output, /*global_rank=*/-1,
        std::move(actual_done));
    NcclManager::instance()->AddToAllReduce(
        std::move(participant),
        {GetCollectiveKey(c),
         /*num_local_devices=*/num_devices(),
         /*num_global_devices=*/num_devices(),
         /*communicator_key=*/"", /*source_rank=*/-1},
        reduction_op());
  }
};
REGISTER_KERNEL_BUILDER(Name("NcclAllReduce").Device(DEVICE_GPU),
                        NcclAllReduceOpKernel);