| 576 | } |
| 577 | |
| 578 | void TensorHandleCopyBetweenDevices(bool async) { |
| 579 | std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status( |
| 580 | TF_NewStatus(), TF_DeleteStatus); |
| 581 | TFE_ContextOptions* opts = TFE_NewContextOptions(); |
| 582 | TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async)); |
| 583 | TFE_Context* ctx = TFE_NewContext(opts, status.get()); |
| 584 | TFE_DeleteContextOptions(opts); |
| 585 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 586 | |
| 587 | TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); |
| 588 | TF_Tensor* t = TFE_TensorHandleResolve(hcpu, status.get()); |
| 589 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 590 | |
| 591 | TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); |
| 592 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 593 | const int num_devices = TF_DeviceListCount(devices); |
| 594 | |
| 595 | const char* kCPUDevice = "CPU:0"; |
| 596 | for (int i = 0; i < num_devices; ++i) { |
| 597 | const string name(TF_DeviceListName(devices, i, status.get())); |
| 598 | if (TF_GetCode(status.get()) != TF_OK) { |
| 599 | ADD_FAILURE() << i << " -- " << TF_Message(status.get()); |
| 600 | continue; |
| 601 | } |
| 602 | auto tag = tensorflow::strings::StrCat("Device #", i, " (", name, ")"); |
| 603 | // Copy to device |
| 604 | TFE_TensorHandle* hdevice = |
| 605 | TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), status.get()); |
| 606 | if (TF_GetCode(status.get()) != TF_OK) { |
| 607 | ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); |
| 608 | continue; |
| 609 | } |
| 610 | // Copy from device to the same device. |
| 611 | TFE_TensorHandle* hdevice2 = |
| 612 | TFE_TensorHandleCopyToDevice(hdevice, ctx, name.c_str(), status.get()); |
| 613 | if (TF_GetCode(status.get()) != TF_OK) { |
| 614 | ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); |
| 615 | continue; |
| 616 | } |
| 617 | TFE_DeleteTensorHandle(hdevice); |
| 618 | // Copy back to CPU |
| 619 | TFE_TensorHandle* hcopy = |
| 620 | TFE_TensorHandleCopyToDevice(hdevice2, ctx, kCPUDevice, status.get()); |
| 621 | if (TF_GetCode(status.get()) != TF_OK) { |
| 622 | ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); |
| 623 | continue; |
| 624 | } |
| 625 | TFE_DeleteTensorHandle(hdevice2); |
| 626 | |
| 627 | // Ensure that the contents are the same! |
| 628 | TF_Tensor* tcopy = TFE_TensorHandleResolve(hcopy, status.get()); |
| 629 | TFE_DeleteTensorHandle(hcopy); |
| 630 | if (TF_GetCode(status.get()) != TF_OK) { |
| 631 | ADD_FAILURE() << tag; |
| 632 | continue; |
| 633 | } |
| 634 | EXPECT_EQ(TF_TensorByteSize(t), TF_TensorByteSize(tcopy)) << tag; |
| 635 | EXPECT_EQ( |
no test coverage detected