[AMD] Use lld library API for linking (#7548)

This commit replaces shell out to ROCm's lld by directly calling the lld
library API.

---------

Co-authored-by: Maksim Levental <[email protected]>

Author: PMylon

third_party/amd/CMakeLists.txt

@@ -3,8 +3,11 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
 add_subdirectory(include)
 add_subdirectory(lib)
 if(TRITON_BUILD_PYTHON_MODULE)
+  find_package(LLD REQUIRED CONFIG PATHS "${MLIR_DIR}/../lld" NO_DEFAULT_PATH)
+  include_directories(${LLD_INCLUDE_DIRS})
+  message(STATUS "Found LLD distro-package @ ${LLD_DIR} and LLD include dirs @ ${LLD_INCLUDE_DIRS}")
   add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms TritonAMDGPUDialectToLLVM)
-  target_link_libraries(TritonAMD PRIVATE Python3::Module pybind11::headers)
+  target_link_libraries(TritonAMD PRIVATE Python3::Module pybind11::headers lldCommon lldELF)
 endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)

third_party/amd/backend/compiler.py

@@ -7,7 +7,6 @@
 import hashlib
 import tempfile
 import re
-import subprocess
 import functools
 import warnings
 from pathlib import Path
@@ -176,26 +175,6 @@ def get_arg_specialization(arg, ty, **kwargs):
             ret += "S"
         return ret
 
-    @staticmethod
-    def path_to_rocm_lld():
-        # Check env path for ld.lld
-        lld_env_path = knobs.amd.lld_path
-        if lld_env_path is not None:
-            lld = Path(lld_env_path)
-            if lld.is_file():
-                return lld
-        # Check backend for ld.lld (used for pytorch wheels)
-        lld = Path(__file__).parent / "llvm/bin/ld.lld"
-        if lld.is_file():
-            return lld
-        lld = Path("/opt/rocm/llvm/bin/ld.lld")
-        if lld.is_file():
-            return lld
-        lld = Path("/usr/bin/ld.lld")
-        if lld.is_file():
-            return lld
-        raise Exception("ROCm linker /opt/rocm/llvm/bin/ld.lld not found. Set 'TRITON_HIP_LLD_PATH' to its path.")
-
     @staticmethod
     def make_ttir(mod, metadata, options):
         pm = ir.pass_manager(mod.context)
@@ -434,14 +413,12 @@ def make_hsaco(src, metadata, options):
         if knobs.compilation.enable_asan:
             target_features = '+xnack'
         hsaco = amd.assemble_amdgcn(src, options.arch, target_features)
-
-        rocm_path = HIPBackend.path_to_rocm_lld()
         with tempfile.NamedTemporaryFile() as tmp_out:
             with tempfile.NamedTemporaryFile() as tmp_in:
-                with open(tmp_in.name, 'wb') as fd_in:
+                with open(tmp_in.name, "wb") as fd_in:
                     fd_in.write(hsaco)
-                subprocess.check_call([rocm_path, '-flavor', 'gnu', '-shared', tmp_in.name, '-o', tmp_out.name])
-            with open(tmp_out.name, 'rb') as fd_out:
+                amd.link_hsaco(tmp_in.name, tmp_out.name)
+            with open(tmp_out.name, "rb") as fd_out:
                 ret = fd_out.read()
         return ret
 
@@ -457,5 +434,4 @@ def add_stages(self, stages, options, language):
 
     @functools.lru_cache()
     def hash(self):
-        version = subprocess.check_output([HIPBackend.path_to_rocm_lld(), "--version"], encoding='utf-8')
-        return f'{version}-{self.target}'
+        return f'{self.target}'

third_party/amd/python/triton_amd.cc

@@ -2,6 +2,7 @@
 #include "TritonAMDGPUToLLVM/Passes.h"
 #include "TritonAMDGPUToLLVM/TargetUtils.h"
 #include "TritonAMDGPUTransforms/Passes.h"
+#include "lld/Common/Driver.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "passes.h"
@@ -27,6 +28,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/TargetParser/TargetParser.h"
+#include <array>
 #include <pybind11/pybind11.h>
 #include <stdexcept>
 
@@ -110,8 +112,31 @@ void addControlConstant(llvm::Module *module, const char *name,
   constant->setUnnamedAddr(GlobalVariable::UnnamedAddr::Local);
   constant->setVisibility(GlobalVariable::VisibilityTypes::ProtectedVisibility);
 }
+
 } // namespace
 
+LLD_HAS_DRIVER(elf)
+
+static std::optional<std::string> lldInvoke(const char *inPath,
+                                            const char *outPath) {
+  // Workaround: Disable parallelism to avoid hangs caused by LLVM's thread pool
+  // when the following code is executed in a forked child process.
+  // Context: lld::elf::LinkerDriver::link uses parallelFor which uses the
+  // LLVM's thread pool. During cleanup at ~TaskGroup() the child process hangs
+  // waiting.
+  std::array args{"ld.lld", "--threads=1", "-shared", inPath, "-o", outPath};
+  std::string errString;
+  llvm::raw_string_ostream errStream(errString);
+  auto lldRes = lld::lldMain(args, llvm::outs(), llvm::errs(),
+                             {{lld::Gnu, &lld::elf::link}});
+  bool noErrors = (!lldRes.retCode && lldRes.canRunAgain);
+  if (!noErrors) {
+    errStream.flush();
+    return errString;
+  }
+  return {};
+}
+
 void init_triton_amd(py::module &&m) {
   m.doc() = "Python bindings to the AMD Triton backend";
 
@@ -305,6 +330,14 @@ void init_triton_amd(py::module &&m) {
     }
   });
 
+  m.def("link_hsaco",
+        [](const std::string &inPath, const std::string &outPath) {
+          if (auto errString = lldInvoke(inPath.c_str(), outPath.c_str()))
+            throw std::runtime_error("LLD failed to link hsaco source " +
+                                     inPath + " into object file " + outPath +
+                                     " because " + errString.value());
+        });
+
   m.def("add_scalarize_packed_fops_llvm_pass", [](llvm::Function *fn) {
     mlir::triton::AMD::runScalarizePackedFOpsPass(*fn);
   });