1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/distributed_runtime/rpc/grpc_tensor_coding.h"
17
18 #include "grpcpp/support/byte_buffer.h"
19 #include "grpcpp/support/slice.h"
20 #include "absl/flags/flag.h"
21 #include "tensorflow/core/common_runtime/dma_helper.h"
22 #include "tensorflow/core/framework/tensor.h"
23 #include "tensorflow/core/framework/tensor.pb.h"
24 #include "tensorflow/core/framework/tensor_reference.h"
25 #include "tensorflow/core/framework/tensor_shape.pb.h"
26 #include "tensorflow/core/lib/gtl/inlined_vector.h"
27 #include "tensorflow/core/lib/io/proto_encode_helper.h"
28 #include "tensorflow/core/platform/env.h"
29 #include "tensorflow/core/protobuf/worker.pb.h"
30
31 // (Omitted internal-only flag)
32
33 namespace tensorflow {
34 namespace grpc {
35
EncodeRecvTensorResponseToByteBuffer(const RecvTensorResponse & proto,::grpc::ByteBuffer * result)36 void EncodeRecvTensorResponseToByteBuffer(const RecvTensorResponse& proto,
37 ::grpc::ByteBuffer* result) {
38 ::grpc::Slice slice(proto.ByteSizeLong());
39 proto.SerializeWithCachedSizesToArray(
40 const_cast<uint8*>(reinterpret_cast<const uint8*>(slice.begin())));
41 ::grpc::ByteBuffer tmp(&slice, 1);
42 result->Swap(&tmp);
43 }
44
45 // We generate a RecvTensorResponse protocol buffer encoding into "*result",
46 // but where possible, we share the underlying Tensor buffer for "val", to
47 // avoid an extra copy.
48 //
49 // We hand-encode the protocol buffer data in the following order, as follows:
50 //
51 // Let R be a RecvTensorResponse object we want to encode, logically
52 // constructed by filling in data from "is_dead" and "val" and filling
53 // in a few other fields as well.
54 //
55 // (Letters here are used in the code to refer back to which part of the
56 // encoding the code is generating).
57 //
58 // A: <protocol buffer encoding of fields except R.tensor()>
59 // B1: <tag encoding for RecvTensorResponse::tensor>
60 // B2: <varint32 length of R.tensor() sub message>
61 // C: <protocol buffer encoding of R.tensor() except for
62 // R.tensor().tensor_content()>
63 // D1: <tag encoding for TensorProto::tensor_content>
64 // D2: <varint32 length of R.tensor().tensor_content() data>
65 // E: <actual data for val's representation>
66 //
67 // If the tensor data is up to "kLargeTensorBytes", then A
68 // through E will all be encoded into "*result" in a single grpc::Slice.
69 //
70 // If the tensor data is larger than "kLargeTensorBytes", then A through
71 // D2 will be encoded in one grpc::Slice, and E will be encoded in a second
72 // grpc::Slice that points to the backing store for the tensor data, to avoid
73 // copying the tensor data (and the grpc::Slice setup will be arrange so as
74 // to dereference the underlying tensor data buffer when it is no longer
75 // needed in the "*result" ByteBuffer).
VarLengthEncodingSize(uint32 tag,size_t bytes)76 static int VarLengthEncodingSize(uint32 tag, size_t bytes) {
77 return core::VarintLength(tag << 3) + core::VarintLength(bytes) + bytes;
78 }
79
80 // Returns an upper bound in bytes of the protocol buffer encoding of
81 // the "skeleton" of "val" (all the data needed for dtype and the shape,
82 // but not the actual contents of "val").
SkeletonEncodingSizeUpperBound(const Tensor & val)83 static int SkeletonEncodingSizeUpperBound(const Tensor& val) {
84 static const int kVarintMax64 = 10; // Max length of varint64 encoding
85 const int ndims = val.shape().dims();
86 return (2 * kVarintMax64) + // dtype
87 (ndims * (4 * kVarintMax64)); // Shape: 4 varints per dim
88 }
89
90 // Encode the skeleton for "val" (the encoded TensorProto contents
91 // (dtype and shape, but not the actual data) into "*e". The backing
92 // store for "*e" must be of appropriate size to hold this encoding.
EncodeSkeleton(const Tensor & val,io::ProtoEncodeHelper * e)93 static void EncodeSkeleton(const Tensor& val, io::ProtoEncodeHelper* e) {
94 // Encode val.dtype()
95 e->WriteUint64(TensorProto::kDtypeFieldNumber, val.dtype());
96
97 // Compute length of val.shape() proto encoding
98 const int ndims = val.shape().dims();
99 int tensor_shape_bytes = 0;
100 for (int d = 0; d < ndims; d++) {
101 int64 dim_size = val.shape().dim_size(d);
102 tensor_shape_bytes +=
103 2 + // TensorShapeProto dim tag + varintlength of submessage
104 1 + // TensorShapeProto_Dim::kSizeFieldNumber
105 core::VarintLength(dim_size);
106 }
107
108 if (tensor_shape_bytes > 0) {
109 e->WriteVarlengthBeginning(TensorProto::kTensorShapeFieldNumber,
110 tensor_shape_bytes);
111 // Encode val.shape()
112 for (int d = 0; d < ndims; d++) {
113 int64 dim_size = val.shape().dim_size(d);
114 int64 dim_varlen = 1 + // TensorShapeProto_Dim::kSizeFieldNumber
115 core::VarintLength(dim_size);
116 e->WriteVarlengthBeginning(TensorShapeProto::kDimFieldNumber, dim_varlen);
117 e->WriteUint64(TensorShapeProto_Dim::kSizeFieldNumber, dim_size);
118 }
119 }
120
121 #ifndef NDEBUG
122 {
123 // Debug-mode only check to make sure the encoding above is
124 // identical to the auto-generated protocol buffer encoding.
125 TensorProto skeleton;
126 skeleton.set_dtype(val.dtype());
127 val.shape().AsProto(skeleton.mutable_tensor_shape());
128 string tensor_except_contents; // tensor() field except contents
129 skeleton.AppendToString(&tensor_except_contents);
130 TensorProto skeleton2;
131 skeleton2.ParseFromString(string(e->data(), e->size()));
132 string out;
133 skeleton.AppendToString(&out);
134 DCHECK_EQ(tensor_except_contents, out) << skeleton.DebugString() << " vs\n"
135 << skeleton2.DebugString();
136 }
137 #endif
138 }
139
EncodeTensorToByteBuffer(bool is_dead,const Tensor & val,bool require_ack,::grpc::ByteBuffer * result)140 void EncodeTensorToByteBuffer(bool is_dead, const Tensor& val, bool require_ack,
141 ::grpc::ByteBuffer* result) {
142 const int kLargeTensorBytes = 1024;
143 RecvTensorResponse response;
144 if (is_dead) {
145 response.set_is_dead(is_dead);
146 }
147 response.set_require_ack(require_ack);
148 response.set_send_start_micros(Env::Default()->NowMicros());
149 if (!DataTypeCanUseMemcpy(val.dtype())) {
150 // Straightforward but slow path for complicated kinds of tensor data
151 // TODO(jeff,sanjay): If this becomes an issue, we could
152 // go directly from val -> ByteBuffer, with some effort.
153 val.AsProtoTensorContent(response.mutable_tensor());
154
155 // Encode full protocol buffer to a ByteBuffer
156 EncodeRecvTensorResponseToByteBuffer(response, result);
157 } else {
158 // skeleton is the encoded TensorProto contents (dtype and shape), but
159 // not the actual data
160 gtl::InlinedVector<char, 128> skeleton(SkeletonEncodingSizeUpperBound(val));
161 io::ProtoEncodeHelper e_skeleton(skeleton.data(), skeleton.size());
162 EncodeSkeleton(val, &e_skeleton);
163
164 StringPiece tdata = val.tensor_data();
165 uint32 overall_tensor_proto_bytesize =
166 (e_skeleton.size() +
167 VarLengthEncodingSize(TensorProto::kTensorContentFieldNumber,
168 tdata.size()));
169 string header; // All of RecvTensorResponse except the tensor() field
170 response.AppendToString(&header);
171
172 size_t expected_size =
173 (header.size() +
174 VarLengthEncodingSize(RecvTensorResponse::kTensorFieldNumber,
175 overall_tensor_proto_bytesize));
176 // If "share_tensor_slice_memory == false", we copy the tensor data to
177 // the end of the buffer we are preparing that holds the rest of the
178 // RecvTensorResponse protocol buffer.
179 //
180 // If "share_tensor_slice_memory == true", we arrange to share the
181 // backing store of the data by creating a slice that also points to the
182 // backing store, with appropriate reference counts to keep the
183 // backing store alive as needed.
184 //
185 // We enable this behavior if the tensor is large.
186 bool share_tensor_slice_memory = (tdata.size() > kLargeTensorBytes);
187
188 // (Omitted internal-only conditional)
189
190 size_t encoder_size = expected_size - tdata.size();
191
192 // Encode all but the actual "tdata", but including the tag and
193 // varlength header for the "tdata"
194 gtl::InlinedVector<char, 1024> space(encoder_size);
195 io::ProtoEncodeHelper e(space.data(), space.size());
196 // (A)
197 e.WriteRawBytes(header);
198
199 // (B1) & (B2)
200 e.WriteVarlengthBeginning(RecvTensorResponse::kTensorFieldNumber,
201 overall_tensor_proto_bytesize);
202 // (C)
203 e.WriteRawBytes(StringPiece(e_skeleton.data(), e_skeleton.size()));
204 // (D1) & (D2)
205 e.WriteVarlengthBeginning(TensorProto::kTensorContentFieldNumber,
206 tdata.size());
207
208 // All but the tensor backing store are serialized now
209
210 // Now allocate memory and put into the ByteBuffer
211 ::grpc::Slice slices[2];
212 int num_slices = 0;
213 {
214 size_t slice_len =
215 e.size() + (share_tensor_slice_memory ? 0 : tdata.size());
216 slices[0] = ::grpc::Slice(slice_len);
217 memcpy(const_cast<uint8_t*>(slices[0].begin()), e.data(), e.size());
218 if (!share_tensor_slice_memory) {
219 // (E)
220 memcpy(const_cast<uint8_t*>(slices[0].begin()) + e.size(), tdata.data(),
221 tdata.size());
222 }
223 num_slices += 1;
224 }
225
226 if (share_tensor_slice_memory) {
227 // (E) Encode tensor data, but by sharing backing store
228 const TensorBuffer* buf = DMAHelper::buffer(&val);
229 buf->Ref();
230 slices[1] = ::grpc::Slice(
231 const_cast<void*>(static_cast<const void*>(tdata.data())),
232 tdata.size(),
233 [](void* backing) { static_cast<TensorBuffer*>(backing)->Unref(); },
234 const_cast<TensorBuffer*>(buf));
235 num_slices += 1;
236 }
237 size_t total_bytes = 0;
238 for (int i = 0; i < num_slices; i++) {
239 total_bytes += slices[i].size();
240 }
241 CHECK_EQ(total_bytes, expected_size);
242
243 ::grpc::ByteBuffer tmp(&slices[0], num_slices);
244 result->Swap(&tmp);
245 }
246 }
247
248 } // namespace grpc
249 } // namespace tensorflow
250