Tensorflow XLA Ahead-of-Time (AOT) compilation with Saved Model (TF2)

makefile

INC = -I$HOME/.local/lib/python3.8/site-packages/tensorflow/include
LIB = -L$HOME/.local/lib/python3.8/site-packages/tensorflow/xla_aot_runtime_src
CXXFLAGS = -D_GLIBCXX_USE_CXX11_ABI=0
LIBS = -lpthread -ltf_xla_runtime

.phony: all clean

all: resnet_test

resnet_test: resnet_test.cc foo.o
    g++ ${CXXFLAGS} ${INC} ${LIB} $^ -o $@ ${LIBS}

clean:
    \rm resnet_test

Saved_model_aot.md

Frozen graph is deprecated in TF2. Using Saved model is the new approach. Instead of creating a configuration file and using bazel in steps 1 and 2 described here, simply use the saved_model_cli. The final compilation step of creating a binary can be performed using a makefile instead of bazel.

1. Create a TF saved_model using tf.saved_model.save.
2. Use the saved_model_cli to perform XLA AOT compilation.

saved_model_cli aot_compile_cpu \
    --dir resnet/1 \
    --tag_set serve \
    --signature_def_key serving_default \
    --output_prefix foo \
    --cpp_class Bar

This will create files foo.h, foo.o, foo_metadata.o and foo_makefile.inc. The file foo.h has a class named Bar.

3. Create libtf_xla_runtime.a. If you installed tensorflow-2.* as user, then

cd $HOME/.local/lib/python3.x/site-packages/tensorflow/xla_aot_runtime_src
cmake .
make

4. Follow the example here and create a resnet_test.cc.

5. Use the makefile to compile resnet_test.cc, the generated object file foo.o and the library libtf_xla_runtime.a to create the final executable.