[OpenMP] Add support for linking AMDGPU images

[jhuber6]

Summary

This patch adds support for linking AMDGPU images using the LLD binary.
AMDGPU files are always bitcode images and will always use the LTO
backend. Additionally we now pass the default architecture found with
the amdgpu-arch tool to the argument list.

Depends on D117156

[JonChesterfield]

A bit nervous about the automatic appending of host library search paths for the device link but I see that's pre-existing in nvptx. I'll check whether -### changes behaviour as expected and approve if it does

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

297

This part is valuable as-is and probably independently testable, in that I think -### will have no march= before this change and will contain march=gfx*** after it. Though for that test to work reliably, I think it would have to run on a machine where amdgpu-arch succeeded, which is probably more bother than it's worth. Will check locally.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

623

Does this mean to take the library paths the host used to link, and pass them unchanged for the device link as well? Doesn't seem a given that this is the right behaviour. Do we have a mechanism for passing arguments to the device linker? Might be ugly, e.g. -Xopenmp-target=sm_70 -Xopenmp-target=-Wl,whatever

[jhuber6]

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

297

Yes, I believe with this we could remove the other calls to checkSystemForAMDGPU and just reference the -march in the driver arguments. Not sure about testing because it requires the binary of course.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

623

That's a good idea for general control over the arguments to the linker wrapper. I don't think the -L arguments are strictly necessary here. Could probably move to your approach if needed, the -L arguments were already present in the clang-nvlink-wrapper so I just replicated that without really caring if it was necessary.

[jhuber6]

Updating after upstreaming a portion of this patch.

[jhuber6]

Update commits.

[tianshilei1992]

@JonChesterfield can we just accept the patch and make it in 14? If we later find anything broken, we could have bug fix.

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

llvm::opt::DerivedArgList *AMDGPUOpenMPToolChain::TranslateArgs(
     Action::OffloadKind DeviceOffloadKind) const {
   DerivedArgList *DAL =
       HostTC.TranslateArgs(Args, BoundArch, DeviceOffloadKind);
   if (!DAL)
     DAL = new DerivedArgList(Args.getBaseArgs());
 
   const OptTable &Opts = getDriver().getOpts();
 
-  if (DeviceOffloadKind != Action::OFK_OpenMP) {
-    for (Arg *A : Args) {
+  if (DeviceOffloadKind == Action::OFK_OpenMP) {
+    for (Arg *A : Args)
+      if (!llvm::is_contained(*DAL, A))
         DAL->append(A);
-    }
+
+    std::string Arch = DAL->getLastArgValue(options::OPT_march_EQ).str();
+    if (Arch.empty()) {
+      checkSystemForAMDGPU(Args, *this, Arch);
+      DAL->AddJoinedArg(nullptr, Opts.getOption(options::OPT_march_EQ), Arch);
+    }
+
+    return DAL;
+
+  for (Arg *A : Args) {
+    DAL->append(A);
   }
 
   if (!BoundArch.empty()) {
     DAL->eraseArg(options::OPT_march_EQ);
     DAL->AddJoinedArg(nullptr, Opts.getOption(options::OPT_march_EQ),
                       BoundArch);
   }
 

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

if (!PtxasPath)
     PtxasPath = sys::findProgramByName("ptxas");
   if (!PtxasPath)
     return createStringError(PtxasPath.getError(),
                              "Unable to find 'ptxas' in path");
 
   // Create a new file to write the linked device image to.
   SmallString<128> TempFile;
   if (std::error_code EC = sys::fs::createTemporaryFile(
-          TheTriple.getArchName() + "-" + Arch, "cubin", TempFile))
+          "lto-" + TheTriple.getArchName() + "-" + Arch, "cubin", TempFile))
     return createFileError(TempFile, EC);
   TempFiles.push_back(static_cast<std::string>(TempFile));
 
   // TODO: Pass in arguments like `-g` and `-v` from the driver.
   SmallVector<StringRef, 16> CmdArgs;
   std::string Opt = "-" + OptLevel;
   CmdArgs.push_back(*PtxasPath);
   CmdArgs.push_back(TheTriple.isArch64Bit() ? "-m64" : "-m32");
for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
   if (sys::ExecuteAndWait(*NvlinkPath, CmdArgs))
     return createStringError(inconvertibleErrorCode(), "'nvlink' failed");
 
   return static_cast<std::string>(TempFile);
 }
 } // namespace nvptx
+namespace amdgcn {
+Expected<std::string> link(ArrayRef<std::string> InputFiles,
+                           ArrayRef<std::string> LinkerArgs, Triple TheTriple,
+                           StringRef Arch) {
+  // AMDGPU uses the lld binary to link device object files.
+  ErrorOr<std::string> LLDPath =
+      sys::findProgramByName("lld", sys::path::parent_path(LinkerExecutable));
+  if (!LLDPath)
+    LLDPath = sys::findProgramByName("lld");
+  if (!LLDPath)
+    return createStringError(LLDPath.getError(),
+                             "Unable to find 'lld' in path");
+
+  // Create a new file to write the linked device image to.
+  SmallString<128> TempFile;
+  if (std::error_code EC = sys::fs::createTemporaryFile(
+          TheTriple.getArchName() + "-" + Arch + "-image", "out", TempFile))
+    return createFileError(TempFile, EC);
+  TempFiles.push_back(static_cast<std::string>(TempFile));
+
+  SmallVector<StringRef, 16> CmdArgs;
+  CmdArgs.push_back(*LLDPath);
+  CmdArgs.push_back("-flavor");
+  CmdArgs.push_back("gnu");
+  CmdArgs.push_back("--no-undefined");
+  CmdArgs.push_back("-shared");
+  CmdArgs.push_back("-o");
+  CmdArgs.push_back(TempFile);
+
+  // Copy system library paths used by the host linker.
+  for (StringRef Arg : LinkerArgs)
+    if (Arg.startswith("-L"))
+      CmdArgs.push_back(Arg);
+
+  // Add extracted input files.
+  for (StringRef Input : InputFiles)
+    CmdArgs.push_back(Input);
+
+  if (sys::ExecuteAndWait(*LLDPath, CmdArgs))
+    return createStringError(inconvertibleErrorCode(), "'lld' failed");
+
+  return static_cast<std::string>(TempFile);
+}
+} // namespace amdgcn
 
 Expected<std::string> linkDevice(ArrayRef<std::string> InputFiles,
                                  ArrayRef<std::string> LinkerArgs,
                                  Triple TheTriple, StringRef Arch) {
   switch (TheTriple.getArch()) {
   case Triple::nvptx:
   case Triple::nvptx64:
     return nvptx::link(InputFiles, LinkerArgs, TheTriple, Arch);
   case Triple::amdgcn:
-    // TODO: AMDGCN linking support.
+    return amdgcn::link(InputFiles, LinkerArgs, TheTriple, Arch);
   case Triple::x86:
   case Triple::x86_64:
     // TODO: x86 linking support.
   default:
     return createStringError(inconvertibleErrorCode(),
                              TheTriple.getArchName() +
                                  " linking is not supported");
   }