diff options
| author |  Mingsheng Hong <hongm@google.com> | 2018-03-19 22:29:33 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-19 22:33:26 -0700 |
| commit | 1a8258e0593270a8e2370517dff8faafce40a687 (patch) | |
| tree | e097a315ba6c21f8f2ac72c5af17564d73efbc4e /tensorflow/c/c_api_experimental.h | |
| parent | 407ddd1c0539cfc5d33ab2629230eab5a958b7d4 (diff) | |
Added infeed support for experimental C APIs associated with TPU graph rewrite.
This initial design of the C API is different from (and mostly higher level
than) the python API counterparts for infeed, in that the python API has
explicit graph construction APIs for generating infeed enqueue/dequeue
ops (e.g. split_inputs_and_generate_enqueue_ops() and generate_dequeue_op()),
while the C API takes an input graph and redirects all input nodes to feed the
infeed enqueue.
One requirement/restriction is that the input nodes in the TF
graph (e.g. Placeholder) must specify their tensor shapes, for infeed enqueue
and dequeue nodes to properly compile with XLA. The API for more general shape
support will be designed and implemented later.
PiperOrigin-RevId: 189693028
Diffstat (limited to 'tensorflow/c/c_api_experimental.h')
| -rw-r--r-- | tensorflow/c/c_api_experimental.h | 13 |
1 files changed, 11 insertions, 2 deletions
diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h index af65123131..2bad278d63 100644 --- a/tensorflow/c/c_api_experimental.h +++ b/tensorflow/c/c_api_experimental.h @@ -63,7 +63,15 @@ TF_CAPI_EXPORT extern void TF_EnableXLACompilation(TF_SessionOptions* options, // Sets up TPU execution, by rewriting the graph accordingly, and initializing // TPU system. // -// On success, returns a shutdown node to be used in a subsequent +// When `infeed_enqueue_node` is non-NULL and there are input tensors, rewrites +// the graph by adding the relevant infeed enqueue/dequeue ops, and returns the +// enqueue op in `infeed_enqueue_node` on success, so that user can run that +// node and feed input tensors. When there are no input tensors, +// `infeed_enqueue_node` is ignored, and user should not run that node later. +// TODO(hongm): In this case, we currently only support input tensors of dim 0 +// shape. Lift that constraint. +// +// On success, also returns a shutdown node to be used in a subsequent // TF_ShutdownTPUExecution(), and sets the new output nodes in // `new_output_nodes` for caller to fetch from. Must be called exactly once // before TF_SessionRun(). @@ -76,7 +84,8 @@ TF_CAPI_EXPORT extern void TF_EnableXLACompilation(TF_SessionOptions* options, TF_CAPI_EXPORT extern TF_Output TF_SetupTPUExecution( TF_Session* session, int num_input_nodes, const TF_Output* input_nodes, int num_output_nodes, const TF_Output* output_nodes, - TF_Output* new_output_nodes, TF_Status* status); + TF_Output* new_output_nodes, TF_Operation** infeed_enqueue_node, + TF_Status* status); // Shuts down TPU system. For any `session` where TF_SetupTPUExecution() has // been successfully called, this call must be made exactly once before the |