Method MakeContext

tensorflow/compiler/xrt/client/xrt_client_test.cc:148–305

}

xla::StatusOr<std::shared_ptr<XrtContext>> XrtClientTest::MakeContext() {
  ChannelCreationFunction channel_func =
      ConvertToChannelCreationFunction(NewHostPortGrpcChannel);
  TF_ASSIGN_OR_RETURN(std::shared_ptr<GrpcChannelCache> channel_cache,
                      GetGrpcChannelCache(cluster_def_, channel_func));

  auto client = std::make_shared<XrtTfClient>(cluster_def_, channel_cache);
  TF_ASSIGN_OR_RETURN(
      std::shared_ptr<XrtTfContext> tf_context,
      XrtTfContext::Create(XrtTfContext::Options(), client, /*job=*/"localhost",
                           /*task=*/0));

  TF_ASSIGN_OR_RETURN(auto context, XrtContext::Create(tf_context, "XLA_CPU"));

  // There should be exactly one XLA_CPU device.
  TF_RET_CHECK(context->device_count() == 1);
  return context;
}

// Tests that we can use the XRT client to perform some simple operations.
TEST_F(XrtClientTest, XrtClientWorks) {
  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<XrtContext> context, MakeContext());

  ASSERT_TRUE(context->tf_context() != nullptr);

  EXPECT_EQ(context->tf_device_ids().size(), 1);

  ASSERT_EQ(context->device_mesh_coordinates().size(), 1);
  ASSERT_EQ(context->device_mesh_coordinates()[0].value_size(), 1);
  EXPECT_EQ(context->device_mesh_coordinates()[0].value(0), 0);

  // Tests sending a literal to and from the device.
  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {3, 4, 5});
  TF_ASSERT_OK_AND_ASSIGN(xla::Literal a,
                          xla::LiteralUtil::CreateRandomLiteral<xla::F32>(
                              shape,
                              /*mean=*/7.0, /*stddev=*/13.5));
  TF_ASSERT_OK_AND_ASSIGN(auto buffer, XrtBuffer::FromLiteral(context, 0, a));
  TF_ASSERT_OK_AND_ASSIGN(xla::Literal b, buffer->ToLiteral());
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(a, b));

  // Run a simple computation, fetch its output, and check it is what we expect.
  auto build_computation = [&]() {
    xla::XlaBuilder builder("test_computation");
    xla::XlaOp p = xla::Parameter(&builder, 0, shape, "param");
    xla::Add(p, p);
    return builder.Build();
  };
  TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation computation, build_computation());

  TF_ASSERT_OK_AND_ASSIGN(xla::DeviceAssignment assignment,
                          xla::ComputationPlacer().AssignDevices(1, 1));
  TF_ASSERT_OK_AND_ASSIGN(auto executable,
                          XrtExecutable::Compile(context, computation.proto(),
                                                 {shape}, shape, assignment));
  EXPECT_EQ(executable->device_assignment(), assignment);
  TF_ASSERT_OK_AND_ASSIGN(auto c_buffer, executable->Execute({buffer}));