Add experimental Session::MakeCallable() API and implement it for DirectSession.

The intent of this new API matches the Python `tf.Session.make_callable()`
method: it splits the two roles of the `Session::Run()` method into separate
methods:

1. `Session::MakeCallable()` takes information about a subgraph (such as the
   names of nodes to feed and fetch), and prunes and optimizes that graph,
   returning a simple handle.
2. `Session::RunCallable()` takes that handle, plus any values to be fed,
   and executes the graph, returning whatever outputs are produced.

This split moves string processing off the critical path of running a
step. We also add a new method `Session::ReleaseCallable()` that makes
it possible to free the resources associated with a cached subgraph,
and could be useful for seldom-executed graphs such as initializers.

PiperOrigin-RevId: 188566635

1 files changed, 35 insertions, 0 deletions

diff --git a/tensorflow/core/public/session.h b/tensorflow/core/public/session.h
index 75ad50f6f2..d58c877cfd 100644
--- a/tensorflow/core/public/session.h
+++ b/tensorflow/core/public/session.h
@@ -195,6 +195,41 @@ class Session {
     return errors::Unimplemented(
         "LocalDeviceManager is not supported for this session.");
   }
+
+  /// \brief A handle to a subgraph, created with `Session::MakeCallable()`.
+  typedef int64 CallableHandle;
+
+  /// \brief Creates a `handle` for invoking the subgraph defined by
+  /// `callable_options`.
+  /// NOTE: This API is still experimental and may change.
+  virtual Status MakeCallable(const CallableOptions& callable_options,
+                              CallableHandle* out_handle) {
+    return errors::Unimplemented(
+        "MakeCallable is not supported for this session.");
+  }
+
+  /// \brief Invokes the subgraph named by `handle` with the given options and
+  /// input tensors.
+  ///
+  /// The order of tensors in `feed_tensors` must and `fetch_tensors` will
+  /// match the order of names in `CallableOptions::feed()` and
+  /// `CallableOptions::fetch()` when this subgraph was created.
+  /// NOTE: This API is still experimental and may change.
+  virtual Status RunCallable(CallableHandle handle,
+                             const std::vector<Tensor>& feed_tensors,
+                             std::vector<Tensor>* fetch_tensors,
+                             RunMetadata* run_metadata) {
+    return errors::Unimplemented(
+        "RunCallable is not supported for this session.");
+  }
+
+  /// \brief Releases resources associated with the given `handle` in this
+  /// session.
+  /// NOTE: This API is still experimental and may change.
+  virtual Status ReleaseCallable(CallableHandle handle) {
+    return errors::Unimplemented(
+        "ReleaseCallable is not supported for this session.");
+  }
 };
 
 /// \brief Create a new session with the given options.