Avoid an extra copy when receiving a tensor over grpc.

A benchmark that sends around 400KB tensors speeds up by ~10-15%.

Details:
* RecvTensor RPC yields a TensorResponse instead of a RecvTensorResponse.
* TensorResponse has custom decoding code that can decode contents
  directly into a Tensor's underlying buffer.
* Wire protocol is unchanged.
Change: 129809238

1 files changed, 3 insertions, 6 deletions

diff --git a/tensorflow/core/distributed_runtime/worker_interface.h b/tensorflow/core/distributed_runtime/worker_interface.h
index 3dd00c559d..4ae4e76ea4 100644
--- a/tensorflow/core/distributed_runtime/worker_interface.h
+++ b/tensorflow/core/distributed_runtime/worker_interface.h
@@ -29,10 +29,8 @@ namespace tensorflow {
 // Status callback.
 typedef std::function<void(const Status&)> StatusCallback;
 
-// Allocator callback for out-of-band transfers.
-class TensorShape;
-typedef std::function<void*(size_t, const DataType&, const TensorShape&)>
-    TensorBufAllocator;
+// Custom decoder for a response to RecvTensorAsync.
+class TensorResponse;
 
 // Interface for talking with the TensorFlow Worker service.
 class WorkerInterface {
@@ -65,8 +63,7 @@ class WorkerInterface {
 
   virtual void RecvTensorAsync(CallOptions* opts,
                                const RecvTensorRequest* request,
-                               RecvTensorResponse* response,
-                               TensorBufAllocator allocator,
+                               TensorResponse* response,
                                StatusCallback done) = 0;
 
   virtual void LoggingAsync(const LoggingRequest* request,