Use the new grpc::ByteBuffer::Swap() operation to avoid a memory allocation in

tensorflow::grpc::EncodeRecvTensorResponseToByteBuffer() and
tensorflow::grpc::EncodeTensorToByteBuffer()

The Swap() operation allows the assignment to the *result ByteBuffer
to be performed without the allocation performed by grpc::ByteBuffer::operator=().

PiperOrigin-RevId: 165504787

1 files changed, 4 insertions, 2 deletions

diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_tensor_coding.cc b/tensorflow/core/distributed_runtime/rpc/grpc_tensor_coding.cc
index 2cd3a33704..5639691804 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_tensor_coding.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_tensor_coding.cc
@@ -34,7 +34,8 @@ void EncodeRecvTensorResponseToByteBuffer(const RecvTensorResponse& proto,
   ::grpc::Slice slice(proto.ByteSizeLong());
   proto.SerializeWithCachedSizesToArray(
       const_cast<uint8*>(reinterpret_cast<const uint8*>(slice.begin())));
-  *result = ::grpc::ByteBuffer(&slice, 1);
+  ::grpc::ByteBuffer tmp(&slice, 1);
+  result->Swap(&tmp);
 }
 
 // We generate a RecvTensorResponse protocol buffer encoding into "*result",
@@ -237,7 +238,8 @@ void EncodeTensorToByteBuffer(bool is_dead, const Tensor& val,
     }
     CHECK_EQ(total_bytes, expected_size);
 
-    *result = ::grpc::ByteBuffer(&slices[0], num_slices);
+    ::grpc::ByteBuffer tmp(&slices[0], num_slices);
+    result->Swap(&tmp);
   }
 }