Method ComputeAsync

tensorflow/core/kernels/sendrecv_ops.cc:163–193

}  // namespace

void RecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  OP_REQUIRES_ASYNC(
      ctx, ctx->rendezvous() != nullptr,
      errors::Internal("Op kernel context needs to provide a rendezvous."),
      done);

  Rendezvous::Args args;
  args.device_context = ctx->op_device_context();
  args.alloc_attrs = ctx->output_alloc_attr(0);
  if (ctx->is_eager()) {
    // NOTE(fishx): Only set cancellation_manager in eager mode. Because in
    // Tensorflow 1.x, session (or graph_mgr) will abort the underlying
    // rendezvous if it encounters any error.
    args.cancellation_manager = ctx->cancellation_manager();
  }

  FrameAndIter frame_iter = GetFrameAndIter(ctx, hostmem_sendrecv_);
  if (frame_iter == FrameAndIter(0, 0)) {
    VLOG(2) << "Recv " << parsed_key_.buf_;
    ctx->rendezvous()->RecvAsync(parsed_key_, args,
                                 make_recv_callback(ctx, std::move(done)));
  } else {
    Rendezvous::ParsedKey in_loop_parsed;
    GetRendezvousKey(key_prefix_, frame_iter, &in_loop_parsed.buf_);
    VLOG(2) << "Recv " << in_loop_parsed.buf_;
    OP_REQUIRES_OK_ASYNC(
        ctx, Rendezvous::ParseKey(in_loop_parsed.buf_, &in_loop_parsed), done);
    ctx->rendezvous()->RecvAsync(in_loop_parsed, args,
                                 make_recv_callback(ctx, std::move(done)));
  }
}

REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_CPU), RecvOp);
REGISTER_KERNEL_BUILDER(Name("_Recv").Device(DEVICE_DEFAULT), RecvOp);