"""Repository rule for ROCm autoconfiguration. `rocm_configure` depends on the following environment variables: * `TF_NEED_ROCM`: Whether to enable building with ROCm. * `TF_SYSROOT`: The sysroot to use when compiling. * `TF_ROCM_AMDGPU_TARGETS`: The AMDGPU targets. * `TF_ROCM_MULTIPLE_PATHS`: Colon-separated list of ROCm installation paths to merge. * `LLVM_PATH`: Path to LLVM installation (used with TF_ROCM_MULTIPLE_PATHS). * `TF_ROCM_RBE_DOCKER_IMAGE`: Docker image to be used in rbe worker to execute the action * `TF_ROCM_RBE_SINGLE_GPU_POOL`: The name of the rbe pool used to execute single gpu tests * `TF_ROCM_RBE_MULTI_GPU_POOL`: The name of the rbe pool used to execute multi gpu tests """ load( "//third_party/gpus/rocm:rocm_redist.bzl", "create_rocm_distro", "rocm_redist", ) load( "//third_party/remote_config:common.bzl", "config_repo_label", "err_out", "execute", "files_exist", "get_bash_bin", "get_cpu_value", "get_host_environ", "get_python_bin", ) load( ":cuda_configure.bzl", "enable_cuda", ) load( ":sycl_configure.bzl", "enable_sycl", ) _TF_ROCM_AMDGPU_TARGETS = "TF_ROCM_AMDGPU_TARGETS" _TF_ROCM_CONFIG_REPO = "TF_ROCM_CONFIG_REPO" _TF_ROCM_MULTIPLE_PATHS = "TF_ROCM_MULTIPLE_PATHS" _LLVM_PATH = "LLVM_PATH" _DISTRIBUTION_PATH = "rocm/rocm_dist" _ROCM_DISTRO_VERSION = "ROCM_DISTRO_VERSION" _ROCM_DISTRO_URL = "ROCM_DISTRO_URL" _ROCM_DISTRO_HASH = "ROCM_DISTRO_HASH" _ROCM_DISTRO_LINKS = "ROCM_DISTRO_LINKS" _TMPDIR = "TMPDIR" # Default hermetic ROCm redistributable version _DEFAULT_ROCM_DISTRO_VERSION = "rocm_7.13.0_gfx94X" _TF_ROCM_RBE_DOCKER_IMAGE = "TF_ROCM_RBE_DOCKER_IMAGE" _TF_ROCM_RBE_POOL = "TF_ROCM_RBE_POOL" _TF_ROCM_RBE_SINGLE_GPU_POOL = "TF_ROCM_RBE_SINGLE_GPU_POOL" _TF_ROCM_RBE_MULTI_GPU_POOL = "TF_ROCM_RBE_MULTI_GPU_POOL" _DEFAULT_TF_ROCM_RBE_POOL = "default" _DEFAULT_TF_ROCM_RBE_SINGLE_GPU_POOL = "linux_x64_gpu" _DEFAULT_TF_ROCM_RBE_MULTI_GPU_POOL = "linux_x64_multigpu" # rocm/tensorflow-build:latest-jammy-python3.11-rocm7.0.2 _DEFAULT_TF_ROCM_RBE_DOCKER_IMAGE = "rocm/tensorflow-build@sha256:a2672ff2510b369b4a5f034272a518dc93c2e492894e3befaeef19649632ccaa" 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)) # ROCm download helper functions def _get_file_name(url): """Extracts filename from URL.""" last_slash_index = url.rfind("/") return url[last_slash_index + 1:] def _download_package(repository_ctx, pkg): """Downloads and extracts a ROCm package. Args: repository_ctx: The repository context. pkg: Package dict with 'url', 'sha256', and optional 'sub_package'. """ file_name = _get_file_name(pkg["url"]) repository_ctx.report_progress("Downloading and extracting {}, expected hash is {}".format(pkg["url"], pkg["sha256"])) repository_ctx.download_and_extract( url = pkg["url"], output = _DISTRIBUTION_PATH, sha256 = pkg["sha256"], type = "zip" if pkg["url"].endswith(".whl") else "", ) if pkg.get("sub_package", None): repository_ctx.report_progress("Extracting {}".format(pkg["sub_package"])) repository_ctx.extract( archive = "{}/{}".format(_DISTRIBUTION_PATH, pkg["sub_package"]), output = _DISTRIBUTION_PATH, ) repository_ctx.delete(file_name) def _setup_rocm_distro_dir_impl(repository_ctx, rocm_distro): """Downloads and sets up a ROCm distribution. Args: repository_ctx: The repository context. rocm_distro: ROCm distribution struct from rocm_redist. Returns: ROCm config struct from _get_rocm_config. """ repository_ctx.file("rocm/.index") for pkg in rocm_distro.packages: _download_package(repository_ctx, pkg) for entry in rocm_distro.required_softlinks: repository_ctx.symlink( "{}/{}".format(_DISTRIBUTION_PATH, entry.target), "{}/{}".format(_DISTRIBUTION_PATH, entry.link), ) bash_bin = get_bash_bin(repository_ctx) # Canonical path helper def _canonical_path(p): parts = [x for x in p.split("/") if x != ""] return "/".join(parts) if parts else "" return _get_rocm_config(repository_ctx, bash_bin, _canonical_path("{}/{}".format(_DISTRIBUTION_PATH, rocm_distro.rocm_root)), "") def verify_build_defines(params): """Verify all variables that crosstool/BUILD.rocm.tpl expects are substituted. Args: params: dict of variables that will be passed to the BUILD.tpl template. """ missing = [] for param in [ "host_compiler_path", "unfiltered_compile_flags", ]: if ("%{" + param + "}") not in params: missing.append(param) if missing: auto_configure_fail( "Missing template parameters: %s" % missing, ) def _enable_rocm(repository_ctx): enable_rocm = get_host_environ(repository_ctx, "TF_NEED_ROCM") if enable_rocm == "1": if get_cpu_value(repository_ctx) != "Linux": repository_ctx.report_progress("Auto-Configuration Warning: ROCm configure is only supported on Linux") return False return True return False def _amdgpu_targets(repository_ctx, rocm_toolkit_path, bash_bin): """Returns a list of strings representing AMDGPU targets.""" amdgpu_targets_str = get_host_environ(repository_ctx, _TF_ROCM_AMDGPU_TARGETS) if not amdgpu_targets_str: cmd = "%s/bin/rocm_agent_enumerator" % rocm_toolkit_path result = execute(repository_ctx, [bash_bin, "-c", cmd]) targets = [target for target in result.stdout.strip().split("\n") if target != "gfx000"] targets = {x: None for x in targets} targets = list(targets.keys()) amdgpu_targets_str = ",".join(targets) amdgpu_targets = [amdgpu for amdgpu in amdgpu_targets_str.split(",") if amdgpu] for amdgpu_target in amdgpu_targets: if amdgpu_target[:3] != "gfx": auto_configure_fail("Invalid AMDGPU target: %s" % amdgpu_target) return amdgpu_targets def _lib_name(lib, version = "", static = False): """Constructs the name of a library on Linux. Args: lib: The name of the library, such as "hip" 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 static: return "lib%s.a" % lib else: if version: version = ".%s" % version return "lib%s.so%s" % (lib, version) def find_rocm_config(repository_ctx, rocm_path): """Returns ROCm config dictionary from running find_rocm_config.py Args: repository_ctx: The repository context. rocm_path: The path to the ROCm installation. Returns: A dictionary containing the ROCm configuration. """ python_bin = get_python_bin(repository_ctx) exec_result = execute(repository_ctx, [python_bin, repository_ctx.attr._find_rocm_config], env_vars = {"ROCM_PATH": rocm_path}) if exec_result.return_code: auto_configure_fail("Failed to run find_rocm_config.py: %s" % err_out(exec_result)) # Parse the dict from stdout. return dict([tuple(x.split(": ")) for x in exec_result.stdout.splitlines()]) def _get_rocm_config(repository_ctx, bash_bin, rocm_path, install_path): """Detects and returns information about the ROCm installation on the system. Args: repository_ctx: The repository context. bash_bin: the path to the path interpreter rocm_path: The path to the ROCm installation. install_path: The install path. 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. rocm_version_number: The version of ROCm on the system. miopen_version_number: The version of MIOpen on the system. hipruntime_version_number: The version of HIP Runtime on the system. """ config = find_rocm_config(repository_ctx, rocm_path) rocm_toolkit_path = config["rocm_toolkit_path"] rocm_version_number = config["rocm_version_number"] miopen_version_number = config["miopen_version_number"] hipruntime_version_number = config["hipruntime_version_number"] return struct( amdgpu_targets = _amdgpu_targets(repository_ctx, rocm_toolkit_path, bash_bin), rocm_toolkit_path = rocm_toolkit_path, rocm_version_number = rocm_version_number, miopen_version_number = miopen_version_number, hipruntime_version_number = hipruntime_version_number, install_path = install_path, ) def _tpl_path(repository_ctx, labelname): return repository_ctx.path(Label("//third_party/gpus/%s.tpl" % labelname)) def _tpl(repository_ctx, tpl, substitutions = {}, out = None): if not out: out = tpl.replace(":", "/") repository_ctx.template( out, _tpl_path(repository_ctx, tpl), substitutions, ) _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): # Set up BUILD file for rocm/. _tpl( repository_ctx, "rocm:build_defs.bzl", { "%{rocm_is_configured}": "False", "%{gpu_is_configured}": "if_true" if enable_cuda(repository_ctx) or enable_sycl(repository_ctx) else "if_false", "%{cuda_or_rocm}": "if_true" if enable_cuda(repository_ctx) else "if_false", "%{rocm_gpu_architectures}": "[]", "%{rocm_version_number}": "0", "%{single_gpu_rbe_pool}": repository_ctx.os.environ.get(_TF_ROCM_RBE_SINGLE_GPU_POOL, _DEFAULT_TF_ROCM_RBE_SINGLE_GPU_POOL), "%{multi_gpu_rbe_pool}": repository_ctx.os.environ.get(_TF_ROCM_RBE_MULTI_GPU_POOL, _DEFAULT_TF_ROCM_RBE_MULTI_GPU_POOL), }, ) _tpl( repository_ctx, "rocm:BUILD", { "%{hip_lib}": _lib_name("hip"), "%{rocblas_lib}": _lib_name("rocblas"), "%{hipblas_lib}": _lib_name("hipblas"), "%{miopen_lib}": _lib_name("miopen"), "%{rccl_lib}": _lib_name("rccl"), "%{hiprand_lib}": _lib_name("hiprand"), "%{hipsparse_lib}": _lib_name("hipsparse"), "%{roctracer_lib}": _lib_name("roctracer64"), "%{rocsolver_lib}": _lib_name("rocsolver"), "%{hipsolver_lib}": _lib_name("hipsolver"), "%{hipblaslt_lib}": _lib_name("hipblaslt"), "%{rocm_headers}": "", }, ) # Create dummy files for the ROCm toolkit since they are still required by # tensorflow/compiler/xla/stream_executor/rocm:rocm_rpath repository_ctx.file("rocm/hip/include/hip/hip_runtime.h", "") # Set up rocm_config.h, which is used by # tensorflow/compiler/xla/stream_executor/dso_loader.cc. _tpl( repository_ctx, "rocm:rocm_config.h", { "%{rocm_toolkit_path}": "/opt/rocm", }, "rocm/rocm_config/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 _remove_root_dir(path, root_dir): if path.startswith(root_dir + "/"): return path[len(root_dir) + 1:] return path def _setup_rocm_distro_dir(repository_ctx): """Sets up the rocm hermetic installation directory to be used in hermetic build""" bash_bin = get_bash_bin(repository_ctx) # Check if ROCM_PATH is set (highest priority) - symlink instead of download rocm_path = repository_ctx.os.environ.get("ROCM_PATH") if rocm_path: repository_ctx.report_progress("Using ROCm from ROCM_PATH: {}".format(rocm_path)) repository_ctx.file("rocm/.index") # Symlink the ROCM_PATH to rocm_dist repository_ctx.symlink(rocm_path, _DISTRIBUTION_PATH) return _get_rocm_config(repository_ctx, bash_bin, _DISTRIBUTION_PATH, rocm_path) # Check for multiple paths support (second priority) multiple_paths = repository_ctx.os.environ.get(_TF_ROCM_MULTIPLE_PATHS) if multiple_paths: repository_ctx.file("rocm/.index") paths_list = multiple_paths.split(":") for rocm_custom_path in paths_list: cmd = "find " + rocm_custom_path + "/* \\( -type f -o -type l \\)" result = execute(repository_ctx, [bash_bin, "-c", cmd]) result_files = result.stdout.strip().split("\n") if result.stdout.strip() else [] for file_path in result_files: relative_path = file_path[len(rocm_custom_path):] symlink_path = _DISTRIBUTION_PATH + relative_path if files_exist(repository_ctx, [symlink_path], bash_bin)[0]: fail("File already present: " + relative_path) else: repository_ctx.symlink(file_path, symlink_path) llvm_path = repository_ctx.os.environ.get(_LLVM_PATH) if llvm_path: repository_ctx.symlink(llvm_path, _DISTRIBUTION_PATH + "/llvm") repository_ctx.symlink(llvm_path, _DISTRIBUTION_PATH + "/lib/llvm") repository_ctx.symlink(llvm_path + "/amdgcn", _DISTRIBUTION_PATH + "/amdgcn") repository_ctx.report_progress("Using ROCm from multiple paths: {}".format(multiple_paths)) return _get_rocm_config(repository_ctx, bash_bin, _DISTRIBUTION_PATH, _DISTRIBUTION_PATH) # Check for custom URL-based distro (third priority) rocm_distro_url = repository_ctx.os.environ.get(_ROCM_DISTRO_URL) if rocm_distro_url: rocm_distro_hash = repository_ctx.os.environ.get(_ROCM_DISTRO_HASH) if not rocm_distro_hash: fail("{} environment variable is required".format(_ROCM_DISTRO_HASH)) rocm_distro_links = repository_ctx.os.environ.get(_ROCM_DISTRO_LINKS, "") rocm_distro = create_rocm_distro(rocm_distro_url, rocm_distro_hash, rocm_distro_links) return _setup_rocm_distro_dir_impl(repository_ctx, rocm_distro) # Check for hermetic redistributable or use default (lowest priority) rocm_distro_version = repository_ctx.os.environ.get(_ROCM_DISTRO_VERSION, _DEFAULT_ROCM_DISTRO_VERSION) if rocm_distro_version not in rocm_redist: fail("Unknown ROCM_DISTRO_VERSION: {}. Available versions: {}".format( rocm_distro_version, ", ".join(rocm_redist.keys()), )) repository_ctx.report_progress("Downloading hermetic ROCm distribution: {}".format(rocm_distro_version)) return _setup_rocm_distro_dir_impl(repository_ctx, rocm_redist[rocm_distro_version]) def _create_local_rocm_repository(repository_ctx): """Creates the repository containing files set up to build with ROCm.""" tpl_paths = {labelname: _tpl_path(repository_ctx, labelname) for labelname in [ "rocm:build_defs.bzl", "rocm:BUILD", "crosstool:BUILD.rocm", "crosstool:hipcc_cc_toolchain_config.bzl", "crosstool:clang/bin/crosstool_wrapper_driver_rocm", "rocm:rocm_config.h", ]} rocm_config = _setup_rocm_distro_dir(repository_ctx) rocm_version_number = int(rocm_config.rocm_version_number) # Copy header and library files to execroot. # rocm_toolkit_path rocm_toolkit_path = _remove_root_dir(rocm_config.rocm_toolkit_path, "rocm") bash_bin = get_bash_bin(repository_ctx) # Set up BUILD file for rocm/ repository_ctx.template( "rocm/build_defs.bzl", tpl_paths["rocm:build_defs.bzl"], { "%{rocm_is_configured}": "True", "%{gpu_is_configured}": "if_true", "%{cuda_or_rocm}": "if_true", "%{single_gpu_rbe_pool}": repository_ctx.os.environ.get(_TF_ROCM_RBE_SINGLE_GPU_POOL, _DEFAULT_TF_ROCM_RBE_SINGLE_GPU_POOL), "%{multi_gpu_rbe_pool}": repository_ctx.os.environ.get(_TF_ROCM_RBE_MULTI_GPU_POOL, _DEFAULT_TF_ROCM_RBE_MULTI_GPU_POOL), "%{rocm_gpu_architectures}": str(rocm_config.amdgpu_targets), "%{rocm_version_number}": str(rocm_version_number), }, ) repository_dict = { "%{rocm_root}": rocm_toolkit_path, "%{rocm_toolkit_path}": str(repository_ctx.path(rocm_config.rocm_toolkit_path)), "%{rocm_rbe_docker_image}": repository_ctx.os.environ.get(_TF_ROCM_RBE_DOCKER_IMAGE, _DEFAULT_TF_ROCM_RBE_DOCKER_IMAGE), "%{rocm_rbe_pool}": repository_ctx.os.environ.get(_TF_ROCM_RBE_POOL, _DEFAULT_TF_ROCM_RBE_POOL), "%{rocm_repo_name}": repository_ctx.name, } repository_ctx.template( "rocm/BUILD", tpl_paths["rocm:BUILD"], repository_dict, ) # Only expand template variables in the BUILD file repository_ctx.template( "crosstool/BUILD", tpl_paths["crosstool:BUILD.rocm"], ) # No templating of cc_toolchain_config - use attributes and templatize the # BUILD file. repository_ctx.template( "crosstool/cc_toolchain_config.bzl", tpl_paths["crosstool:hipcc_cc_toolchain_config.bzl"], ) repository_ctx.template( "crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc", tpl_paths["crosstool:clang/bin/crosstool_wrapper_driver_rocm"], { "%{tmpdir}": get_host_environ( repository_ctx, _TMPDIR, "", ), }, ) # Set up rocm_config.h, which is used by # tensorflow/compiler/xla/stream_executor/dso_loader.cc. repository_ctx.template( "rocm/rocm_config/rocm_config.h", tpl_paths["rocm:rocm_config.h"], { "%{rocm_toolkit_path}": rocm_config.install_path, "%{rocm_version_number}": rocm_config.rocm_version_number, "%{miopen_version_number}": rocm_config.miopen_version_number, "%{hipruntime_version_number}": rocm_config.hipruntime_version_number, }, ) # Set up rocm_config.h, which is used by # tensorflow/compiler/xla/stream_executor/dso_loader.cc. repository_ctx.template( "rocm/rocm_config_hermetic/rocm_config.h", tpl_paths["rocm:rocm_config.h"], { "%{rocm_toolkit_path}": str(repository_ctx.path(rocm_config.rocm_toolkit_path)), "%{rocm_version_number}": rocm_config.rocm_version_number, "%{miopen_version_number}": rocm_config.miopen_version_number, "%{hipruntime_version_number}": rocm_config.hipruntime_version_number, }, ) 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", "%{gpu_is_configured}": "if_true", "%{cuda_or_rocm}": "if_true", }, ) repository_ctx.template( "rocm/BUILD", config_repo_label(remote_config_repo, "rocm:BUILD"), {}, ) repository_ctx.template( "rocm/build_defs.bzl", config_repo_label(remote_config_repo, "rocm:build_defs.bzl"), {}, ) repository_ctx.template( "rocm/rocm/rocm_config.h", config_repo_label(remote_config_repo, "rocm:rocm/rocm_config.h"), {}, ) repository_ctx.template( "crosstool/BUILD", config_repo_label(remote_config_repo, "crosstool:BUILD"), {}, ) repository_ctx.template( "crosstool/cc_toolchain_config.bzl", config_repo_label(remote_config_repo, "crosstool:cc_toolchain_config.bzl"), {}, ) repository_ctx.template( "crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc", config_repo_label(remote_config_repo, "crosstool:clang/bin/crosstool_wrapper_driver_is_not_gcc"), {}, ) 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 get_host_environ(repository_ctx, _TF_ROCM_CONFIG_REPO) != None: _create_remote_rocm_repository( repository_ctx, get_host_environ(repository_ctx, _TF_ROCM_CONFIG_REPO), ) else: _create_local_rocm_repository(repository_ctx) _ENVIRONS = [ "TF_NEED_ROCM", "TF_NEED_CUDA", # Needed by the `if_gpu_is_configured` macro "ROCM_PATH", _TF_ROCM_AMDGPU_TARGETS, _TF_ROCM_MULTIPLE_PATHS, _LLVM_PATH, _TF_ROCM_RBE_DOCKER_IMAGE, _TF_ROCM_RBE_POOL, _TF_ROCM_RBE_SINGLE_GPU_POOL, _TF_ROCM_RBE_MULTI_GPU_POOL, _ROCM_DISTRO_VERSION, _ROCM_DISTRO_URL, _ROCM_DISTRO_HASH, _ROCM_DISTRO_LINKS, _TMPDIR, ] remote_rocm_configure = repository_rule( implementation = _create_local_rocm_repository, environ = _ENVIRONS, remotable = True, attrs = { "environ": attr.string_dict(), "_find_rocm_config": attr.label( default = Label("//third_party/gpus:find_rocm_config.py"), ), }, ) rocm_configure = repository_rule( implementation = _rocm_autoconf_impl, environ = _ENVIRONS + [_TF_ROCM_CONFIG_REPO], attrs = { "_find_rocm_config": attr.label( default = Label("//third_party/gpus:find_rocm_config.py"), ), }, )