[XLA][Feature] - Pass config flags for LLVM runtime.

[annanay25]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 14.04
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): ('v1.0.0-1783-g4c3bb1a', '1.0.0')
• Bazel version (if compiling from source): 0.4.5
• CUDA/cuDNN version: -
• GPU model and memory: -
• Exact command to reproduce: -

Problem description

As part of my Google Summer of Code project, I am trying to build TensorFlow with Polly-enabled LLVM. To do this, I have written my own BUILD file which runs TensorFlow with a custom repository of LLVM that has Polly checked out as well. I have managed to get a clean build and am now looking to incorporate Polly's passes in the Optimization pipeline of XLA.

In XLA, the llvm module passes are registered here.

Polly register's its passes in LLVM through the following steps

static llvm::RegisterStandardPasses RegisterPollyOptimizerEarly(
    llvm::PassManagerBuilder::EP_ModuleOptimizerEarly,
    registerPollyEarlyAsPossiblePasses);

Corresponding file - <polly-src>/lib/Support/RegisterPasses.cpp.

polly::initializePollyPasses(Registry); 

Corresponding file - <polly-src>/lib/Polly.cpp

I have built the object files for both these files. But I want to check if Polly is actually being invoked in the pipeline, and so my question is -

• Are these steps enough to use Polly in the bazel build of TensorFlow?
• How can I pass configuration flags to LLVM in TensorFlow to check for Polly usage?

As a reference, please find my BUILD file here.

cc @phawkins @eliben

[jart]

Thank you for working to bring LLVM Polly optimizations to TensorFlow! Since @phawkins owns our llvm.BUILD file, I'll assign to him, since helping you accomplish this will probably be quite beneficial to the community.

[annanay25]

Thank you so much! @jart

Update: I have figured out the answer to the second question -
How can I pass configuration flags to LLVM in TensorFlow to check for Polly usage?

Its by adding TF_XLA_FLAGS="--xla_cpu_llvm_cl_opts=-polly" to the python command.

However, the -polly flag is not being recognized by llvm yet, and that means that the pass manager has not yet been linked correctly. I will need some help regarding the same.

Thanks.

[brianretford]

I've been working w/ Annanay to get the llvm build files modified -- he's done some great work. Annanay - can you share your repo here? It'll be a lot easier for people to help you if they can reproduce your work easily.

[annanay25]

Hi @brianretford,
Thanks for this.

These are the steps to reproduce my state right now:

• Clone my tensorflow repository or
  1. copy the llvm build file
  2. edit the xla build file to add this line to your local tensorflow repository.
• Clone this version of llvm - it is rolled back to the required commit and has polly checked out as well.
• Modify the tensorflow/tensorflow/workspace.bzl file and change the llvm rule by replacing it with this:

  native.new_local_repository (
       name = "llvm",
       path = "<path to llvm clone>",
       build_file = str(Label("//third_party/llvm:llvm.BUILD")),
   )

• Run bazel build --verbose_failures -c opt //tensorflow/tools/pip_package:build_pip_package and bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
  in the tensorflow root folder.

• Navigate to a folder with a python file that uses tensorflow.
  Run TF_XLA_FLAGS="--xla_cpu_llvm_cl_opts=-disable-inlined-alloca-merging" python <filename.py> --xla=true. It should run without any errors.
  Run
  TF_XLA_FLAGS="--xla_cpu_llvm_cl_opts=-polly" python <filename.py> --xla=true. The command should fail with "unrecognized flag -polly".

Thanks.

[annanay25]

Please check #10471.