Sets up a rendezvous for eager op execution.

PiperOrigin-RevId: 165019327

tensorflow/c/eager/BUILD

@@ -15,6 +15,7 @@ cc_library(
         "//tensorflow/c:c_api_internal",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",

tensorflow/c/eager/c_api.cc

@@ -27,6 +27,8 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device_factory.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
+#include "tensorflow/core/framework/rendezvous.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/lib/core/refcount.h"
 #include "tensorflow/core/lib/gtl/map_util.h"
@@ -53,6 +55,7 @@ struct TFE_Context {
 
   // TFE_Context is an extension of TF_Session. And TF_Session needs a TF_Graph.
   TF_Session* session;
+  tensorflow::Rendezvous* rendezvous;
 
   tensorflow::mutex functions_mu;
   tensorflow::FunctionLibraryDefinition func_lib_def GUARDED_BY(functions_mu){
@@ -135,6 +138,8 @@ TFE_Context* TFE_NewContext(const TF_SessionOptions* opts, TF_Status* status) {
         ret->session->device_mgr, opts->options.env, ret->devices()[i],
         TF_GRAPH_DEF_VERSION, &ret->func_lib_def, {});
   }
+  ret->rendezvous =
+      new tensorflow::IntraProcessRendezvous(ret->session->device_mgr);
 
   return ret;
 }
@@ -145,6 +150,7 @@ void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) {
   TF_Graph* graph = ctx->session->graph;
   TF_DeleteSession(ctx->session, status);
   TF_DeleteGraph(graph);
+  ctx->rendezvous->Unref();
   delete ctx;
 }
 
@@ -470,7 +476,7 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals,
       tensorflow::gtl::FindPtrOrNull(ctx->kernel_cache, cache_key);
   if (kernel == nullptr) {
     const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef();
-    kernel = new tensorflow::KernelAndDevice();
+    kernel = new tensorflow::KernelAndDevice(ctx->rendezvous);
     if (!op->is_function()) {
       status->status =
           tensorflow::KernelAndDevice::InitOp(device, ndef, kernel);

tensorflow/c/eager/runtime.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/c/eager/runtime.h"
 
 #include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -270,6 +271,7 @@ Status KernelAndDevice::Run(std::vector<Tensor>* input_tensors,
   params.output_attr_array = gtl::vector_as_array(&out_attrs);
   params.function_library = flib_;
   params.slice_reader_cache = &slice_reader_cache_;
+  params.rendezvous = rendez_;
   // TODO(apassos): use a thread pool.
   std::function<void(std::function<void()>)> runner =
       [](std::function<void()> f) { f(); };

tensorflow/c/eager/runtime.h

@@ -174,7 +174,8 @@ class KernelAndDevice {
   static Status InitFn(const NodeDef& ndef, FunctionLibraryRuntime* flib,
                        KernelAndDevice* out);
 
-  KernelAndDevice() : device_(nullptr), flib_(nullptr) {}
+  KernelAndDevice(tensorflow::Rendezvous* rendez)
+      : device_(nullptr), flib_(nullptr), rendez_(rendez) {}
 
   // TODO(ashankar): Handle list-valued inputs.
   Status Run(std::vector<Tensor>* inputs, std::vector<Tensor>* outputs);
@@ -186,6 +187,7 @@ class KernelAndDevice {
   tensorflow::Device* device_;
   tensorflow::FunctionLibraryRuntime* flib_;
   tensorflow::checkpoint::TensorSliceReaderCacheWrapper slice_reader_cache_;
+  tensorflow::Rendezvous* rendez_;
 };
 
 }  // namespace tensorflow

tensorflow/c/eager/runtime_test.cc

@@ -70,7 +70,7 @@ TEST(KernelAndDevice, Run) {
                    .NumInputs(inputs.size())
                    .BuildNodeDef());
   std::unique_ptr<Device> device(CPUDevice());
-  KernelAndDevice kernel;
+  KernelAndDevice kernel(nullptr);
   Status s = KernelAndDevice::InitOp(device.get(), ndef, &kernel);
   ASSERT_TRUE(s.ok()) << s;
   std::vector<Tensor> outputs;