Add XlaCompiledFunction, a lightweight API for calling XLA computations that are

compiled down to functions. The API is based on a generic form of the original
AOT auto-generated header.

For AOT (tfcompile), this API has been slotted into the auto-generated header.

For JIT, a new XlaCompiledFunctionJit class has been added, which compiles a
tensorflow::GraphDef and allows the user to create XlaCompiledFunction objects.

XlaCompiledFunction contains optional metadata; mappings from arg/result names
to their index, and the program shape. This data is always available via JIT,
but only provided via AOT if the tfcompile --gen_name_to_index and
--gen_program_shape flags are set. We don't enable by default for AOT to keep
binary sizes smaller; the ProgramShape proto pulls in lots of code, and may also
be large.

PiperOrigin-RevId: 170811579

1 files changed, 2 insertions, 0 deletions

diff --git a/tensorflow/compiler/aot/tfcompile_main.cc b/tensorflow/compiler/aot/tfcompile_main.cc
index cc499c3284..6ab3d47418 100644
--- a/tensorflow/compiler/aot/tfcompile_main.cc
+++ b/tensorflow/compiler/aot/tfcompile_main.cc
@@ -94,6 +94,8 @@ Status Main(const MainFlags& flags) {
   TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_object,
                                        StringPiece(obj.data(), obj.size())));
   HeaderOpts header_opts;
+  header_opts.gen_name_to_index = flags.gen_name_to_index;
+  header_opts.gen_program_shape = flags.gen_program_shape;
   if (flags.cpp_class.empty()) {
     return errors::InvalidArgument("Must specify --cpp_class");
   }