Method Compute

tensorflow/core/kernels/slice_sendrecv_ops.cc:45–86

}

void SliceSendOp::Compute(OpKernelContext* ctx) {
  OP_REQUIRES(
    ctx, ctx->rendezvous() != nullptr,
    errors::Internal("Op kernel context needs to provide a rendezvous."));

  const Tensor& input_t = ctx->input(0);
  FrameAndIter frame_iter = \
    slice_sendrecv::GetFrameAndIter(ctx, hostmem_sendrecv_);

  // send total_bytes.
  OP_REQUIRES_OK(ctx, SendTotalBytes(ctx, frame_iter, input_t));
  // if input is dead, only send total_bytes dead tensor.
  if (ctx->is_input_dead()) {
    return;
  }

  // if total bytes is smaller than slice size, send directly.
  if (input_t.TotalBytes() <= slice_size_) {
    Rendezvous::Args args;
    args.device_context = ctx->op_device_context();
    args.alloc_attrs = ctx->input_alloc_attr(0);

    Rendezvous::ParsedKey parsed_key;
    slice_sendrecv::GetSliceRendezvousKey(key_prefix_, "_transfer_data",
                                          frame_iter, &parsed_key.buf_);
    VLOG(2) << "SliceSend " << parsed_key.buf_;
    OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_key.buf_, &parsed_key));
    OP_REQUIRES_OK(ctx, ctx->rendezvous()->Send(parsed_key, args, input_t,
                                                ctx->is_input_dead()));
    return;
  }

  // send shape.
  OP_REQUIRES_OK(ctx, SendShape(ctx, frame_iter, input_t));

  // send data.
  if (dtype_ == DT_STRING) {
    OP_REQUIRES_OK(ctx, SendString(ctx, frame_iter, input_t));
  } else {
    OP_REQUIRES_OK(ctx, SendBasicType(ctx, frame_iter, input_t));
  }
}

Status SliceSendOp::SendTotalBytes(OpKernelContext* ctx,
                                   const FrameAndIter& frame_iter,