Add ability to fetch return nodes and unused input mappings from C API GraphDef import

This change introduces yet another ImportGraphDef function to the C
API (TF_GraphImportGraphDefWithResults), but this one has extensible
return values so we shouldn't have to add more in the future.

This change also modifies the ImportGraphDef C interface to manage all
string data for the user.

PiperOrigin-RevId: 173894710

1 files changed, 56 insertions, 1 deletions

diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h
index 1e8bfdc7b0..ca5c934634 100644
--- a/tensorflow/c/c_api.h
+++ b/tensorflow/c/c_api.h
@@ -914,7 +914,62 @@ TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsAddReturnOutput(
 TF_CAPI_EXPORT extern int TF_ImportGraphDefOptionsNumReturnOutputs(
     const TF_ImportGraphDefOptions* opts);
 
+// Add an operation in `graph_def` to be returned via the `return_opers` output
+// parameter of TF_GraphImportGraphDef().
+TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsAddReturnOperation(
+    TF_ImportGraphDefOptions* opts, const char* oper_name);
+
+// Returns the number of return operations added via
+// TF_ImportGraphDefOptionsAddReturnOperation().
+TF_CAPI_EXPORT extern int TF_ImportGraphDefOptionsNumReturnOperations(
+    const TF_ImportGraphDefOptions* opts);
+
+// TF_ImportGraphDefResults holds results that are generated by
+// TF_GraphImportGraphDefWithResults().
+typedef struct TF_ImportGraphDefResults TF_ImportGraphDefResults;
+
+// Fetches the return outputs requested via
+// TF_ImportGraphDefOptionsAddReturnOutput(). The number of fetched outputs is
+// returned in `num_outputs`. The array of return outputs is returned in
+// `outputs`. `*outputs` is owned by and has the lifetime of `results`.
+TF_CAPI_EXPORT extern void TF_ImportGraphDefResultsReturnOutputs(
+    TF_ImportGraphDefResults* results, int* num_outputs, TF_Output** outputs);
+
+// Fetches the return operations requested via
+// TF_ImportGraphDefOptionsAddReturnOperation(). The number of fetched
+// operations is returned in `num_opers`. The array of return operations is
+// returned in `opers`. `*opers` is owned by and has the lifetime of `results`.
+TF_CAPI_EXPORT extern void TF_ImportGraphDefResultsReturnOperations(
+    TF_ImportGraphDefResults* results, int* num_opers, TF_Operation*** opers);
+
+// Fetches any input mappings requested via
+// TF_ImportGraphDefOptionsAddInputMapping() that weren't used as input to any
+// node in the imported graph def. The number of fetched mappings is returned in
+// `num_unused_input_mappings`. The array of each mapping's source node name is
+// returned in `src_names`, and the array of each mapping's source index is
+// returned in `src_indexes`.
+//
+// `*src_names`, `*src_indexes`, and the memory backing each string in
+// `src_names` are owned by and have the lifetime of `results`.
+TF_CAPI_EXPORT extern void TF_ImportGraphDefResultsUnusedInputMappings(
+    TF_ImportGraphDefResults* results, int* num_unused_input_mappings,
+    const char*** src_names, int** src_indexes);
+
+// Deletes a results object returned by TF_GraphImportGraphDefWithResults().
+TF_CAPI_EXPORT extern void TF_DeleteImportGraphDefResults(
+    TF_ImportGraphDefResults* results);
+
+// Import the graph serialized in `graph_def` into `graph`.  Returns nullptr and
+// a bad status on error. Otherwise, returns a populated
+// TF_ImportGraphDefResults instance. The returned instance must be deleted via
+// TF_DeleteImportGraphDefResults().
+TF_CAPI_EXPORT extern TF_ImportGraphDefResults*
+TF_GraphImportGraphDefWithResults(TF_Graph* graph, const TF_Buffer* graph_def,
+                                  const TF_ImportGraphDefOptions* options,
+                                  TF_Status* status);
+
 // Import the graph serialized in `graph_def` into `graph`.
+// Convenience function for when only return outputs are needed.
 //
 // `num_return_outputs` must be the number of return outputs added (i.e. the
 // result of TF_ImportGraphDefOptionsNumReturnOutputs()).  If
@@ -926,7 +981,7 @@ TF_CAPI_EXPORT extern void TF_GraphImportGraphDefWithReturnOutputs(
     int num_return_outputs, TF_Status* status);
 
 // Import the graph serialized in `graph_def` into `graph`.
-// Convenience function for when no return outputs have been added.
+// Convenience function for when no results are needed.
 TF_CAPI_EXPORT extern void TF_GraphImportGraphDef(
     TF_Graph* graph, const TF_Buffer* graph_def,
     const TF_ImportGraphDefOptions* options, TF_Status* status);