Method FuseRecvOp

tensorflow/core/kernels/fuserecv_ops.cc:54–87

}

FuseRecvOp::FuseRecvOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx), fuse_count_(0) {
  // send devices are not single
  std::vector<string> tensor_names;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_names", &tensor_names));

  std::vector<string> send_devices;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_devices", &send_devices));

  std::vector<int64> send_device_incarnations;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("send_device_incarnations",
                                   &send_device_incarnations));

  std::vector<string> recv_devices;
  OP_REQUIRES_OK(ctx, ctx->GetAttr("recv_devices", &recv_devices));

  fuse_count_ = tensor_names.size();
  key_prefixs_.resize(fuse_count_);
  parsed_keys_.resize(fuse_count_);

  for (int i = 0; i < fuse_count_; ++i) {
    key_prefixs_[i] = GetRendezvousKeyPrefix(
        send_devices[i], recv_devices[i],
        static_cast<uint64>(send_device_incarnations[i]),
        tensor_names[i]);
    // The vast majority of Recv nodes are outside any loop context, so
    // proactively cache the rendezvous key for the top-level.
    GetRendezvousKey(key_prefixs_[i], {0, 0}, &parsed_keys_[i].buf_);
    OP_REQUIRES_OK(ctx, Rendezvous::ParseKey(parsed_keys_[i].buf_, &parsed_keys_[i]));
  }

  if (!ctx->GetAttr("_hostmem_sendrecv", &hostmem_sendrecv_).ok()) {
    hostmem_sendrecv_ = false;
  }
}

void FuseRecvOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  //LOG(INFO) << "FuseRecvOp::ComputeAsync";