github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/toolchains/preconfig/generate/generate.bzl (about)

     1  load(
     2      "@bazel_toolchains//rules:docker_config.bzl",
     3      "docker_toolchain_autoconfig",
     4  )
     5  
     6  def _tensorflow_rbe_config(name, compiler, python_version, os, rocm_version = None, cuda_version = None, cudnn_version = None, tensorrt_version = None, tensorrt_install_path = None, cudnn_install_path = None, compiler_prefix = None, build_bazel_src = False):
     7      base = "@%s//image" % os
     8      config_repos = [
     9          "local_config_python",
    10          "local_config_cc",
    11      ]
    12      env = {
    13          "ABI_VERSION": "gcc",
    14          "ABI_LIBC_VERSION": "glibc_2.19",
    15          "BAZEL_COMPILER": compiler,
    16          "BAZEL_HOST_SYSTEM": "i686-unknown-linux-gnu",
    17          "BAZEL_TARGET_LIBC": "glibc_2.19",
    18          "BAZEL_TARGET_CPU": "k8",
    19          "BAZEL_TARGET_SYSTEM": "x86_64-unknown-linux-gnu",
    20          "CC_TOOLCHAIN_NAME": "linux_gnu_x86",
    21          "CC": compiler,
    22          "PYTHON_BIN_PATH": "/usr/bin/python%s" % python_version,
    23          "CLEAR_CACHE": "1",
    24          "HOST_CXX_COMPILER": compiler,
    25          "HOST_C_COMPILER": compiler,
    26      }
    27  
    28      if cuda_version != None and rocm_version != None:
    29          fail("Specifying both cuda_version and rocm_version is not supported.")
    30  
    31      if cuda_version != None:
    32          base = "@cuda%s-cudnn%s-%s//image" % (cuda_version, cudnn_version, os)
    33  
    34          # The cuda toolchain currently contains its own C++ toolchain definition,
    35          # so we do not fetch local_config_cc.
    36          config_repos = [
    37              "local_config_python",
    38              "local_config_cuda",
    39              "local_config_tensorrt",
    40          ]
    41          env.update({
    42              "TF_NEED_CUDA": "1",
    43              "TF_CUDA_CLANG": "1" if compiler == "clang" else "0",
    44              "TF_CUDA_COMPUTE_CAPABILITIES": "3.0,6.0",
    45              "TF_ENABLE_XLA": "1",
    46              "TF_CUDNN_VERSION": cudnn_version,
    47              "TF_CUDA_VERSION": cuda_version,
    48              "CUDNN_INSTALL_PATH": cudnn_install_path if cudnn_install_path != None else "/usr/lib/x86_64-linux-gnu",
    49              "TF_NEED_TENSORRT": "1",
    50              "TF_TENSORRT_VERSION": tensorrt_version,
    51              "TENSORRT_INSTALL_PATH": tensorrt_install_path if tensorrt_install_path != None else "/usr/lib/x86_64-linux-gnu",
    52              "GCC_HOST_COMPILER_PATH": compiler if compiler != "clang" else "",
    53              "GCC_HOST_COMPILER_PREFIX": compiler_prefix if compiler_prefix != None else "/usr/bin",
    54          })
    55  
    56      if rocm_version != None:
    57          base = "@rocm-%s//image" % (os)
    58  
    59          # The rocm toolchain currently contains its own C++ toolchain definition,
    60          # so we do not fetch local_config_cc.
    61          config_repos = [
    62              "local_config_python",
    63              "local_config_rocm",
    64          ]
    65          env.update({
    66              "TF_NEED_ROCM": "1",
    67              "TF_ENABLE_XLA": "0",
    68          })
    69  
    70      docker_toolchain_autoconfig(
    71          name = name,
    72          base = base,
    73          bazel_version = "0.24.1",
    74          build_bazel_src = build_bazel_src,
    75          config_repos = config_repos,
    76          env = env,
    77          mount_project = "$(mount_project)",
    78          tags = ["manual"],
    79      )
    80  
    81  tensorflow_rbe_config = _tensorflow_rbe_config