Skip to content

Instantly share code, notes, and snippets.

@adityaiitb
Last active April 6, 2022 14:15
Show Gist options
  • Save adityaiitb/73272a239a1a81a2e71c4baaba4dd87a to your computer and use it in GitHub Desktop.
Save adityaiitb/73272a239a1a81a2e71c4baaba4dd87a to your computer and use it in GitHub Desktop.
Tensorflow XLA Ahead-of-Time (AOT) compilation with Saved Model (TF2)
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

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.

  1. 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
  1. Follow the example here and create a resnet_test.cc.

  2. 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment