tvm/contrib/rocm: improve finding of ld.lld

python/tvm/contrib/rocm.py

@@ -19,9 +19,43 @@
 from os.path import join, exists
 from . import util
 from .._ffi.base import py_str
+from .. import codegen
 from ..api import register_func, convert
 
-def rocm_link(in_file, out_file):
+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 = []
+    if hasattr(codegen, "llvm_version_major"):
+        major = codegen.llvm_version_major()
+        lld_list += ["ld.lld-%d.0" % major]
+        lld_list += ["ld.lld-%d" % major]
+    lld_list += ["lld"]
+    valid_list = [util.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
@@ -31,8 +65,12 @@ def rocm_link(in_file, out_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.
     """
-    args = ["ld.lld", "-shared", in_file, "-o", out_file]
+    args = [lld if lld is not None else find_lld()[0], "-shared", in_file, "-o", out_file]
     proc = subprocess.Popen(
         args,
         stdout=subprocess.PIPE,