1 /* Copyright 2018 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/c/eager/c_api_test_util.h"
17
18 #include "tensorflow/c/eager/c_api.h"
19 #include "tensorflow/c/eager/c_api_experimental.h"
20 #include "tensorflow/c/tf_datatype.h"
21 #include "tensorflow/c/tf_tensor.h"
22 #include "tensorflow/core/platform/logging.h"
23 #include "tensorflow/core/platform/strcat.h"
24 #include "tensorflow/core/platform/test.h"
25 #include "tensorflow/core/platform/tstring.h"
26 #include "tensorflow/core/protobuf/cluster.pb.h"
27
28 using tensorflow::string;
29 using tensorflow::tstring;
30
TestScalarTensorHandle(TFE_Context * ctx,float value)31 TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, float value) {
32 float data[] = {value};
33 TF_Status* status = TF_NewStatus();
34 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, nullptr, 0, status);
35 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
36 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
37 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
38 TF_DeleteTensor(t);
39 TF_DeleteStatus(status);
40 return th;
41 }
42
TestScalarTensorHandle(TFE_Context * ctx,const tensorflow::tstring & value)43 TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx,
44 const tensorflow::tstring& value) {
45 TF_Status* status = TF_NewStatus();
46 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_STRING, nullptr, 0, status);
47 tstring* data = static_cast<tstring*>(TF_TensorData(t));
48 *data = value;
49 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
50 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
51 TF_DeleteTensor(t);
52 TF_DeleteStatus(status);
53 return th;
54 }
55
TestScalarTensorHandle(TFE_Context * ctx,int value)56 TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, int value) {
57 int data[] = {value};
58 TF_Status* status = TF_NewStatus();
59 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_INT32, nullptr, 0, status);
60 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
61 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
62 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
63 TF_DeleteTensor(t);
64 TF_DeleteStatus(status);
65 return th;
66 }
67
TestScalarTensorHandle(TFE_Context * ctx,bool value)68 TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, bool value) {
69 bool data[] = {value};
70 TF_Status* status = TF_NewStatus();
71 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_BOOL, nullptr, 0, status);
72 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
73 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
74 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
75 TF_DeleteTensor(t);
76 TF_DeleteStatus(status);
77 return th;
78 }
79
DoubleTestMatrixTensorHandle(TFE_Context * ctx)80 TFE_TensorHandle* DoubleTestMatrixTensorHandle(TFE_Context* ctx) {
81 int64_t dims[] = {2, 2};
82 double data[] = {1.0, 2.0, 3.0, 4.0};
83 TF_Status* status = TF_NewStatus();
84 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_DOUBLE, &dims[0],
85 sizeof(dims) / sizeof(int64_t), status);
86 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
87 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
88 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
89 TF_DeleteTensor(t);
90 TF_DeleteStatus(status);
91 return th;
92 }
93
TestMatrixTensorHandle(TFE_Context * ctx)94 TFE_TensorHandle* TestMatrixTensorHandle(TFE_Context* ctx) {
95 int64_t dims[] = {2, 2};
96 float data[] = {1.0f, 2.0f, 3.0f, 4.0f};
97 TF_Status* status = TF_NewStatus();
98 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
99 sizeof(dims) / sizeof(int64_t), status);
100 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
101 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
102 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
103 TF_DeleteTensor(t);
104 TF_DeleteStatus(status);
105 return th;
106 }
107
TestMatrixTensorHandleWithInput(TFE_Context * ctx,float data[],int64_t dims[],int num_dims)108 TFE_TensorHandle* TestMatrixTensorHandleWithInput(TFE_Context* ctx,
109 float data[], int64_t dims[],
110 int num_dims) {
111 TF_Status* status = TF_NewStatus();
112 TF_Tensor* t =
113 TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0], num_dims, status);
114 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
115 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
116 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
117 TF_DeleteTensor(t);
118 TF_DeleteStatus(status);
119 return th;
120 }
121
TestTensorHandleWithDimsFloat(TFE_Context * ctx,float data[],int64_t dims[],int num_dims)122 TFE_TensorHandle* TestTensorHandleWithDimsFloat(TFE_Context* ctx, float data[],
123 int64_t dims[], int num_dims) {
124 TF_Status* status = TF_NewStatus();
125 TF_Tensor* t =
126 TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0], num_dims, status);
127 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
128 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
129 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
130 TF_DeleteTensor(t);
131 TF_DeleteStatus(status);
132 return th;
133 }
134
TestTensorHandleWithDimsInt(TFE_Context * ctx,int data[],int64_t dims[],int num_dims)135 TFE_TensorHandle* TestTensorHandleWithDimsInt(TFE_Context* ctx, int data[],
136 int64_t dims[], int num_dims) {
137 TF_Status* status = TF_NewStatus();
138 TF_Tensor* t =
139 TFE_AllocateHostTensor(ctx, TF_INT32, &dims[0], num_dims, status);
140 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
141 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
142 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
143 TF_DeleteTensor(t);
144 TF_DeleteStatus(status);
145 return th;
146 }
147
TestMatrixTensorHandle100x100(TFE_Context * ctx)148 TFE_TensorHandle* TestMatrixTensorHandle100x100(TFE_Context* ctx) {
149 constexpr int64_t dims[] = {100, 100};
150 constexpr int num_elements = dims[0] * dims[1];
151 float data[num_elements];
152 for (int i = 0; i < num_elements; ++i) {
153 data[i] = 1.0f;
154 }
155 TF_Status* status = TF_NewStatus();
156 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
157 sizeof(dims) / sizeof(int64_t), status);
158 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
159 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
160 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
161 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
162 TF_DeleteTensor(t);
163 TF_DeleteStatus(status);
164 return th;
165 }
166
DoubleTestMatrixTensorHandle3X2(TFE_Context * ctx)167 TFE_TensorHandle* DoubleTestMatrixTensorHandle3X2(TFE_Context* ctx) {
168 int64_t dims[] = {3, 2};
169 double data[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
170 TF_Status* status = TF_NewStatus();
171 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
172 sizeof(dims) / sizeof(int64_t), status);
173 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
174 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
175 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
176 TF_DeleteTensor(t);
177 TF_DeleteStatus(status);
178 return th;
179 }
180
TestMatrixTensorHandle3X2(TFE_Context * ctx)181 TFE_TensorHandle* TestMatrixTensorHandle3X2(TFE_Context* ctx) {
182 int64_t dims[] = {3, 2};
183 float data[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
184 TF_Status* status = TF_NewStatus();
185 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
186 sizeof(dims) / sizeof(int64_t), status);
187 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
188 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
189 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
190 TF_DeleteTensor(t);
191 TF_DeleteStatus(status);
192 return th;
193 }
194
TestVariable(TFE_Context * ctx,float value,const tensorflow::string & device_name)195 TFE_TensorHandle* TestVariable(TFE_Context* ctx, float value,
196 const tensorflow::string& device_name) {
197 TF_Status* status = TF_NewStatus();
198 // Create the variable handle.
199 TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
200 if (TF_GetCode(status) != TF_OK) return nullptr;
201 TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
202 TFE_OpSetAttrShape(op, "shape", {}, 0, status);
203 TFE_OpSetAttrString(op, "container", "localhost", 0);
204 TFE_OpSetAttrString(op, "shared_name", "", 0);
205 if (!device_name.empty()) {
206 TFE_OpSetDevice(op, device_name.c_str(), status);
207 }
208 if (TF_GetCode(status) != TF_OK) return nullptr;
209 TFE_TensorHandle* var_handle = nullptr;
210 int num_retvals = 1;
211 TFE_Execute(op, &var_handle, &num_retvals, status);
212 if (TF_GetCode(status) != TF_OK) return nullptr;
213 TFE_DeleteOp(op);
214 if (TF_GetCode(status) != TF_OK) return nullptr;
215 CHECK_EQ(1, num_retvals);
216
217 // Assign 'value' to it.
218 op = TFE_NewOp(ctx, "AssignVariableOp", status);
219 if (TF_GetCode(status) != TF_OK) return nullptr;
220 TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
221 TFE_OpAddInput(op, var_handle, status);
222
223 // Convert 'value' to a TF_Tensor then a TFE_TensorHandle.
224 std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> t(
225 TF_AllocateTensor(TF_FLOAT, nullptr, 0, sizeof(value)), TF_DeleteTensor);
226 memcpy(TF_TensorData(t.get()), &value, TF_TensorByteSize(t.get()));
227
228 std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)>
229 value_handle(TFE_NewTensorHandle(t.get(), status),
230 TFE_DeleteTensorHandle);
231 if (TF_GetCode(status) != TF_OK) return nullptr;
232
233 TFE_OpAddInput(op, value_handle.get(), status);
234 if (TF_GetCode(status) != TF_OK) return nullptr;
235
236 num_retvals = 0;
237 TFE_Execute(op, nullptr, &num_retvals, status);
238 TFE_DeleteOp(op);
239 if (TF_GetCode(status) != TF_OK) return nullptr;
240 CHECK_EQ(0, num_retvals);
241
242 TF_DeleteStatus(status);
243
244 return var_handle;
245 }
246
AddOp(TFE_Context * ctx,TFE_TensorHandle * a,TFE_TensorHandle * b)247 TFE_Op* AddOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
248 TF_Status* status = TF_NewStatus();
249
250 TFE_Op* op = TFE_NewOp(ctx, "AddV2", status);
251 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
252 TFE_OpAddInput(op, a, status);
253 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
254 TFE_OpAddInput(op, b, status);
255 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
256 TF_DeleteStatus(status);
257 TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
258
259 return op;
260 }
261
MatMulOp(TFE_Context * ctx,TFE_TensorHandle * a,TFE_TensorHandle * b)262 TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
263 TF_Status* status = TF_NewStatus();
264
265 TFE_Op* op = TFE_NewOp(ctx, "MatMul", status);
266 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
267 TFE_OpAddInput(op, a, status);
268 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
269 TFE_OpAddInput(op, b, status);
270 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
271 TF_DeleteStatus(status);
272 TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
273
274 return op;
275 }
276
IdentityOp(TFE_Context * ctx,TFE_TensorHandle * a)277 TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a) {
278 TF_Status* status = TF_NewStatus();
279
280 TFE_Op* op = TFE_NewOp(ctx, "Identity", status);
281 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
282 TFE_OpAddInput(op, a, status);
283 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
284 TF_DeleteStatus(status);
285 TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
286
287 return op;
288 }
289
ShapeOp(TFE_Context * ctx,TFE_TensorHandle * a)290 TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a) {
291 TF_Status* status = TF_NewStatus();
292
293 TFE_Op* op = TFE_NewOp(ctx, "Shape", status);
294 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
295 TFE_OpAddInput(op, a, status);
296 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
297 TF_DeleteStatus(status);
298 TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
299
300 return op;
301 }
302
TestAxisTensorHandle(TFE_Context * ctx)303 TFE_TensorHandle* TestAxisTensorHandle(TFE_Context* ctx) {
304 int64_t dims[] = {1};
305 int data[] = {1};
306 TF_Status* status = TF_NewStatus();
307 TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_INT32, &dims[0],
308 sizeof(dims) / sizeof(int64_t), status);
309 memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
310 TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
311 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
312 TF_DeleteTensor(t);
313 TF_DeleteStatus(status);
314 return th;
315 }
316
MinOp(TFE_Context * ctx,TFE_TensorHandle * input,TFE_TensorHandle * axis)317 TFE_Op* MinOp(TFE_Context* ctx, TFE_TensorHandle* input,
318 TFE_TensorHandle* axis) {
319 TF_Status* status = TF_NewStatus();
320
321 TFE_Op* op = TFE_NewOp(ctx, "Min", status);
322 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
323 TFE_OpAddInput(op, input, status);
324 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
325 TFE_OpAddInput(op, axis, status);
326 CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
327 TFE_OpSetAttrBool(op, "keep_dims", 1);
328 TFE_OpSetAttrType(op, "Tidx", TF_INT32);
329 TF_DeleteStatus(status);
330 TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(input));
331
332 return op;
333 }
334
GetDeviceName(TFE_Context * ctx,string * device_name,const char * device_type)335 bool GetDeviceName(TFE_Context* ctx, string* device_name,
336 const char* device_type) {
337 std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
338 TF_NewStatus(), TF_DeleteStatus);
339 TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get());
340 CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
341
342 const int num_devices = TF_DeviceListCount(devices);
343 for (int i = 0; i < num_devices; ++i) {
344 const string dev_type(TF_DeviceListType(devices, i, status.get()));
345 CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
346 const string dev_name(TF_DeviceListName(devices, i, status.get()));
347 CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
348 if (dev_type == device_type) {
349 *device_name = dev_name;
350 LOG(INFO) << "Found " << device_type << " device " << *device_name;
351 TF_DeleteDeviceList(devices);
352 return true;
353 }
354 }
355 TF_DeleteDeviceList(devices);
356 return false;
357 }
358
GetServerDef(const string & job_name,int num_tasks)359 tensorflow::ServerDef GetServerDef(const string& job_name, int num_tasks) {
360 tensorflow::ServerDef server_def;
361 server_def.set_protocol("grpc");
362 server_def.set_job_name(job_name);
363 server_def.set_task_index(0);
364 tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
365 tensorflow::JobDef* job_def = cluster_def->add_job();
366 job_def->set_name(job_name);
367 for (int i = 0; i < num_tasks; i++) {
368 int port = tensorflow::testing::PickUnusedPortOrDie();
369 job_def->mutable_tasks()->insert(
370 {i, tensorflow::strings::StrCat("localhost:", port)});
371 }
372 return server_def;
373 }
374
GetServerDef(int num_tasks)375 tensorflow::ServerDef GetServerDef(int num_tasks) {
376 return GetServerDef("localhost", num_tasks);
377 }
378