1 /* Copyright 2019 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/common_runtime/eager/tensor_handle.h"
17
18 #include "tensorflow/core/common_runtime/composite_device.h"
19 #include "tensorflow/core/common_runtime/device_mgr.h"
20 #include "tensorflow/core/framework/tensor_shape.h"
21 #include "tensorflow/core/lib/core/status_test_util.h"
22 #include "tensorflow/core/platform/random.h"
23 #include "tensorflow/core/platform/test.h"
24
25 namespace tensorflow {
26 namespace {
27
TEST(TensorHandle_ShapeTest,AsyncShape)28 TEST(TensorHandle_ShapeTest, AsyncShape) {
29 Tensor t(DT_UINT16, TensorShape({2, 2}));
30 EXPECT_TRUE(t.shape().IsSameSize(TensorShape({2, 2})));
31 for (int64_t a = 0; a < t.shape().dim_size(0); a++) {
32 for (int64_t b = 0; b < t.shape().dim_size(1); b++) {
33 t.matrix<uint16>()(a, b) = uint16(a * b);
34 }
35 }
36
37 StaticDeviceMgr device_mgr(DeviceFactory::NewDevice(
38 "CPU", {}, "/job:localhost/replica:0/task:0/device:CPU:0"));
39 auto ctx = new EagerContext(
40 SessionOptions(),
41 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT, false,
42 &device_mgr, false, nullptr, nullptr);
43 TensorHandle* sync_th =
44 TensorHandle::CreateLocalHandle(std::move(t), nullptr, nullptr, ctx);
45 TensorHandle* async_th = TensorHandle::CreateEmptyLocalHandle(
46 nullptr, nullptr, nullptr, DataType::DT_UINT16, ctx);
47
48 EXPECT_TRUE(async_th->CopyInferenceShape(sync_th).ok());
49
50 TensorShape sync_shape;
51 TensorShape async_shape;
52 EXPECT_TRUE(sync_th->Shape(&sync_shape).ok());
53 EXPECT_TRUE(async_th->Shape(&async_shape).ok());
54 EXPECT_EQ(sync_shape, async_shape);
55
56 int num_dims = -1;
57 EXPECT_TRUE(async_th->NumDims(&num_dims).ok());
58 EXPECT_EQ(num_dims, 2);
59
60 int64_t num_elements = -1;
61 EXPECT_TRUE(async_th->NumElements(&num_elements).ok());
62 EXPECT_EQ(num_elements, 4);
63
64 sync_th->Unref();
65 async_th->Unref();
66 ctx->Unref();
67 }
68
CreateDevice(const char * type,const char * name,bool is_local=true)69 static Device* CreateDevice(const char* type, const char* name,
70 bool is_local = true) {
71 class FakeDevice : public Device {
72 public:
73 explicit FakeDevice(const DeviceAttributes& attr, bool is_local)
74 : Device(nullptr, attr), is_local_(is_local) {}
75 Status Sync() override { return OkStatus(); }
76 Allocator* GetAllocator(AllocatorAttributes) override { return nullptr; }
77 bool IsLocal() const override { return is_local_; }
78
79 private:
80 const bool is_local_;
81 };
82 DeviceAttributes attr;
83 attr.set_name(name);
84 attr.set_device_type(type);
85 int64_t incarnation = random::New64();
86 while (incarnation == 0) {
87 incarnation = random::New64();
88 }
89 attr.set_incarnation(incarnation);
90 return new FakeDevice(attr, is_local);
91 }
92
93 } // namespace
94
95 class PackedTensorHandleTest : public ::testing::Test {
96 public:
PackedTensorHandleTest()97 PackedTensorHandleTest() {
98 std::vector<std::unique_ptr<Device>> devices;
99 for (const char* name : device_names_) {
100 devices.emplace_back(CreateDevice("GPU", name));
101 }
102 devices.emplace_back(CreateDevice("CPU", host_name_));
103 device_mgr_ = new StaticDeviceMgr(std::move(devices));
104
105 context_ = new EagerContext(
106 SessionOptions(),
107 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
108 /* async= */ false, device_mgr_,
109 /* device_mgr_owned= */ false, /* rendezvous= */ nullptr,
110 /* cluster_flr= */ nullptr);
111 }
112
~PackedTensorHandleTest()113 ~PackedTensorHandleTest() override {
114 delete device_mgr_;
115 context_->Unref();
116 }
117
context()118 EagerContext* context() { return context_; }
119
ListDevices() const120 std::vector<Device*> ListDevices() const {
121 return device_mgr_->ListDevices();
122 }
123
IsReady(TensorHandle * handle) const124 bool IsReady(TensorHandle* handle) const { return handle->IsReady(); }
125
126 private:
127 const std::vector<const char*> device_names_ = {
128 "/job:worker/replica:0/task:0/device:GPU:0",
129 "/job:worker/replica:0/task:0/device:GPU:1",
130 "/job:worker/replica:0/task:1/device:GPU:0",
131 "/job:worker/replica:0/task:1/device:GPU:1"};
132
133 const char* host_name_ = "/job:worker/replica:0/task:0/device:CPU:0";
134
135 StaticDeviceMgr* device_mgr_;
136 EagerContext* context_;
137 };
138
TEST_F(PackedTensorHandleTest,PackedHandle)139 TEST_F(PackedTensorHandleTest, PackedHandle) {
140 tensorflow::DataType dtype = DT_RESOURCE;
141 TensorShape shape = {};
142 DtypeAndPartialTensorShape dtype_and_shape = {DT_FLOAT, {2, 2}};
143
144 // Create 2 local TensorHandles (ready)
145 std::vector<TensorHandle*> handles;
146 Tensor t0(dtype, shape);
147 Device* d0 = ListDevices().at(0);
148 TensorHandle* h0 =
149 TensorHandle::CreateLocalHandle(std::move(t0), d0, d0, d0, context());
150 h0->SetResourceHandleDtypeAndShape({dtype_and_shape});
151 handles.push_back(h0);
152 Tensor t1(dtype, shape);
153 Device* d1 = ListDevices().at(1);
154 TensorHandle* h1 =
155 TensorHandle::CreateLocalHandle(std::move(t1), d1, d1, d1, context());
156 h1->SetResourceHandleDtypeAndShape({dtype_and_shape});
157 handles.push_back(h1);
158
159 // Create 2 remote TensorHandles (not ready).
160 const string remote_task = "/job:worker/replica:0/task:1";
161 Device* d2 = ListDevices().at(2);
162 TensorHandle* h2 = TensorHandle::CreateUnshapedRemoteHandle(
163 /*op_id=*/0, /*output_num=*/0, remote_task, dtype, d2, context());
164 handles.push_back(h2);
165 Device* d3 = ListDevices().at(3);
166 TensorHandle* h3 = TensorHandle::CreateUnshapedRemoteHandle(
167 /*op_id=*/1, /*output_num=*/0, remote_task, dtype, d3, context());
168 handles.push_back(h3);
169
170 TensorHandle* packed_handle = nullptr;
171 TF_EXPECT_OK(TensorHandle::CreatePackedHandle(std::move(handles), context(),
172 &packed_handle));
173
174 h0->Unref();
175 h1->Unref();
176 h2->Unref();
177 h3->Unref();
178
179 EXPECT_EQ(packed_handle->NumPackedHandles(), 4);
180 EXPECT_EQ(packed_handle->Type(), TensorHandle::PACKED);
181 EXPECT_EQ(packed_handle->dtype, dtype);
182 TensorShape packed_shape;
183 TF_ASSERT_OK(packed_handle->Shape(&packed_shape));
184 EXPECT_EQ(packed_shape, shape);
185 std::vector<DtypeAndPartialTensorShape> dtypes_and_shapes;
186 TF_ASSERT_OK(
187 packed_handle->GetResourceHandleDtypesAndShapes(&dtypes_and_shapes));
188 EXPECT_EQ(dtypes_and_shapes.size(), 1);
189 EXPECT_EQ(dtypes_and_shapes.at(0).dtype, DT_FLOAT);
190 EXPECT_EQ(dtypes_and_shapes.at(0).shape.IsIdenticalTo({2, 2}), true);
191
192 CompositeDevice* device =
193 reinterpret_cast<CompositeDevice*>(packed_handle->device());
194 EXPECT_EQ(device->name(), "/job:worker/replica:0/task:0/device:COMPOSITE:0");
195 EXPECT_EQ(device->underlying_devices()->size(), 4);
196
197 const std::vector<TensorHandle::HandleType> expected_handle_types = {
198 TensorHandle::LOCAL, TensorHandle::LOCAL, TensorHandle::REMOTE,
199 TensorHandle::REMOTE};
200 for (int i = 0; i < packed_handle->NumPackedHandles(); ++i) {
201 TensorHandle* h = nullptr;
202 TF_ASSERT_OK(packed_handle->ExtractPackedHandle(i, &h));
203 EXPECT_EQ(h->device(), ListDevices().at(i));
204 EXPECT_EQ(h->Type(), expected_handle_types.at(i));
205 }
206 EXPECT_FALSE(IsReady(packed_handle));
207
208 TF_ASSERT_OK(h2->SetRemoteShape(shape, ListDevices().at(2),
209 context()->GetContextViewId()));
210 EXPECT_FALSE(IsReady(packed_handle));
211 TF_ASSERT_OK(h3->SetRemoteShape(shape, ListDevices().at(3),
212 context()->GetContextViewId()));
213 EXPECT_TRUE(IsReady(packed_handle));
214
215 packed_handle->Unref();
216 }
217
TEST_F(PackedTensorHandleTest,PackedSingleHandle)218 TEST_F(PackedTensorHandleTest, PackedSingleHandle) {
219 tensorflow::DataType dtype = DT_RESOURCE;
220 TensorShape shape = {};
221
222 Tensor t(dtype, shape);
223 Device* d = ListDevices().at(0);
224 TensorHandle* h =
225 TensorHandle::CreateLocalHandle(std::move(t), d, d, d, context());
226 std::vector<TensorHandle*> handles = {h};
227
228 TensorHandle* packed_handle = nullptr;
229 TF_EXPECT_OK(TensorHandle::CreatePackedHandle(std::move(handles), context(),
230 &packed_handle));
231 h->Unref();
232
233 EXPECT_EQ(packed_handle->Type(), TensorHandle::PACKED);
234 EXPECT_EQ(packed_handle->dtype, dtype);
235 TensorShape packed_shape;
236 TF_ASSERT_OK(packed_handle->Shape(&packed_shape));
237 EXPECT_EQ(packed_shape, shape);
238
239 CompositeDevice* device =
240 reinterpret_cast<CompositeDevice*>(packed_handle->device());
241 EXPECT_EQ(device->name(), "/job:worker/replica:0/task:0/device:COMPOSITE:0");
242 EXPECT_EQ(device->underlying_devices()->size(), 1);
243 EXPECT_EQ(packed_handle->NumPackedHandles(), 1);
244 TensorHandle* h0 = nullptr;
245 TF_ASSERT_OK(packed_handle->ExtractPackedHandle(0, &h0));
246 EXPECT_EQ(h0->device(), d);
247 EXPECT_TRUE(IsReady(packed_handle));
248 packed_handle->Unref();
249 }
250
TEST(TensorHandle_ResourceDeviceTest,OnLocalDevice)251 TEST(TensorHandle_ResourceDeviceTest, OnLocalDevice) {
252 std::unique_ptr<Device> d0(
253 CreateDevice("CPU", "/job:localhost/replica:0/task:0/device:CPU:0"));
254 StaticDeviceMgr local_device_mgr(std::move(d0));
255 auto ctx = new EagerContext(
256 SessionOptions(),
257 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT, false,
258 &local_device_mgr, false, nullptr, nullptr);
259
260 tensorflow::DataType dtype = DT_RESOURCE;
261 TensorShape shape = {2};
262 Tensor t(dtype, shape);
263
264 Device* d = local_device_mgr.ListDevices()[0];
265 TensorHandle* th =
266 TensorHandle::CreateLocalHandle(std::move(t), d, d, d, ctx);
267 // Remote device incarnation for local resource should be 0 (invalid)
268 EXPECT_EQ(0, th->resource_remote_device_incarnation());
269 // Local device manager must contain the resource device.
270 EXPECT_TRUE(local_device_mgr.ContainsDevice(
271 th->resource_device()->attributes().incarnation()));
272
273 std::unique_ptr<Device> d1(
274 CreateDevice("CPU", "/job:localhost/replica:0/task:0/device:CPU:0"));
275 StaticDeviceMgr new_device_mgr(std::move(d1));
276 EXPECT_FALSE(new_device_mgr.ContainsDevice(
277 th->resource_device()->attributes().incarnation()));
278
279 th->Unref();
280 ctx->Unref();
281 }
282
TEST(TensorHandle_ResourceDeviceTest,OnRemoteDevice)283 TEST(TensorHandle_ResourceDeviceTest, OnRemoteDevice) {
284 std::unique_ptr<Device> d_local(
285 CreateDevice("CPU", "/job:localhost/replica:0/task:0/device:CPU:0"));
286 StaticDeviceMgr local_device_mgr(std::move(d_local));
287 auto ctx = new EagerContext(
288 SessionOptions(),
289 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT, false,
290 &local_device_mgr, false, nullptr, nullptr);
291
292 std::unique_ptr<Device> d0(
293 CreateDevice("CPU", "/job:worker/task:0/device:CPU:0", false));
294 Device* d0_ptr = d0.get();
295 std::unique_ptr<Device> d1(
296 CreateDevice("CPU", "/job:worker/task:1/device:CPU:0", false));
297 Device* d1_ptr = d1.get();
298
299 DynamicDeviceMgr remote_device_mgr;
300 std::vector<std::unique_ptr<Device>> vector_d0;
301 vector_d0.emplace_back(std::move(d0));
302 TF_ASSERT_OK(remote_device_mgr.AddDevices(std::move(vector_d0)));
303
304 TensorHandle* th0 = TensorHandle::CreateUnshapedRemoteHandle(
305 0, 0, "", DT_RESOURCE, d0_ptr, ctx);
306 EXPECT_TRUE(remote_device_mgr.ContainsDevice(
307 th0->resource_remote_device_incarnation()));
308
309 std::vector<std::unique_ptr<Device>> vector_d1;
310 vector_d1.emplace_back(std::move(d1));
311 TF_ASSERT_OK(remote_device_mgr.AddDevices(std::move(vector_d1)));
312 EXPECT_TRUE(remote_device_mgr.ContainsDevice(
313 th0->resource_remote_device_incarnation()));
314
315 TensorHandle* th1 = TensorHandle::CreateUnshapedRemoteHandle(
316 0, 0, "", DT_RESOURCE, d1_ptr, ctx);
317 EXPECT_TRUE(remote_device_mgr.ContainsDevice(
318 th1->resource_remote_device_incarnation()));
319
320 std::vector<Device*> remove_d1{d1_ptr};
321 TF_ASSERT_OK(remote_device_mgr.RemoveDevices(std::move(remove_d1)));
322 EXPECT_FALSE(remote_device_mgr.ContainsDevice(
323 th1->resource_remote_device_incarnation()));
324 EXPECT_TRUE(remote_device_mgr.ContainsDevice(
325 th0->resource_remote_device_incarnation()));
326
327 th0->Unref();
328 th1->Unref();
329 ctx->Unref();
330 }
331
332 class RemoteTensorHandleTest : public ::testing::Test {
333 public:
RemoteTensorHandleTest()334 RemoteTensorHandleTest() {
335 std::vector<std::unique_ptr<Device>> devices;
336 for (const char* name : device_names_) {
337 devices.emplace_back(CreateDevice("CPU", name));
338 }
339 device_mgr_ = new StaticDeviceMgr(std::move(devices));
340
341 context_ = new EagerContext(
342 SessionOptions(),
343 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
344 /* async= */ false, device_mgr_,
345 /* device_mgr_owned= */ false, /* rendezvous= */ nullptr,
346 /* cluster_flr= */ nullptr);
347 }
348
~RemoteTensorHandleTest()349 ~RemoteTensorHandleTest() override {
350 delete device_mgr_;
351 context_->Unref();
352 }
353
context()354 EagerContext* context() { return context_; }
355
ListDevices() const356 std::vector<Device*> ListDevices() const {
357 return device_mgr_->ListDevices();
358 }
359
360 private:
361 const std::vector<const char*> device_names_ = {
362 "/job:worker/replica:0/task:0/device:CPU:0",
363 "/job:worker/replica:0/task:1/device:CPU:0",
364 "/job:worker/replica:0/task:2/device:CPU:0"};
365
366 StaticDeviceMgr* device_mgr_;
367 EagerContext* context_;
368 };
369
TEST_F(RemoteTensorHandleTest,UnknownRemoteDevice)370 TEST_F(RemoteTensorHandleTest, UnknownRemoteDevice) {
371 std::vector<std::unique_ptr<Device>> devices;
372 devices.emplace_back(
373 CreateDevice("CPU", "/job:worker/replica:0/task:0/device:CPU:0"));
374 devices.emplace_back(
375 CreateDevice("CPU", "/job:worker/replica:0/task:1/device:CPU:0"));
376 devices.emplace_back(
377 CreateDevice("CPU", "/job:worker/replica:0/task:2/device:CPU:0"));
378 StaticDeviceMgr device_mgr(std::move(devices));
379
380 EagerContext* context = new EagerContext(
381 SessionOptions(),
382 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
383 /* async= */ false, &device_mgr,
384 /* device_mgr_owned= */ false, /* rendezvous= */ nullptr,
385 /* cluster_flr= */ nullptr);
386
387 tensorflow::DataType dtype = DT_FLOAT;
388 TensorShape shape = {};
389
390 const string remote_task = "/job:worker/replica:0/task:1";
391 Device* d1 = device_mgr.ListDevices().at(1);
392 TensorHandle* h = TensorHandle::CreateUnshapedRemoteHandle(
393 /*op_id=*/0, /*output_num=*/0, remote_task, dtype, d1, context,
394 /*unknown_device=*/true);
395 EXPECT_EQ(h->device(), d1);
396
397 Device* d2 = device_mgr.ListDevices().at(2);
398 TF_ASSERT_OK(h->SetRemoteShapeAndDevice(
399 shape, d1, context->GetContextViewId(), d2->name()));
400 Status s;
401 EXPECT_EQ(h->BackingDeviceName(&s), d2->name());
402 TF_EXPECT_OK(s);
403 EXPECT_EQ(h->device(), d2);
404 h->Unref();
405 context->Unref();
406 }
407
TEST(TensorHandle_DeviceNameTest,OnLocalDevice)408 TEST(TensorHandle_DeviceNameTest, OnLocalDevice) {
409 std::vector<std::unique_ptr<Device>> devices;
410 devices.emplace_back(
411 CreateDevice("CPU", "/job:localhost/replica:0/task:0/device:CPU:0"));
412 devices.emplace_back(
413 CreateDevice("GPU", "/job:localhost/replica:0/task:0/device:GPU:0"));
414 StaticDeviceMgr local_device_mgr(std::move(devices));
415 auto ctx = new EagerContext(
416 SessionOptions(),
417 tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT, false,
418 &local_device_mgr, false, nullptr, nullptr);
419
420 Device* dcpu = local_device_mgr.ListDevices()[0];
421 Device* dgpu = local_device_mgr.ListDevices()[1];
422 tensorflow::DataType dtype = DT_RESOURCE;
423 TensorShape shape = {2};
424 Tensor tcpu(dtype, shape);
425 Tensor tgpu(dtype, shape);
426 Status s;
427
428 TensorHandle* th_cpu =
429 TensorHandle::CreateLocalHandle(std::move(tcpu), dcpu, dcpu, dcpu, ctx);
430 const char* device_name = th_cpu->DeviceName(&s);
431 TF_EXPECT_OK(s);
432 ASSERT_TRUE(absl::StrContains(device_name, "CPU")) << device_name;
433 const char* backing_device_name = th_cpu->BackingDeviceName(&s);
434 TF_EXPECT_OK(s);
435 ASSERT_TRUE(absl::StrContains(backing_device_name, "CPU"))
436 << backing_device_name;
437 const char* device_type = th_cpu->DeviceType(&s);
438 TF_EXPECT_OK(s);
439 ASSERT_TRUE(absl::StrContains(device_type, "CPU")) << device_type;
440 int device_id = th_cpu->DeviceId(&s);
441 TF_EXPECT_OK(s);
442 ASSERT_EQ(0, device_id) << device_id;
443
444 TensorHandle* th_gpu =
445 TensorHandle::CreateLocalHandle(std::move(tgpu), dgpu, dgpu, dgpu, ctx);
446 device_name = th_gpu->DeviceName(&s);
447 TF_EXPECT_OK(s);
448 ASSERT_TRUE(absl::StrContains(device_name, "GPU")) << device_name;
449 backing_device_name = th_gpu->BackingDeviceName(&s);
450 TF_EXPECT_OK(s);
451 std::cout << "backing_device_name for GPU: " << backing_device_name
452 << std::endl;
453 ASSERT_TRUE(absl::StrContains(backing_device_name, "GPU"))
454 << backing_device_name;
455 device_type = th_gpu->DeviceType(&s);
456 TF_EXPECT_OK(s);
457 ASSERT_TRUE(absl::StrContains(device_type, "GPU")) << device_type;
458 device_id = th_gpu->DeviceId(&s);
459 TF_EXPECT_OK(s);
460 ASSERT_EQ(0, device_id) << device_id;
461
462 th_cpu->Unref();
463 th_gpu->Unref();
464 ctx->Unref();
465 }
466
467 } // namespace tensorflow
468