# -*- Python -*- """Repository rule for ROCm autoconfiguration. `rocm_configure` depends on the following environment variables: * `TF_NEED_ROCM`: Whether to enable building with ROCm. * `GCC_HOST_COMPILER_PATH`: The GCC host compiler path * `ROCM_TOOLKIT_PATH`: The path to the ROCm toolkit. Default is `/opt/rocm`. * `TF_ROCM_VERSION`: The version of the ROCm toolkit. If this is blank, then use the system default. * `TF_MIOPEN_VERSION`: The version of the MIOpen library. * `TF_ROCM_AMDGPU_TARGETS`: The AMDGPU targets. Default is `gfx803,gfx900`. """ _GCC_HOST_COMPILER_PATH = "GCC_HOST_COMPILER_PATH" _ROCM_TOOLKIT_PATH = "ROCM_TOOLKIT_PATH" _TF_ROCM_VERSION = "TF_ROCM_VERSION" _TF_MIOPEN_VERSION = "TF_MIOPEN_VERSION" _TF_ROCM_AMDGPU_TARGETS = "TF_ROCM_AMDGPU_TARGETS" _TF_ROCM_CONFIG_REPO = "TF_ROCM_CONFIG_REPO" _DEFAULT_ROCM_VERSION = "" _DEFAULT_MIOPEN_VERSION = "" _DEFAULT_ROCM_TOOLKIT_PATH = "/opt/rocm" _DEFAULT_ROCM_AMDGPU_TARGETS = ["gfx803", "gfx900"] def find_cc(repository_ctx): """Find the C++ compiler.""" # Return a dummy value for GCC detection here to avoid error 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 <...>" def _cxx_inc_convert(path): """Convert path returned by cc -E xc++ in a complete path.""" path = path.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 rocm_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 [ str(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 = depset(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 rocm configuration fails.""" red = "\033[0;31m" no_color = "\033[0m" fail("\n%sROCm 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) # Add numpy headers inc_dirs.append("/usr/lib/python2.7/dist-packages/numpy/core/include") entries = [] for inc_dir in inc_dirs: entries.append(" cxx_builtin_include_directory: \"%s\"" % inc_dir) # define TENSORFLOW_USE_ROCM entries.append(" unfiltered_cxx_flag: \"-DTENSORFLOW_USE_ROCM\"") return "\n".join(entries) def _rocm_include_path(repository_ctx, rocm_config): """Generates the cxx_builtin_include_directory entries for rocm 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 = [] # general ROCm include path inc_dirs.append(rocm_config.rocm_toolkit_path + "/include") # Add HSA headers inc_dirs.append("/opt/rocm/hsa/include") # Add HIP headers inc_dirs.append("/opt/rocm/include/hip") inc_dirs.append("/opt/rocm/include/hip/hcc_detail") # Add rocrand and hiprand headers inc_dirs.append("/opt/rocm/rocrand/include") inc_dirs.append("/opt/rocm/hiprand/include") # Add rocfft headers inc_dirs.append("/opt/rocm/rocfft/include") # Add rocBLAS headers inc_dirs.append("/opt/rocm/rocblas/include") # Add MIOpen headers inc_dirs.append("/opt/rocm/miopen/include") # Add hcc headers inc_dirs.append("/opt/rocm/hcc/include") inc_dirs.append("/opt/rocm/hcc/compiler/lib/clang/7.0.0/include/") inc_dirs.append("/opt/rocm/hcc/lib/clang/7.0.0/include") # Newer hcc builds use/are based off of clang 8.0.0. inc_dirs.append("/opt/rocm/hcc/compiler/lib/clang/8.0.0/include/") inc_dirs.append("/opt/rocm/hcc/lib/clang/8.0.0/include") inc_entries = [] for inc_dir in inc_dirs: inc_entries.append(" cxx_builtin_include_directory: \"%s\"" % inc_dir) return "\n".join(inc_entries) 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_toolkit_path(repository_ctx): """Finds the rocm toolkit directory. Args: repository_ctx: The repository context. Returns: A speculative real path of the rocm toolkit install directory. """ rocm_toolkit_path = _DEFAULT_ROCM_TOOLKIT_PATH if _ROCM_TOOLKIT_PATH in repository_ctx.os.environ: rocm_toolkit_path = repository_ctx.os.environ[_ROCM_TOOLKIT_PATH].strip() if not repository_ctx.path(rocm_toolkit_path).exists: auto_configure_fail("Cannot find rocm toolkit path.") return str(repository_ctx.path(rocm_toolkit_path).realpath) def _amdgpu_targets(repository_ctx): """Returns a list of strings representing AMDGPU targets.""" if _TF_ROCM_AMDGPU_TARGETS not in repository_ctx.os.environ: return _DEFAULT_ROCM_AMDGPU_TARGETS amdgpu_targets_str = repository_ctx.os.environ[_TF_ROCM_AMDGPU_TARGETS] amdgpu_targets = amdgpu_targets_str.split(",") for amdgpu_target in amdgpu_targets: if amdgpu_target[:3] != "gfx" or not amdgpu_target[3:].isdigit(): auto_configure_fail("Invalid AMDGPU target: %s" % amdgpu_target) return amdgpu_targets 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 _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 "hip" 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"): 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 elif version: version = ".%s" % version return "lib%s%s.dylib" % (lib, version) else: auto_configure_fail("Invalid cpu_value: %s" % cpu_value) def _find_rocm_lib( lib, repository_ctx, cpu_value, basedir, version = "", static = False): """Finds the given ROCm libraries on the system. Args: lib: The name of the library, such as "hip" repository_ctx: The repository context. cpu_value: The name of the host operating system. basedir: The install directory of ROCm. 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)) 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 rocm library %s" % file_name) def _find_libs(repository_ctx, rocm_config): """Returns the ROCm libraries on the system. Args: repository_ctx: The repository context. rocm_config: The ROCm config as returned by _get_rocm_config Returns: Map of library names to structs of filename and path as returned by _find_rocm_lib. """ cpu_value = rocm_config.cpu_value return { "hip": _find_rocm_lib( "hip_hcc", repository_ctx, cpu_value, rocm_config.rocm_toolkit_path, ), "rocblas": _find_rocm_lib( "rocblas", repository_ctx, cpu_value, rocm_config.rocm_toolkit_path + "/rocblas", ), "rocfft": _find_rocm_lib( "rocfft", repository_ctx, cpu_value, rocm_config.rocm_toolkit_path + "/rocfft", ), "hiprand": _find_rocm_lib( "hiprand", repository_ctx, cpu_value, rocm_config.rocm_toolkit_path + "/hiprand", ), "miopen": _find_rocm_lib( "MIOpen", repository_ctx, cpu_value, rocm_config.rocm_toolkit_path + "/miopen", ), } def _get_rocm_config(repository_ctx): """Detects and returns information about the ROCm installation on the system. Args: repository_ctx: The repository context. Returns: A struct containing the following fields: rocm_toolkit_path: The ROCm toolkit installation directory. amdgpu_targets: A list of the system's AMDGPU targets. cpu_value: The name of the host operating system. """ cpu_value = _cpu_value(repository_ctx) rocm_toolkit_path = _rocm_toolkit_path(repository_ctx) return struct( rocm_toolkit_path = rocm_toolkit_path, amdgpu_targets = _amdgpu_targets(repository_ctx), 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=rocm 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 rocm/. _tpl( repository_ctx, "rocm:build_defs.bzl", { "%{rocm_is_configured}": "False", "%{rocm_extra_copts}": "[]", }, ) _tpl( repository_ctx, "rocm:BUILD", { "%{hip_lib}": _lib_name("hip", cpu_value), "%{rocblas_lib}": _lib_name("rocblas", cpu_value), "%{miopen_lib}": _lib_name("miopen", cpu_value), "%{rocfft_lib}": _lib_name("rocfft", cpu_value), "%{hiprand_lib}": _lib_name("hiprand", cpu_value), "%{rocm_include_genrules}": "", "%{rocm_headers}": "", }, ) # Create dummy files for the ROCm toolkit since they are still required by # tensorflow/core/platform/default/build_config:rocm. repository_ctx.file("rocm/hip/include/hip/hip_runtime.h", "") # Set up rocm_config.h, which is used by # tensorflow/stream_executor/dso_loader.cc. _tpl( repository_ctx, "rocm:rocm_config.h", { "%{rocm_toolkit_path}": _DEFAULT_ROCM_TOOLKIT_PATH, }, "rocm/rocm/rocm_config.h", ) # If rocm_configure is not configured to build with GPU support, and the user # attempts to build with --config=rocm, 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 = [] # We clear folders that might have been generated previously to avoid # undesired inclusions command.append('if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi') command.append('if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi') 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 = "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 ],\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. """ find_result = _execute( repository_ctx, ["find", src_dir, "-follow", "-type", "f"], empty_stdout_fine = True, ) result = find_result.stdout return result def _compute_rocm_extra_copts(repository_ctx, amdgpu_targets): if False: amdgpu_target_flags = ["--amdgpu-target=" + amdgpu_target for amdgpu_target in amdgpu_targets] else: # AMDGPU targets are handled in the "crosstool_wrapper_driver_is_not_gcc" amdgpu_target_flags = [] return str(amdgpu_target_flags) def _create_local_rocm_repository(repository_ctx): """Creates the repository containing files set up to build with ROCm.""" rocm_config = _get_rocm_config(repository_ctx) # Set up symbolic links for the rocm toolkit by creating genrules to do # symlinking. We create one genrule for each directory we want to track under # rocm_toolkit_path rocm_toolkit_path = rocm_config.rocm_toolkit_path rocm_include_path = rocm_toolkit_path + "/include" genrules = [_symlink_genrule_for_dir( repository_ctx, rocm_include_path, "rocm/include", "rocm-include", )] genrules.append(_symlink_genrule_for_dir( repository_ctx, rocm_toolkit_path + "/rocfft/include", "rocm/include/rocfft", "rocfft-include", )) genrules.append(_symlink_genrule_for_dir( repository_ctx, rocm_toolkit_path + "/rocblas/include", "rocm/include/rocblas", "rocblas-include", )) genrules.append(_symlink_genrule_for_dir( repository_ctx, rocm_toolkit_path + "/miopen/include", "rocm/include/miopen", "miopen-include", )) 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("rocm/lib/" + lib.file_name) genrules.append(_symlink_genrule_for_dir( repository_ctx, None, "", "rocm-lib", rocm_lib_src, rocm_lib_dest, )) included_files = _read_dir(repository_ctx, rocm_include_path).replace( rocm_include_path, "", ).splitlines() # 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.amdgpu_targets, ), }, ) _tpl( repository_ctx, "rocm:BUILD", { "%{hip_lib}": rocm_libs["hip"].file_name, "%{rocblas_lib}": rocm_libs["rocblas"].file_name, "%{rocfft_lib}": rocm_libs["rocfft"].file_name, "%{hiprand_lib}": rocm_libs["hiprand"].file_name, "%{miopen_lib}": rocm_libs["miopen"].file_name, "%{rocm_include_genrules}": "\n".join(genrules), "%{rocm_headers}": ('":rocm-include",\n' + '":rocfft-include",\n' + '":rocblas-include",\n' + '":miopen-include",'), }, ) # Set up crosstool/ _tpl(repository_ctx, "crosstool:BUILD", {"%{linker_files}": ":empty", "%{win_linker_files}": ":empty"}) 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, "%{clang_path}": str(cc), } _tpl(repository_ctx, "crosstool:CROSSTOOL_hipcc", rocm_defines, out = "crosstool/CROSSTOOL") _tpl( repository_ctx, "crosstool:clang/bin/crosstool_wrapper_driver_rocm", { "%{cpu_compiler}": str(cc), "%{hipcc_path}": "/opt/rocm/bin/hipcc", "%{gcc_host_compiler_path}": str(cc), "%{rocm_amdgpu_targets}": ",".join( ["\"%s\"" % c for c in rocm_config.amdgpu_targets], ), }, ) # Set up rocm_config.h, which is used by # tensorflow/stream_executor/dso_loader.cc. _tpl( repository_ctx, "rocm:rocm_config.h", { "%{rocm_amdgpu_targets}": ",".join( ["\"%s\"" % c for c in rocm_config.amdgpu_targets], ), "%{rocm_toolkit_path}": rocm_config.rocm_toolkit_path, }, "rocm/rocm/rocm_config.h", ) def _create_remote_rocm_repository(repository_ctx, remote_config_repo): """Creates pointers to a remotely configured repo set up to build with ROCm.""" _tpl( repository_ctx, "rocm:build_defs.bzl", { "%{rocm_is_configured}": "True", "%{rocm_extra_copts}": _compute_rocm_extra_copts( repository_ctx, #_compute_capabilities(repository_ctx) ), }, ) _tpl( repository_ctx, "rocm:remote.BUILD", { "%{remote_rocm_repo}": remote_config_repo, }, "rocm/BUILD", ) _tpl(repository_ctx, "crosstool:remote.BUILD", { "%{remote_rocm_repo}": remote_config_repo, }, "crosstool/BUILD") def _rocm_autoconf_impl(repository_ctx): """Implementation of the rocm_autoconf repository rule.""" if not _enable_rocm(repository_ctx): _create_dummy_repository(repository_ctx) elif _TF_ROCM_CONFIG_REPO in repository_ctx.os.environ: _create_remote_rocm_repository( repository_ctx, repository_ctx.os.environ[_TF_ROCM_CONFIG_REPO], ) else: _create_local_rocm_repository(repository_ctx) rocm_configure = repository_rule( implementation = _rocm_autoconf_impl, environ = [ _GCC_HOST_COMPILER_PATH, "TF_NEED_ROCM", _ROCM_TOOLKIT_PATH, _TF_ROCM_VERSION, _TF_MIOPEN_VERSION, _TF_ROCM_AMDGPU_TARGETS, _TF_ROCM_CONFIG_REPO, ], ) """Detects and configures the local ROCm toolchain. Add the following to your WORKSPACE FILE: ```python rocm_configure(name = "local_config_rocm") ``` Args: name: A unique name for this workspace rule. """