rocm_configure.bzl

# -*- Python -*-
"""Repository rule for ROCM autoconfiguration.

`rocm_configure` depends on the following environment variables:

  * `TF_NEED_ROCM`: Whether to enable building with CUDA.
  * `GCC_HOST_COMPILER_PATH`: The GCC host compiler path
  * `ROCM_INSTALL_PATH`: The path to the CUDA toolkit. Default is
    `/opt/rocm`.
"""

_GCC_HOST_COMPILER_PATH = "GCC_HOST_COMPILER_PATH"
_ROCM_INSTALL_PATH = "ROCM_INSTALL_PATH"

_DEFAULT_ROCM_INSTALL_PATH = "/opt/rocm"


# TODO(dzc): Once these functions have been factored out of Bazel's
# cc_configure.bzl, load them from @bazel_tools instead.
# BEGIN cc_configure common functions.
def find_cc(repository_ctx):
  """Find the C++ compiler."""
  # On Windows, we use Bazel's MSVC CROSSTOOL for GPU build
  # Return a dummy value for GCC detection here to avoid error
  if _is_windows(repository_ctx):
    return "/use/--config=win-cuda --cpu=x64_windows_msvc/instead"

  target_cc_name = "gcc"
  cc_path_envvar = _GCC_HOST_COMPILER_PATH
  cc_name = target_cc_name

  if cc_path_envvar in repository_ctx.os.environ:
    cc_name_from_env = repository_ctx.os.environ[cc_path_envvar].strip()
    if cc_name_from_env:
      cc_name = cc_name_from_env
  if cc_name.startswith("/"):
    # Absolute path, maybe we should make this supported by our which function.
    return cc_name
  cc = repository_ctx.which(cc_name)
  if cc == None:
    fail(("Cannot find {}, either correct your path or set the {}" +
          " environment variable").format(target_cc_name, cc_path_envvar))
  return cc


_INC_DIR_MARKER_BEGIN = "#include <...>"


# OSX add " (framework directory)" at the end of line, strip it.
_OSX_FRAMEWORK_SUFFIX = " (framework directory)"
_OSX_FRAMEWORK_SUFFIX_LEN =  len(_OSX_FRAMEWORK_SUFFIX)
def _cxx_inc_convert(path):
  """Convert path returned by cc -E xc++ in a complete path."""
  path = path.strip()
  if path.endswith(_OSX_FRAMEWORK_SUFFIX):
    path = path[:-_OSX_FRAMEWORK_SUFFIX_LEN].strip()
  return path

def _get_cxx_inc_directories_impl(repository_ctx, cc, lang_is_cpp):
  """Compute the list of default C or C++ include directories."""
  if lang_is_cpp:
    lang = "c++"
  else:
    lang = "c"
  # TODO: We pass -no-canonical-prefixes here to match the compiler flags,
  #       but in cuda_clang CROSSTOOL file that is a `feature` and we should
  #       handle the case when it's disabled and no flag is passed
  result = repository_ctx.execute([cc, "-no-canonical-prefixes",
                                   "-E", "-x" + lang, "-", "-v"])
  index1 = result.stderr.find(_INC_DIR_MARKER_BEGIN)
  if index1 == -1:
    return []
  index1 = result.stderr.find("\n", index1)
  if index1 == -1:
    return []
  index2 = result.stderr.rfind("\n ")
  if index2 == -1 or index2 < index1:
    return []
  index2 = result.stderr.find("\n", index2 + 1)
  if index2 == -1:
    inc_dirs = result.stderr[index1 + 1:]
  else:
    inc_dirs = result.stderr[index1 + 1:index2].strip()

  return [repository_ctx.path(_cxx_inc_convert(p))
          for p in inc_dirs.split("\n")]

def get_cxx_inc_directories(repository_ctx, cc):
  """Compute the list of default C and C++ include directories."""
  # For some reason `clang -xc` sometimes returns include paths that are
  # different from the ones from `clang -xc++`. (Symlink and a dir)
  # So we run the compiler with both `-xc` and `-xc++` and merge resulting lists
  includes_cpp = _get_cxx_inc_directories_impl(repository_ctx, cc, True)
  includes_c = _get_cxx_inc_directories_impl(repository_ctx, cc, False)

  includes_cpp_set = set(includes_cpp)
  return includes_cpp + [inc for inc in includes_c
                         if inc not in includes_cpp_set]


def auto_configure_fail(msg):
  """Output failure message when cuda configuration fails."""
  red = "\033[0;31m"
  no_color = "\033[0m"
  fail("\n%sCuda Configuration Error:%s %s\n" % (red, no_color, msg))
# END cc_configure common functions (see TODO above).


def _host_compiler_includes(repository_ctx, cc):
  """Generates the cxx_builtin_include_directory entries for gcc inc dirs.

  Args:
    repository_ctx: The repository context.
    cc: The path to the gcc host compiler.

  Returns:
    A string containing the cxx_builtin_include_directory for each of the gcc
    host compiler include directories, which can be added to the CROSSTOOL
    file.
  """
  inc_dirs = get_cxx_inc_directories(repository_ctx, cc)
  inc_entries = []
  for inc_dir in inc_dirs:
    inc_entries.append("  cxx_builtin_include_directory: \"%s\"" % inc_dir)
  return "\n".join(inc_entries)

def _rocm_include_path(repository_ctx, rocm_config):
  """Generates the cxx_builtin_include_directory entries for cuda inc dirs.

  Args:
    repository_ctx: The repository context.
    cc: The path to the gcc host compiler.

  Returns:
    A string containing the cxx_builtin_include_directory for each of the gcc
    host compiler include directories, which can be added to the CROSSTOOL
    file.
  """
  return "\n".join("/opt/rocm/include")


def _enable_rocm(repository_ctx):
  if "TF_NEED_ROCM" in repository_ctx.os.environ:
    enable_rocm = repository_ctx.os.environ["TF_NEED_ROCM"].strip()
    return enable_rocm == "1"
  return False


def _rocm_install_path(repository_ctx):
  """Finds the cuda toolkit directory.

  Args:
    repository_ctx: The repository context.

  Returns:
    A speculative real path of the cuda toolkit install directory.
  """
  rocm_install_path = _DEFAULT_ROCM_INSTALL_PATH
  if _ROCM_INSTALL_PATH in repository_ctx.os.environ:
    rocm_install_path = repository_ctx.os.environ[_ROCM_INSTALL_PATH].strip()
  if not repository_ctx.path(rocm_install_path).exists:
    auto_configure_fail("Cannot find rocm install path.")
  return str(repository_ctx.path(rocm_install_path).realpath)


def _matches_version(environ_version, detected_version):
  """Checks whether the user-specified version matches the detected version.

  This function performs a weak matching so that if the user specifies only the
  major or major and minor versions, the versions are still considered matching
  if the version parts match. To illustrate:

      environ_version  detected_version  result
      -----------------------------------------
      5.1.3            5.1.3             True
      5.1              5.1.3             True
      5                5.1               True
      5.1.3            5.1               False
      5.2.3            5.1.3             False

  Args:
    environ_version: The version specified by the user via environment
      variables.
    detected_version: The version autodetected from the CUDA installation on
      the system.

  Returns: True if user-specified version matches detected version and False
    otherwise.
  """
  environ_version_parts = environ_version.split(".")
  detected_version_parts = detected_version.split(".")
  if len(detected_version_parts) < len(environ_version_parts):
    return False
  for i, part in enumerate(detected_version_parts):
    if i >= len(environ_version_parts):
      break
    if part != environ_version_parts[i]:
      return False
  return True


_NVCC_VERSION_PREFIX = "Cuda compilation tools, release "


def _find_cuda_define(repository_ctx, cudnn_header_dir, define):
  """Returns the value of a #define in cudnn.h

  Greps through cudnn.h and returns the value of the specified #define. If the
  #define is not found, then raise an error.

  Args:
    repository_ctx: The repository context.
    cudnn_header_dir: The directory containing the cuDNN header.
    define: The #define to search for.

  Returns:
    The value of the #define found in cudnn.h.
  """
  # Confirm location of cudnn.h and grep for the line defining CUDNN_MAJOR.
  cudnn_h_path = repository_ctx.path("%s/cudnn.h" % cudnn_header_dir)
  if not cudnn_h_path.exists:
    auto_configure_fail("Cannot find cudnn.h at %s" % str(cudnn_h_path))
  result = repository_ctx.execute(["grep", "--color=never", "-E", define, str(cudnn_h_path)])
  if result.stderr:
    auto_configure_fail("Error reading %s: %s" %
                        (result.stderr, str(cudnn_h_path)))

  # Parse the cuDNN major version from the line defining CUDNN_MAJOR
  lines = result.stdout.splitlines()
  if len(lines) == 0 or lines[0].find(define) == -1:
    auto_configure_fail("Cannot find line containing '%s' in %s" %
                        (define, str(cudnn_h_path)))
  return lines[0].replace(define, "").strip()


def _cpu_value(repository_ctx):
  """Returns the name of the host operating system.

  Args:
    repository_ctx: The repository context.

  Returns:
    A string containing the name of the host operating system.
  """
  os_name = repository_ctx.os.name.lower()
  if os_name.startswith("mac os"):
    return "Darwin"
  if os_name.find("windows") != -1:
    return "Windows"
  result = repository_ctx.execute(["uname", "-s"])
  return result.stdout.strip()


def _is_windows(repository_ctx):
  """Returns true if the host operating system is windows."""
  return _cpu_value(repository_ctx) == "Windows"

def _lib_name(lib, cpu_value, version="", static=False):
  """Constructs the platform-specific name of a library.

  Args:
    lib: The name of the library, such as "cudart"
    cpu_value: The name of the host operating system.
    version: The version of the library.
    static: True the library is static or False if it is a shared object.

  Returns:
    The platform-specific name of the library.
  """
  if cpu_value in ("Linux", "FreeBSD"):
    if static:
      return "lib%s.a" % lib
    else:
      if version:
        version = ".%s" % version
      return "lib%s.so%s" % (lib, version)
  elif cpu_value == "Windows":
    return "%s.lib" % lib
  elif cpu_value == "Darwin":
    if static:
      return "lib%s.a" % lib
    else:
      if version:
        version = ".%s" % version
    return "lib%s%s.dylib" % (lib, version)
  else:
    auto_configure_fail("Invalid cpu_value: %s" % cpu_value)


def _find_cuda_lib(lib, repository_ctx, cpu_value, basedir, version="",
                   static=False):
  """Finds the given CUDA or cuDNN library on the system.

  Args:
    lib: The name of the library, such as "cudart"
    repository_ctx: The repository context.
    cpu_value: The name of the host operating system.
    basedir: The install directory of CUDA or cuDNN.
    version: The version of the library.
    static: True if static library, False if shared object.

  Returns:
    Returns a struct with the following fields:
      file_name: The basename of the library found on the system.
      path: The full path to the library.
  """
  file_name = _lib_name(lib, cpu_value, version, static)
  if cpu_value == "Linux":
    path = repository_ctx.path("%s/lib64/%s" % (basedir, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))
    path = repository_ctx.path("%s/lib64/stubs/%s" % (basedir, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))
    path = repository_ctx.path(
        "%s/lib/x86_64-linux-gnu/%s" % (basedir, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))

  elif cpu_value == "Windows":
    path = repository_ctx.path("%s/lib/x64/%s" % (basedir, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))

  path = repository_ctx.path("%s/lib/%s" % (basedir, file_name))
  if path.exists:
    return struct(file_name=file_name, path=str(path.realpath))
  path = repository_ctx.path("%s/%s" % (basedir, file_name))
  if path.exists:
    return struct(file_name=file_name, path=str(path.realpath))

  auto_configure_fail("Cannot find cuda library %s" % file_name)


def _find_cupti_lib(repository_ctx, cuda_config):
  """Finds the cupti library on the system.

  On most systems, the cupti library is not installed in the same directory as
  the other CUDA libraries but rather in a special extras/CUPTI directory.

  Args:
    repository_ctx: The repository context.
    cuda_config: The cuda configuration as returned by _get_cuda_config.

  Returns:
    Returns a struct with the following fields:
      file_name: The basename of the library found on the system.
      path: The full path to the library.
  """
  file_name = _lib_name("cupti", cuda_config.cpu_value,
                        cuda_config.cuda_version)
  if cuda_config.cpu_value == "Linux":
    path = repository_ctx.path(
        "%s/extras/CUPTI/lib64/%s" % (cuda_config.cuda_toolkit_path, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))

    path = repository_ctx.path(
        "%s/lib/x86_64-linux-gnu/%s" % (cuda_config.cuda_toolkit_path,
                                        file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))

  elif cuda_config.cpu_value == "Windows":
    path = repository_ctx.path(
        "%s/extras/CUPTI/libx64/%s" %
        (cuda_config.cuda_toolkit_path, file_name))
    if path.exists:
      return struct(file_name=file_name, path=str(path.realpath))

  path = repository_ctx.path(
      "%s/extras/CUPTI/lib/%s" % (cuda_config.cuda_toolkit_path, file_name))
  if path.exists:
    return struct(file_name=file_name, path=str(path.realpath))

  path = repository_ctx.path(
      "%s/lib/%s" % (cuda_config.cuda_toolkit_path, file_name))
  if path.exists:
    return struct(file_name=file_name, path=str(path.realpath))

  auto_configure_fail("Cannot find cupti library %s" % file_name)

def _find_libs(repository_ctx, cuda_config):
  """Returns the CUDA and cuDNN libraries on the system.

  Args:
    repository_ctx: The repository context.
    cuda_config: The CUDA config as returned by _get_cuda_config

  Returns:
    Map of library names to structs of filename and path as returned by
    _find_cuda_lib and _find_cupti_lib.
  """
  cudnn_version = cuda_config.cudnn_version
  cudnn_ext = ".%s" % cudnn_version if cudnn_version else ""
  cpu_value = cuda_config.cpu_value
  return {
      "cuda": _find_cuda_lib("cuda", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path),
      "cudart": _find_cuda_lib(
          "cudart", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
          cuda_config.cuda_version),
      "cudart_static": _find_cuda_lib(
          "cudart_static", repository_ctx, cpu_value,
          cuda_config.cuda_toolkit_path, cuda_config.cuda_version, static=True),
      "cublas": _find_cuda_lib(
          "cublas", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
          cuda_config.cuda_version),
      "cusolver": _find_cuda_lib(
          "cusolver", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
          cuda_config.cuda_version),
      "curand": _find_cuda_lib(
          "curand", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
          cuda_config.cuda_version),
      "cufft": _find_cuda_lib(
          "cufft", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path,
          cuda_config.cuda_version),
      "cudnn": _find_cuda_lib(
          "cudnn", repository_ctx, cpu_value, cuda_config.cudnn_install_basedir,
          cuda_config.cudnn_version),
      "cupti": _find_cupti_lib(repository_ctx, cuda_config),
  }


def _find_cudnn_header_dir(repository_ctx, cudnn_install_basedir):
  """Returns the path to the directory containing cudnn.h

  Args:
    repository_ctx: The repository context.
    cudnn_install_basedir: The cudnn install directory as returned by
      _cudnn_install_basedir.

  Returns:
    The path of the directory containing the cudnn header.
  """
  if repository_ctx.path(cudnn_install_basedir + "/cudnn.h").exists:
    return cudnn_install_basedir
  if repository_ctx.path(cudnn_install_basedir + "/include/cudnn.h").exists:
    return cudnn_install_basedir + "/include"
  if repository_ctx.path("/usr/include/cudnn.h").exists:
    return "/usr/include"
  auto_configure_fail("Cannot find cudnn.h under %s" % cudnn_install_basedir)


def _find_cudnn_lib_path(repository_ctx, cudnn_install_basedir, symlink_files):
  """Returns the path to the directory containing libcudnn

  Args:
    repository_ctx: The repository context.
    cudnn_install_basedir: The cudnn install dir as returned by
      _cudnn_install_basedir.
    symlink_files: The symlink files as returned by _cuda_symlink_files.

  Returns:
    The path of the directory containing the cudnn libraries.
  """
  lib_dir = cudnn_install_basedir + "/" + symlink_files.cuda_dnn_lib
  if repository_ctx.path(lib_dir).exists:
    return lib_dir
  alt_lib_dir = cudnn_install_basedir + "/" + symlink_files.cuda_dnn_lib_alt
  if repository_ctx.path(alt_lib_dir).exists:
    return alt_lib_dir

  auto_configure_fail("Cannot find %s or %s under %s" %
       (symlink_files.cuda_dnn_lib, symlink_files.cuda_dnn_lib_alt,
        cudnn_install_basedir))


def _cudart_static_linkopt(cpu_value):
  """Returns additional platform-specific linkopts for cudart."""
  return "" if cpu_value == "Darwin" else "\"-lrt\","

def _get_rocm_config(repository_ctx):
  """Detects and returns information about the CUDA installation on the system.

  Args:
    repository_ctx: The repository context.

  Returns:
    A struct containing the following fields:
      cuda_toolkit_path: The CUDA toolkit installation directory.
      cudnn_install_basedir: The cuDNN installation directory.
      cuda_version: The version of CUDA on the system.
      cudnn_version: The version of cuDNN on the system.
      compute_capabilities: A list of the system's CUDA compute capabilities.
      cpu_value: The name of the host operating system.
  """
  cpu_value = _cpu_value(repository_ctx)
  rocm_install_path = _rocm_install_path(repository_ctx)
  return struct(
      rocm_install_path = rocm_install_path,
      cpu_value = cpu_value)


def _tpl(repository_ctx, tpl, substitutions={}, out=None):
  if not out:
    out = tpl.replace(":", "/")
  repository_ctx.template(
      out,
      Label("//third_party/gpus/%s.tpl" % tpl),
      substitutions)


def _file(repository_ctx, label):
  repository_ctx.template(
      label.replace(":", "/"),
      Label("//third_party/gpus/%s.tpl" % label),
      {})


_DUMMY_CROSSTOOL_BZL_FILE = """
def error_gpu_disabled():
  fail("ERROR: Building with --config=cuda but TensorFlow is not configured " +
       "to build with GPU support. Please re-run ./configure and enter 'Y' " +
       "at the prompt to build with GPU support.")

  native.genrule(
      name = "error_gen_crosstool",
      outs = ["CROSSTOOL"],
      cmd = "echo 'Should not be run.' && exit 1",
  )

  native.filegroup(
      name = "crosstool",
      srcs = [":CROSSTOOL"],
      output_licenses = ["unencumbered"],
  )
"""


_DUMMY_CROSSTOOL_BUILD_FILE = """
load("//crosstool:error_gpu_disabled.bzl", "error_gpu_disabled")

error_gpu_disabled()
"""


def _create_dummy_repository(repository_ctx):
  cpu_value = _cpu_value(repository_ctx)

  # Set up BUILD file for cuda/.
  _tpl(repository_ctx, "rocm:build_defs.bzl",
       {
           "%{rocm_is_configured}": "False",
           "%{rocm_extra_copts}": "[]"
       })
  _tpl(repository_ctx, "cuda:BUILD",
       {
           "%{rocm_driver_lib}": _lib_name("hip_hcc", cpu_value),
           "%{rocblas_lib}": _lib_name("rocblas", cpu_value),
           "%{rocm_include_genrules}": '',
           "%{rocm_headers}": '',
       })

  # Create dummy files for the CUDA toolkit since they are still required by
  # tensorflow/core/platform/default/build_config:cuda.
  repository_ctx.file("rocm/include/hip_runtime.h", "")
  repository_ctx.file("rocm/rocblas/include/rocblas.h", "")
  repository_ctx.file("rocm/hip/lib/%s" % _lib_name("hip_hcc", cpu_value))
  repository_ctx.file("rocm/rocblas/lib/%s" % _lib_name("rocblas", cpu_value))

  # Set up cuda_config.h, which is used by
  # tensorflow/stream_executor/dso_loader.cc.
  _tpl(repository_ctx, "rocm:rocm_config.h",
       {
           "%{rocm_install_path}": _DEFAULT_ROCM_INSTALL_PATH,
       })

  # If cuda_configure is not configured to build with GPU support, and the user
  # attempts to build with --config=cuda, add a dummy build rule to intercept
  # this and fail with an actionable error message.
  repository_ctx.file("crosstool/error_gpu_disabled.bzl",
                      _DUMMY_CROSSTOOL_BZL_FILE)
  repository_ctx.file("crosstool/BUILD", _DUMMY_CROSSTOOL_BUILD_FILE)


def _execute(repository_ctx, cmdline, error_msg=None, error_details=None,
             empty_stdout_fine=False):
  """Executes an arbitrary shell command.

  Args:
    repository_ctx: the repository_ctx object
    cmdline: list of strings, the command to execute
    error_msg: string, a summary of the error if the command fails
    error_details: string, details about the error or steps to fix it
    empty_stdout_fine: bool, if True, an empty stdout result is fine, otherwise
      it's an error
  Return:
    the result of repository_ctx.execute(cmdline)
  """
  result = repository_ctx.execute(cmdline)
  if result.stderr or not (empty_stdout_fine or result.stdout):
    auto_configure_fail(
        "\n".join([
            error_msg.strip() if error_msg else "Repository command failed",
            result.stderr.strip(),
            error_details if error_details else ""]))
  return result


def _norm_path(path):
  """Returns a path with '/' and remove the trailing slash."""
  path = path.replace("\\", "/")
  if path[-1] == "/":
    path = path[:-1]
  return path


def _symlink_genrule_for_dir(repository_ctx, src_dir, dest_dir, genrule_name,
    src_files = [], dest_files = []):
  """Returns a genrule to symlink(or copy if on Windows) a set of files.

  If src_dir is passed, files will be read from the given directory; otherwise
  we assume files are in src_files and dest_files
  """
  if src_dir != None:
    src_dir = _norm_path(src_dir)
    dest_dir = _norm_path(dest_dir)
    files = _read_dir(repository_ctx, src_dir)
    # Create a list with the src_dir stripped to use for outputs.
    dest_files = files.replace(src_dir, '').splitlines()
    src_files = files.splitlines()
  command = []
  outs = []
  for i in range(len(dest_files)):
    if dest_files[i] != "":
      # If we have only one file to link we do not want to use the dest_dir, as
      # $(@D) will include the full path to the file.
      dest = '$(@D)/' + dest_dir + dest_files[i] if len(dest_files) != 1 else '$(@D)/' + dest_files[i]
      # On Windows, symlink is not supported, so we just copy all the files.
      cmd = 'cp -f' if _is_windows(repository_ctx) else 'ln -s'
      command.append(cmd + ' "%s" "%s"' % (src_files[i] , dest))
      outs.append('      "' + dest_dir + dest_files[i] + '",')
  genrule = _genrule(src_dir, genrule_name, " && ".join(command),
                     "\n".join(outs))
  return genrule


def _genrule(src_dir, genrule_name, command, outs):
  """Returns a string with a genrule.

  Genrule executes the given command and produces the given outputs.
  """
  return (
      'genrule(\n' +
      '    name = "' +
      genrule_name + '",\n' +
      '    outs = [\n' +
      outs +
      '    ],\n' +
      '    cmd = """\n' +
      command +
      '    """,\n' +
      ')\n\n'
  )


def _read_dir(repository_ctx, src_dir):
  """Returns a string with all files in a directory.

  Finds all files inside a directory, traversing subfolders and following
  symlinks. The returned string contains the full path of all files
  separated by line breaks.
  """
  if _is_windows(repository_ctx):
    src_dir = src_dir.replace("/", "\\")
    find_result = _execute(
        repository_ctx, ["cmd.exe", "/c", "dir", src_dir, "/b", "/s", "/a-d"],
        empty_stdout_fine=True)
    # src_files will be used in genrule.outs where the paths must
    # use forward slashes.
    result = find_result.stdout.replace("\\", "/")
  else:
    find_result = _execute(
        repository_ctx, ["find", src_dir, "-follow", "-type", "f"],
        empty_stdout_fine=True)
    result = find_result.stdout
  return result


def _use_cuda_clang(repository_ctx):
  if "TF_CUDA_CLANG" in repository_ctx.os.environ:
    enable_cuda = repository_ctx.os.environ["TF_CUDA_CLANG"].strip()
    return enable_cuda == "1"
  return False

def _compute_rocm_extra_copts(repository_ctx, rocm_config):
  capability_flags = []
  return str(capability_flags)

def _create_rocm_repository(repository_ctx):
  """Creates the repository containing files set up to build with CUDA."""
  rocm_config = _get_rocm_config(repository_ctx)

  # Set up symbolic links for the cuda toolkit by creating genrules to do
  # symlinking. We create one genrule for each directory we want to track under
  # cuda_toolkit_path
  rocm_install_path = rocm_config.rocm_install_path
  rocm_include_path = rocm_install_path + "/include"
  genrules = [_symlink_genrule_for_dir(repository_ctx,
      rocm_include_path, "include", "rocm-include")]
  genrules.append(_symlink_genrule_for_dir(repository_ctx,
      rocm_install_path + "/nvvm", "nvvm", "cuda-nvvm"))

  rocm_libs = _find_libs(repository_ctx, rocm_config)
  rocm_lib_src = []
  rocm_lib_dest = []
  for lib in rocm_libs.values():
    rocm_lib_src.append(lib.path)
    rocm_lib_dest.append("lib/" + lib.file_name)
  genrules.append(_symlink_genrule_for_dir(repository_ctx, None, "", "rocm-lib",
                                       rocm_lib_src, rocm_lib_dest))

  # Set up BUILD file for rocm/
  _tpl(repository_ctx, "rocm:build_defs.bzl",
       {
           "%{rocm_is_configured}": "True",
           "%{rocm_extra_copts}": _compute_rocm_extra_copts(repository_ctx, rocm_config),

       })
  _tpl(repository_ctx, "rocm:BUILD",
       {
           "%{rocm_driver_lib}": rocm_libs["hip_hcc"].file_name,
           "%{rocblas_lib}": rocm_libs["rocblas"].file_name,
           "%{rocm_headers}": ('":rocm-include",\n')
       })
  # Set up crosstool/
  _file(repository_ctx, "crosstool:BUILD")
  cc = find_cc(repository_ctx)
  host_compiler_includes = _host_compiler_includes(repository_ctx, cc)
  rocm_defines = {
           "%{rocm_include_path}": _rocm_include_path(repository_ctx,
                                                      rocm_config),
           "%{host_compiler_includes}": host_compiler_includes,
       }
  rocm_defines["%{clang_path}"] = cc
  _tpl(repository_ctx, "crosstool:CROSSTOOL_clang", rocm_defines, out="crosstool/CROSSTOOL")

  # Set up rocm_config.h, which is used by
  # tensorflow/stream_executor/dso_loader.cc.
  _tpl(repository_ctx, "rocm:rocm_config.h",
       {
           "%{rocm_install_path}": rocm_config.rocm_install_path,
       })


def _rocm_autoconf_impl(repository_ctx):
  """Implementation of the cuda_autoconf repository rule."""
  if not _enable_rocm(repository_ctx):
    _create_dummy_repository(repository_ctx)
  else:
    _create_rocm_repository(repository_ctx)



rocm_configure = repository_rule(
    implementation = _rocm_autoconf_impl,
    environ = [
        _GCC_HOST_COMPILER_PATH,
        "TF_NEED_ROCM",
        _ROCM_INSTALL_PATH,
    ],
)

"""Detects and configures the local CUDA toolchain.

Add the following to your WORKSPACE FILE:

```python
cuda_configure(name = "local_config_cuda")
```

Args:
  name: A unique name for this workspace rule.
"""