void ComputeAsync()

in src/nccl_op.cc [454:528]
52 lines of code
9 McCabe index (conditional complexity)

    void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override
    {
        if (op_event == NULL)
            cuEventCreate(&op_event, CU_EVENT_DISABLE_TIMING ); // | CU_EVENT_BLOCKING_SYNC

        const Tensor& input = ctx->input(0);

        char debug_message[1000];
        bool verbose = !debug_str.empty();
        if (verbose)
        {
            sprintf(debug_message, " op_num: %03d size: %s debug_str: %s\n", op_num, input.shape().DebugString().c_str(), debug_str.c_str());
            std::cout << "START";
            std::cout << debug_message;
            fflush(stdout);
        }

        // TODO: Turn on for all ranks
        if (!logfile.empty())
            if (FILE* log = fopen(logfile.c_str(), "a"))
            {
              fprintf(log, "%03d %10d %12s %d %d %d\n",
                op_num,
                (int)input.NumElements(),
                input.shape().DebugString().c_str(),
                (int)ctx->step_id(),
                (int)ctx->frame_iter().frame_id,
                (int)ctx->frame_iter().iter_id
              );
              fclose(log);
            }

        // Nccl ops work in place, so save some memory here.
        // ctx->set_output(0, input);
        Tensor* output;
        OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));

        auto actual_done = [ctx, done, debug_message, verbose](Status s) {
            if (verbose)
            {
                std::cout << "END  ";
                std::cout << debug_message;
                fflush(stdout);
            }
            OP_REQUIRES_OK_ASYNC(ctx, s, done);
            done();
        };

        auto stream    = ctx->op_device_context()->stream();
        auto executor  = stream->parent();
        //auto gpu_info  = ctx->device()->tensorflow_gpu_device_info();
        //auto event_mgr = gpu_info->event_mgr;
        //int  gpu_id    = gpu_info->gpu_id;

        CUstream cu_stream = reinterpret_cast<CUstream>(stream->implementation()->GpuStreamHack());

        // Record an event to mark this ops insertion in the compute stream.
        // The reduce op will be executed on the nccl stream but
        // not before this event is reached on the compute stream.
        cuEventRecord(op_event, cu_stream);

        // this keeps the backward pass from gobbling up buffers for every reduce op.
        // By syncing we let big ones finish first.
        // It's ok to greedily execute small buffers.
        if (sync_size && input.NumElements() > sync_size)
            cuStreamSynchronize(cu_stream);

        auto manager = NcclManager::instance(num_comms, prereduce, executor);

        if (!mpi_ranks.empty())
        {
            manager->CreateCustomComm(op_num, mpi_ranks, comm_id, executor);
        }
        manager->AddEntry(op_num, comm_idx, &input, output, std::move(actual_done), op_event);
    }