[XLA][WIP] Add support for Polyhedral compilation through Polly

Patch by @annanay25 tensorflow#12408

Author: ARPITH C. JACOB acjacob@us.ibm.com

tensorflow/compiler/xla/service/cpu/BUILD

@@ -19,6 +19,15 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
 load("//tensorflow/compiler/xla:xla.bzl", "ORC_JIT_MEMORY_MAPPER_TARGETS")
 
+config_setting(
+ name = "using_polly",
+ values = {
+ "polly": "true",
+ },
+ visibility = ["//tensorflow:internal"]
+)
+ #visibility = ["@%ws%//tensorflow:internal"]
+
 # Filegroup used to collect source files for dependency checking.
 filegroup(
  name = "c_srcs",
@@ -368,6 +377,7 @@ cc_library(
  "@llvm//:object",
  "@llvm//:support",
  "@llvm//:target",
+ "@llvm//:polly",
  ],
 )
 

tensorflow/compiler/xla/service/cpu/build_defs.bzl

@@ -10,3 +10,14 @@ def runtime_copts():
  "//tensorflow:android": ["-O2"],
  "//conditions:default": []
  }))
+
+def if_polly(basic_deps):
+ """Shorthand for select()'ing on whether we're building with Polly.
+ Returns a select statement which evaluates to basic_deps if we're building
+ with Polly enabled. Otherwise, the select statement evaluates to basic_deps.
+ """
+ return select({
+ "//third_party/llvm:using_polly": basic_deps.append("@llvm//:polly"),
+ "//conditions:default": basic_deps
+ })
+

tensorflow/compiler/xla/service/cpu/compiler_functor.cc

@@ -22,6 +22,9 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "external/llvm/tools/polly/include/polly/LinkAllPasses.h"
+#include "external/llvm/tools/polly/include/polly/RegisterPasses.h"
+
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -312,6 +315,11 @@ void CompilerFunctor::AddOptimizationPasses(
  builder.SLPVectorize = opt_level > 1 && size_level == 0;
 
  builder.populateFunctionPassManager(*function_passes);
+
+ llvm::PassRegistry& Registry = *llvm::PassRegistry::getPassRegistry();
+ ::polly::initializePollyPasses(Registry);
+ ::polly::registerPollyPasses(*module_passes);
+
  builder.populateModulePassManager(*module_passes);
 }
 

third_party/llvm/BUILD

@@ -6,3 +6,10 @@ py_binary(
  srcs_version = "PY2AND3",
  visibility = ["@llvm//:__subpackages__"],
 )
+
+py_binary(
+ name = "expand_isl_cmake_vars",
+ srcs = ["expand_isl_cmake_vars.py"],
+ srcs_version = "PY2AND3",
+ visibility = ["@llvm//:__subpackages__"],
+)

third_party/llvm/llvm.bzl

@@ -85,6 +85,26 @@ def cmake_var_string(cmake_vars):
  return " ".join([_quote("{}={}".format(k, str(v)))
  for (k, v) in cmake_vars.items()])
 
+def expand_isl_cmake_vars(name, src, dst, cmake_vars):
+ """Expands #cmakedefine, #cmakedefine01, and CMake variables in a text file.
+
+ Args:
+ name: the name of the rule
+ src: the input of the rule
+ dst: the output of the rule
+ cmake_vars: a string containing the CMake variables, as generated by
+ cmake_var_string.
+ """
+ expand_isl_cmake_vars_tool = "@//third_party/llvm:expand_isl_cmake_vars"
+ native.genrule(
+ name = name,
+ srcs = [src],
+ tools = [expand_isl_cmake_vars_tool],
+ outs = [dst],
+ cmd = ("$(location {}) ".format(expand_isl_cmake_vars_tool) + cmake_vars +
+ "< $< > $@")
+ )
+
 def expand_cmake_vars(name, src, dst, cmake_vars):
  """Expands #cmakedefine, #cmakedefine01, and CMake variables in a text file.