# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Utility for ROCm backend""" import os import re import subprocess from os.path import exists, join import tvm_ffi import tvm.runtime import tvm.target from . import utils def find_lld(required=True): """Find ld.lld in system. Parameters ---------- required : bool Whether it is required, runtime error will be raised if the compiler is required. Returns ------- valid_list : list of str List of possible paths. Note ---- This function will first search ld.lld that matches the major llvm version that built with tvm """ lld_list = [] major = tvm.target.codegen.llvm_version_major(allow_none=True) if major is not None: lld_list += [f"ld.lld-{major}.0"] lld_list += [f"ld.lld-{major}"] lld_list += ["ld.lld"] lld_list += [f"/opt/rocm/llvm/bin/{x}" for x in lld_list] valid_list = [utils.which(x) for x in lld_list] valid_list = [x for x in valid_list if x] if not valid_list and required: raise RuntimeError("cannot find ld.lld, candidates are: " + str(lld_list)) return valid_list def rocm_link(in_file, out_file, lld=None): """Link relocatable ELF object to shared ELF object using lld Parameters ---------- in_file : str Input file name (relocatable ELF object file) out_file : str Output file name (shared ELF object file) lld : str, optional The lld linker, if not specified, we will try to guess the matched clang version. """ # if our result has undefined symbols, it will fail to load # (hipModuleLoad/hipModuleLoadData), but with a somewhat opaque message # so we have ld.lld check this here. # If you get a complaint about missing symbols you might want to check the # list of bitcode files below. args = [ lld if lld is not None else find_lld()[0], "--no-undefined", "-shared", in_file, "-o", out_file, ] proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) (out, _) = proc.communicate() if proc.returncode != 0: msg = "Linking error using ld.lld:\n" msg += out.decode("utf-8", errors="replace") raise RuntimeError(msg) @tvm_ffi.register_global_func("tvm_callback_rocm_link") def callback_rocm_link(obj_bin): """Links object file generated from LLVM to HSA Code Object Parameters ---------- obj_bin : bytearray The object file Return ------ cobj_bin : bytearray The HSA Code Object """ tmp_dir = utils.tempdir() tmp_obj = tmp_dir.relpath("rocm_kernel.o") tmp_cobj = tmp_dir.relpath("rocm_kernel.co") with open(tmp_obj, "wb") as out_file: out_file.write(bytes(obj_bin)) rocm_link(tmp_obj, tmp_cobj) cobj_bin = bytearray(open(tmp_cobj, "rb").read()) return cobj_bin @tvm_ffi.register_global_func("tvm_callback_rocm_bitcode_path") def callback_rocm_bitcode_path(rocdl_dir=None): """Utility function to find ROCm device library bitcodes Parameters ---------- rocdl_dir : str The path to rocm library directory The default value is the standard location """ # seems link order matters. if rocdl_dir is None: rocm_path = find_rocm_path() amdgcn_path = f"{rocm_path}/amdgcn/bitcode/" if exists(amdgcn_path): rocdl_dir = amdgcn_path # starting with rocm 3.9 else: rocdl_dir = "/opt/rocm/lib/" # until rocm 3.8 bitcode_names = [ "oclc_daz_opt_on", "ocml", "irif", # this does not exist in rocm 3.9, drop eventually "oclc_correctly_rounded_sqrt_off", "oclc_correctly_rounded_sqrt_on", "oclc_daz_opt_off", "oclc_finite_only_off", "oclc_finite_only_on", # todo (t-vi): an alternative might be to scan for the "oclc_isa_version_803", "oclc_isa_version_900", # isa version files (if the linker throws out "oclc_isa_version_906", # the unneeded ones or we filter for the arch we need) "oclc_isa_version_1030", "oclc_unsafe_math_off", "oclc_unsafe_math_on", "oclc_wavefrontsize64_on", "oclc_abi_version_500", ] bitcode_files = [] for n in bitcode_names: p = join(rocdl_dir, n + ".bc") # rocm >= 3.9 if not exists(p): # rocm <= 3.8 p = join(rocdl_dir, n + ".amdgcn.bc") if exists(p): bitcode_files.append(p) elif "isa_version" not in n and n not in {"irif"}: raise RuntimeError("could not find bitcode " + n) return tvm.runtime.convert(bitcode_files) def parse_compute_version(compute_version): """Parse compute capability string to divide major and minor version Parameters ---------- compute_version : str compute capability of a GPU (e.g. "6.0") Returns ------- major : int major version number minor : int minor version number """ split_ver = compute_version.split(".") try: major = int(split_ver[0]) minor = int(split_ver[1]) return major, minor except (IndexError, ValueError) as err: # pylint: disable=raise-missing-from raise RuntimeError("Compute version parsing error: " + str(err)) def have_matrixcore(compute_version=None): """Either MatrixCore support is provided in the compute capability or not Parameters ---------- compute_version : str, optional compute capability of a GPU (e.g. "7.0"). Returns ------- have_matrixcore : bool True if MatrixCore support is provided, False otherwise """ if compute_version is None: if tvm.rocm(0).exist: compute_version = tvm.rocm(0).compute_version else: raise RuntimeError("No ROCm runtime found") major, _ = parse_compute_version(compute_version) # matrix core first introduced in 8.0 if major >= 8: return True return False @tvm_ffi.register_global_func("tvm_callback_rocm_get_arch") def get_rocm_arch(rocm_path=None): """Utility function to get the AMD GPU architecture Parameters ---------- rocm_path : str The path to rocm installation directory Returns ------- gpu_arch : str The AMD GPU architecture """ if rocm_path is None: try: rocm_path = find_rocm_path() except RuntimeError: rocm_path = None gpu_arch = "gfx900" # check if rocm is installed if rocm_path is None or not os.path.exists(rocm_path): print("ROCm not detected, using default gfx900") return gpu_arch try: # Execute rocminfo command rocminfo_output = subprocess.check_output([f"{rocm_path}/bin/rocminfo"]).decode("utf-8") # Use regex to match the "Name" field match = re.search(r"Name:\s+(gfx\d+[a-zA-Z]*)", rocminfo_output) if match: gpu_arch = match.group(1) return gpu_arch except subprocess.CalledProcessError: print( f"Unable to execute rocminfo command, \ please ensure ROCm is installed and you have an AMD GPU on your system.\ using default {gpu_arch}." ) return gpu_arch def find_rocm_path(): """Utility function to find ROCm path Returns ------- path : str Path to ROCm root. """ if "ROCM_PATH" in os.environ: return os.environ["ROCM_PATH"] cmd = ["which", "hipcc"] proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) (out, _) = proc.communicate() out = out.decode("utf-8", errors="replace").strip() if proc.returncode == 0: return os.path.realpath(os.path.join(out, "../..")) rocm_path = "/opt/rocm" if os.path.exists(os.path.join(rocm_path, "bin/hipcc")): return rocm_path raise RuntimeError("Cannot find ROCm path")