Method ComputeAsync

tensorflow/core/kernels/partitioned_function_ops.cc:85–131

}

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

  // The function body's graph is placed and partitioned the first time
  // `ComputeAsync` is invoked; every subsequent invocation calls each
  // of the function shards yielded by partitioning.
  //
  // The partitioning step yields a set of devices on which to run the
  // function, and exactly one function shard is created for each device
  // Inputs and outputs are pinned to the local device, for simplicity.
  //
  // TODO(akshayka): Support re-sharding the function on subsequent calls,
  // via, e.g., virtual device annotations and a list of device names
  // supplied through an attribute.
  //
  // TODO(akshayka): Add a fastpath for functions that execute on a single
  // device.
  FunctionLibraryRuntime::Handle handle;
  // If we are instantiating the function, we can efficiently extract the
  // inputs while instantiating. Else, we extract them separately below.
  std::vector<Tensor> inputs;
  bool inputs_extracted = false;
  {
    mutex_lock l(mu_);
    auto it = handles_.find(lib);
    if (it == handles_.end()) {
      OP_REQUIRES_OK_ASYNC(ctx, Instantiate(lib, ctx, &inputs, &handle), done);
      inputs_extracted = true;
      handles_[lib] = handle;
    } else {
      handle = it->second;
    }
  }

  if (!inputs_extracted) {
    OpInputList args;
    OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("args", &args), done);
    inputs.reserve(args.size());
    for (const Tensor& tensor : args) {
      inputs.push_back(tensor);
    }
  }

  RunFunction(handle, inputs, lib, ctx, done);
}

Status PartitionedCallOp::FillOutputDevices(
    const FunctionLibraryRuntime& lib, const Device& cpu_device,