| 693 | } |
| 694 | |
| 695 | void TensorHandleCopyBetweenTwoGPUDevices(bool async) { |
| 696 | std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status( |
| 697 | TF_NewStatus(), TF_DeleteStatus); |
| 698 | TFE_ContextOptions* opts = TFE_NewContextOptions(); |
| 699 | TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async)); |
| 700 | TFE_Context* ctx = TFE_NewContext(opts, status.get()); |
| 701 | TFE_DeleteContextOptions(opts); |
| 702 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 703 | |
| 704 | TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); |
| 705 | TF_Tensor* t = TFE_TensorHandleResolve(hcpu, status.get()); |
| 706 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 707 | |
| 708 | TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); |
| 709 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 710 | const int num_devices = TF_DeviceListCount(devices); |
| 711 | bool has_gpu0 = false; |
| 712 | bool has_gpu1 = false; |
| 713 | for (int i = 0; i < num_devices; ++i) { |
| 714 | const char* dev = TF_DeviceListName(devices, i, status.get()); |
| 715 | ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); |
| 716 | string device_name(dev); |
| 717 | if (device_name.find("GPU:0") != string::npos) { |
| 718 | has_gpu0 = true; |
| 719 | } |
| 720 | if (device_name.find("GPU:1") != string::npos) { |
| 721 | has_gpu1 = true; |
| 722 | } |
| 723 | } |
| 724 | |
| 725 | const char* kCPUDevice = "CPU:0"; |
| 726 | if (!has_gpu0 || !has_gpu1) { |
| 727 | TF_DeleteDeviceList(devices); |
| 728 | TF_DeleteTensor(t); |
| 729 | TFE_DeleteTensorHandle(hcpu); |
| 730 | TFE_DeleteContext(ctx); |
| 731 | return; |
| 732 | } |
| 733 | const string gpu_1_name(TF_DeviceListName(devices, 1, status.get())); |
| 734 | ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK); |
| 735 | const string gpu_2_name(TF_DeviceListName(devices, 2, status.get())); |
| 736 | ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK); |
| 737 | TFE_TensorHandle* hdevice = |
| 738 | TFE_TensorHandleCopyToDevice(hcpu, ctx, gpu_1_name.c_str(), status.get()); |
| 739 | ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK); |
| 740 | |
| 741 | TFE_TensorHandle* hdevice2 = TFE_TensorHandleCopyToDevice( |
| 742 | hdevice, ctx, gpu_2_name.c_str(), status.get()); |
| 743 | ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK); |
| 744 | TFE_DeleteTensorHandle(hdevice); |
| 745 | // Copy back to CPU |
| 746 | TFE_TensorHandle* hcopy = |
| 747 | TFE_TensorHandleCopyToDevice(hdevice2, ctx, kCPUDevice, status.get()); |
| 748 | ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK); |
| 749 | TFE_DeleteTensorHandle(hdevice2); |
| 750 | |
| 751 | // Ensure that the contents are the same! |
| 752 | TF_Tensor* tcopy = TFE_TensorHandleResolve(hcopy, status.get()); |
no test coverage detected