c++tensorflow-xla

How do I run a minimal XLA C++ computation?


I'm using the XLA C++ API, and I've managed to run a simple addition, but I've no idea if I'm doing it right. There seem to be an awful lot of classes that I've not used. Here's my example

auto builder = new XlaBuilder("XlaBuilder");
auto one = ConstantR0(builder, 1);
auto two = ConstantR0(builder, 2);
auto res = one + two;

ValueInferenceMode value_inf_mode;
auto value_inf = new ValueInference(builder_);

auto lit = value_inf
    ->AnalyzeConstant(res, value_inf_mode)
    ->GetValue()
    ->Clone();

// I'm using `untyped_data` because I can't express arbitrary array types.
// I guess I could use `data<int32>` in this simple case
auto data = lit.untyped_data();

std::cout << ((int32*) data)[0] << std::endl;  // prints 3

Solution

  • I suspect I didn't actually run that computation through XLA. Here's a different approach based on a sample harness in the XLA source code

    XlaComputation computation = res.builder()->Build().ConsumeValueOrDie();
    ExecutionProfile profile;
    Literal lit = ClientLibrary::LocalClientOrDie()
        ->ExecuteAndTransfer(computation, {}, nullptr, &profile)
        .ConsumeValueOrDie();
    
    data = lit.untyped_data()