Provide a step container to OpKernelContexts when running eagerly.

This lets us run some ops that require step containers (e.g.
TensorArray-related ops).

 Before change:
 Benchmark                    Time(ns)        CPU(ns)     Iterations
-------------------------------------------------------------------
 BM_CreateGraph                  61292          80915           8581
 BM_RunGraph                      7899          13398          51251
 BM_CreateAndDestroySession       2588           2594         269838
 BM_KernelAndDeviceInit           2971           2976         235908
 BM_KernelAndDeviceRun             505            506        1000000

 After change:
 Benchmark                    Time(ns)        CPU(ns)     Iterations
 -------------------------------------------------------------------
 BM_CreateGraph                  78295         105539           8698
 BM_RunGraph                      9907          16988          47908
 BM_CreateAndDestroySession       2773           2778         247635
 BM_KernelAndDeviceInit           2678           2682         270054
 BM_KernelAndDeviceRun             553            554        1000000

PiperOrigin-RevId: 198125630

diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.cc b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
index a63b2b9..2a43a31 100644 (file)
--- a/tensorflow/core/common_runtime/eager/kernel_and_device.cc
+++ b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #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/framework/resource_mgr.h"
 #include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/gtl/map_util.h"
@@ -86,6 +87,11 @@ Status KernelAndDevice::Run(std::vector<Tensor>* input_tensors,
       [](std::function<void()> f) { f(); };
   params.runner = &runner;
 
+  ScopedStepContainer step_container(0, [this](const string& name) {
+    device_->resource_manager()->Cleanup(name).IgnoreError();
+  });
+  params.step_container = &step_container;
+
   OpKernelContext context(&params);
 
   if (kernel_->def().op() == "_Recv") {