Method ComputeAsync

tensorflow/core/kernels/fuserecv_ops.cc:89–137

}

void FuseRecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  //LOG(INFO) << "FuseRecvOp::ComputeAsync";
  OP_REQUIRES(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."));

  // NOTE(rangeng.llb): recv args.
  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->output_alloc_attr(0);

  using namespace std::placeholders;
  Rendezvous::FuseDoneCallback done_cb = std::bind(
      [ctx](DoneCallback done,
            // Begin unbound arguments.
            const Status& s, const std::vector<Rendezvous::Args>& send_args,
            const Rendezvous::Args& recv_args, const std::vector<Tensor>& vals,
            const std::vector<bool>& is_deads) {
        ctx->SetStatus(s);
        if (s.ok()) {
          OpOutputList output;
          OP_REQUIRES_OK(ctx, ctx->output_list("tensors", &output));

          bool all_deads = true;
          for (int i = 0; i < vals.size(); ++i) {
            if (!is_deads[i]) {
              output.set(i, vals[i]);
              all_deads = false;
            }
          }
          // is_output_dead means all the outputs is dead.
          *ctx->is_output_dead() = all_deads;
        }
        done();
      },
      std::move(done), _1, _2, _3, _4, _5);

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (!(frame_iter == FrameAndIter(0, 0))) {
    // if frame_iter has changed, change the parsed keys.
    for (int i = 0; i < fuse_count_; ++i) {
      GetRendezvousKey(key_prefixs_[i], frame_iter, &parsed_keys_[i].buf_);
      OP_REQUIRES_OK_ASYNC(ctx,
                           Rendezvous::ParseKey(parsed_keys_[i].buf_, &parsed_keys_[i]),
                           done);
    }
  }
  ctx->rendezvous()->FuseRecvAsync(parsed_keys_, args, std::move(done_cb));
}

REGISTER_KERNEL_BUILDER(Name("_FuseRecv").Device(DEVICE_CPU), FuseRecvOp);
REGISTER_KERNEL_BUILDER(Name("_FuseRecv").Device(DEVICE_GPU), FuseRecvOp);