Method ComputeAsync

tensorflow/core/kernels/function_ops.cc:310–416

}

void RemoteCallOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  FunctionLibraryRuntime* lib = ctx->function_library();
  OP_REQUIRES_ASYNC(ctx, lib != nullptr,
                    errors::Internal("No function library is provided."), done);

  const string& source_device = lib->device()->name();
  const Tensor* target;
  OP_REQUIRES_OK_ASYNC(ctx, ctx->input("target", &target), done);
  string target_device;
  OP_REQUIRES_OK_ASYNC(
      ctx,
      DeviceNameUtils::CanonicalizeDeviceName(target->scalar<tstring>()(),
                                              source_device, &target_device),
      done);

  AttrValueMap attr_values = func_.attr();

  FunctionLibraryRuntime::InstantiateOptions instantiate_opts;
  // TODO(b/141576771): Propagate ctxt's ConfigProto into opts.config_proto.
  instantiate_opts.target = target_device;

  FunctionTarget function_target = {target_device, lib};

  FunctionLibraryRuntime::Handle handle;
  {
    mutex_lock l(mu_);
    auto cached_entry = handle_cache_.find(function_target);
    if (cached_entry != handle_cache_.end()) {
      handle = cached_entry->second;
    } else {
      VLOG(1) << "Instantiating " << func_.name() << " on " << target_device;
      profiler::TraceMe activity(
          [&] {
            return strings::StrCat("RemoteCall: Instantiate: ", func_.name(),
                                   " on ", target_device);
          },
          profiler::TraceMeLevel::kInfo);
      OP_REQUIRES_OK_ASYNC(
          ctx,
          lib->Instantiate(func_.name(), AttrSlice(&attr_values),
                           instantiate_opts, &handle),
          done);
      auto insert_result = handle_cache_.insert({function_target, handle});
      CHECK(insert_result.second) << "Insert unsuccessful.";
      VLOG(1) << "Instantiated " << func_.name() << " on " << target_device
              << ", resulting in handle: " << handle << " flr: " << lib;
    }
  }

  OpInputList arguments;
  OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("args", &arguments), done);

  FunctionLibraryRuntime::Options opts;
  opts.runner = ctx->runner();
  opts.source_device = source_device;
  if (opts.source_device != target_device) {
    opts.remote_execution = true;
  }