| 38 | class AsyncOutOpTest : public OpsTestBase {}; |
| 39 | |
| 40 | TEST_F(AsyncOutOpTest, Recv) { |
| 41 | SetDevice(DEVICE_GPU, |
| 42 | std::unique_ptr<tensorflow::Device>( |
| 43 | DeviceFactory::NewDevice("GPU", {}, kTestDeviceName))); |
| 44 | |
| 45 | Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 1, 6, 2})); |
| 46 | test::FillValues<float>( |
| 47 | &expected, {-1.0f, -2.0f, -3.0f, -5.0f, -8.0f, -13.0f, -21.0f, -34.0f, |
| 48 | -55.0f, -89.0f, -143.0f, -231.0f}); |
| 49 | string key = |
| 50 | AsyncIoRendezvous::GetRendezvousKey(kTestDeviceName, kTestTensorName); |
| 51 | uint64 key_hash = AsyncIoRendezvous::GetRendezvousKeyHash(key); |
| 52 | GetXlaAsyncIORendezvous()->InitializeRendezvousQueue(key_hash); |
| 53 | AsyncIoRendezvous::TensorPayload val; |
| 54 | val.tensor = expected; |
| 55 | TF_ASSERT_OK(GetXlaAsyncIORendezvous()->Send(key_hash, val)); |
| 56 | |
| 57 | TF_EXPECT_OK(NodeDefBuilder("async_out_recv", "_XlaAsyncOutRecv") |
| 58 | .Device(kTestDeviceName) |
| 59 | .Attr("T", DT_FLOAT) |
| 60 | .Attr("device_name", kTestDeviceName) |
| 61 | .Attr("tensor_name", kTestTensorName) |
| 62 | .Finalize(node_def())); |
| 63 | TF_EXPECT_OK(InitOpWithGraphVersion(8)); |
| 64 | TF_ASSERT_OK(RunOpKernel()); |
| 65 | GetXlaAsyncIORendezvous()->FinalizeRendezvousQueue(key_hash); |
| 66 | |
| 67 | test::ExpectTensorEqual<float>(expected, *GetOutput(0)); |
| 68 | } |
| 69 | |
| 70 | TEST_F(AsyncOutOpTest, Send) { |
| 71 | std::vector<float> golden{-1.0f, -2.0f, -3.0f, -5.0f, -8.0f, -13.0f}; |
nothing calls this directly
no test coverage detected