Commit 067be1a
Part 2 of Swift<->TF sends/recvs: support receiving tensors in TF from
Swift via direct session.
The changes are:
1. Added a TF experimental C API for Swift host to enqueue a tensor for sending
to TF. Again, the C APIs can be removed once the Fifo-queue based design
proves stable later.
2. TFLowerGraph is extended to generate Fifo related nodes for TF to receive
tensors. This is similar to the extension for TF to send tensors.
3. TFPartition is extended to support host send (createHostSend()), which does
the tensor send via a new protocol method TensorSendableReceivable.sendToDevice().
The main complexity is in sending a scalar, where a new protocol method
TensorSendableReceivable.createScalarTensor() is called to first create a tensor
out of it, and then send it over to TF.
Also removed code for protocol conformance on AccelerableByTensorFlow. Instead
have compiler look up that conformance from the SILFunction on sending/receiving
tensors.
AccelerableByTensorFlow could be removed from the compiler-known protocol list
now, but we'll defer that till things can stabilized more (in the past this
protocol has been added to and removed from the list at different times).
PiperOrigin-RevId: 1955394361 parent 150089e commit 067be1a
2 files changed
Lines changed: 69 additions & 2 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -8407,3 +8407,51 @@ TF_Tensor* TF_DequeueNamedTensor(TF_Session* session, int tensor_id, | |||
| 8407 | 8407 | } | |
| 8408 | 8408 | return ret; | |
| 8409 | 8409 | } | |
| 8410 | + | ||
| 8411 | + void TF_EnqueueNamedTensor(TF_Session* session, int tensor_id, | ||
| 8412 | + TF_Tensor* tensor, TF_Status* status) { | ||
| 8413 | + assert(session); | ||
| 8414 | + { | ||
| 8415 | + tensorflow::mutex_lock c(session->graph->mu); | ||
| 8416 | + if (VLOG_IS_ON(1)) { | ||
| 8417 | + VLOG(1) << "Enqueuing named tensor with id " << tensor_id | ||
| 8418 | + << ", with input graph: " | ||
| 8419 | + << session->graph->graph.ToGraphDefDebug().DebugString(); | ||
| 8420 | + tensorflow::Tensor internal_tensor; | ||
| 8421 | + if (tensorflow::TF_TensorToTensor(tensor, &internal_tensor).ok()) { | ||
| 8422 | + VLOG(1) << "Enqueu'ing tensor content: " | ||
| 8423 | + << internal_tensor.DebugString(); | ||
| 8424 | + } | ||
| 8425 | + } | ||
| 8426 | + } | ||
| 8427 | + | ||
| 8428 | + TF_Operation* enqueue_op = TF_GraphOperationByName( | ||
| 8429 | + session->graph, | ||
| 8430 | + tensorflow::strings::StrCat("fifo_queue_enqueue_", tensor_id).c_str()); | ||
| 8431 | + if (enqueue_op == nullptr) { | ||
| 8432 | + status->status = tensorflow::errors::Internal( | ||
| 8433 | + "Unable to find the enqueue node in the TF graph."); | ||
| 8434 | + return; | ||
| 8435 | + } | ||
| 8436 | + | ||
| 8437 | + TF_Operation* placeholder_op = TF_GraphOperationByName( | ||
| 8438 | + session->graph, | ||
| 8439 | + tensorflow::strings::StrCat("arg_tensor_enqueue_", tensor_id).c_str()); | ||
| 8440 | + if (placeholder_op == nullptr) { | ||
| 8441 | + status->status = tensorflow::errors::Internal( | ||
| 8442 | + "Unable to find the placeholder node as input to enqueue in the TF " | ||
| 8443 | + "graph."); | ||
| 8444 | + return; | ||
| 8445 | + } | ||
| 8446 | + | ||
| 8447 | + VLOG(1) << "Running the enqueue op"; | ||
| 8448 | + TF_Output input{placeholder_op, 0}; | ||
| 8449 | + TF_SessionRun(session, /*run_options*/ nullptr, | ||
| 8450 | + // input related parameters | ||
| 8451 | + /*inputs*/ &input, /*input_values*/ &tensor, /*ninputs*/ 1, | ||
| 8452 | + // output related parameters | ||
| 8453 | + /*outputs*/ nullptr, /*output_values*/ nullptr, /*noutputs*/ 0, | ||
| 8454 | + /*targets*/ &enqueue_op, /*ntargets*/ 1, | ||
| 8455 | + /*run_metadata*/ nullptr, status); | ||
| 8456 | + VLOG(1) << "Enqueuing is done."; | ||
| 8457 | + } | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -87,15 +87,34 @@ TF_CAPI_EXPORT extern TF_Operation* TF_MakeFileBasedIteratorGetNextWithDatasets( | |||
| 87 | 87 | unsigned char is_mnist, TF_Status* status); | |
| 88 | 88 | ||
| 89 | 89 | // On success, dequeues a tensor from a TF-managed FifoQueue given by | |
| 90 | - // `tensor_id`, associated with `session`. Caller must call TF_DeleteTensor() | ||
| 91 | - // over the returned tensor. If the queue is empty, this call is blocked. | ||
| 90 | + // `tensor_id`, associated with `session`. There must be a graph node named | ||
| 91 | + // "fifo_queue_dequeue_<tensor_id>", to be executed by this API call. | ||
| 92 | + | ||
| 93 | + // Caller must call TF_DeleteTensor() over the returned tensor. If the queue is | ||
| 94 | + // empty, this call is blocked. | ||
| 92 | 95 | // | |
| 93 | 96 | // Tensors are enqueued via the corresponding TF enqueue op. | |
| 94 | 97 | // TODO(hongm): Add support for `timeout_ms`. | |
| 95 | 98 | TF_CAPI_EXPORT extern TF_Tensor* TF_DequeueNamedTensor(TF_Session* session, | |
| 96 | 99 | int tensor_id, | |
| 97 | 100 | TF_Status* status); | |
| 98 | 101 | ||
| 102 | + // On success, enqueues `tensor` into a TF-managed FifoQueue given by | ||
| 103 | + // `tensor_id`, associated with `session`. There must be a graph node named | ||
| 104 | + // "fifo_queue_enqueue_<tensor_id>", to be executed by this API call. It reads | ||
| 105 | + // from a placeholder node "arg_tensor_enqueue_<tensor_id>". | ||
| 106 | + // | ||
| 107 | + // `tensor` is still owned by the caller. This call will be blocked if the queue | ||
| 108 | + // has reached its capacity, and will be unblocked when the queued tensors again | ||
| 109 | + // drop below the capacity due to dequeuing. | ||
| 110 | + // | ||
| 111 | + // Tensors are dequeued via the corresponding TF dequeue op. | ||
| 112 | + // TODO(hongm): Add support for `timeout_ms`. | ||
| 113 | + TF_CAPI_EXPORT extern void TF_EnqueueNamedTensor(TF_Session* session, | ||
| 114 | + int tensor_id, | ||
| 115 | + TF_Tensor* tensor, | ||
| 116 | + TF_Status* status); | ||
| 117 | + | ||
| 99 | 118 | #ifdef __cplusplus | |
| 100 | 119 | } /* end extern "C" */ | |
| 101 | 120 | #endif | |
0 commit comments