[OpenMP] Add linking of OpenMP math library

[jhuber6]

Summary

This patch adds support for linking in the OpenMP math wrappers library.
The math library first replaces all math calls with an OpenMP wrapper
call. This wrapper call is linked early with a library that associates
the wrapper call to the original math function. This is necessary to
have access to the math function symbols without including the math
header which includes incompatible code with the GPU. These wrapper
functions have another library which maps them to the device library
version late when doing LTO. Linking libdevice must be done after this
library and done late when doing LTO. Unfortunately doing LTO with
libdevice is very slow right now.

Depends on D121467

[jdoerfert]

clang/include/clang/Driver/Options.td

1056

FWIW, I don't believe this amdgpu/nvpxt separation is actually improving things, on the contrary.

[JonChesterfield]

clang/include/clang/Driver/Options.td

1056

libomptarget_gpu? Definitely simplifies clang which is a win. I think it's an artefact of the previous device runtime layout.

[jhuber6]

clang/include/clang/Driver/Options.td

1056

The device library has different triples so we don't get errors when linking, not sure if it's possible to compile to generic IR that doesn't care about the triple. And I thought the same but wasn't too concerned with addressing it in the off chance someone wants to compile for amdgpu and nvptx64 at the same time and also specify the libraries.

clang/include/clang/Driver/Options.td

HelpText<"Method to generate ID's for compilation units for single source "
            "random numbers) | 'none' (disabled). Default is 'hash'. This option "
            "will be overridden by option '-cuid=[ID]' if it is specified." >;
 def libomptarget_amdgpu_bc_path_EQ : Joined<["--"], "libomptarget-amdgpu-bc-path=">, Group<i_Group>,
   HelpText<"Path to libomptarget-amdgcn bitcode library">;
 def libomptarget_amdgcn_bc_path_EQ : Joined<["--"], "libomptarget-amdgcn-bc-path=">, Group<i_Group>,
   HelpText<"Path to libomptarget-amdgcn bitcode library">, Alias<libomptarget_amdgpu_bc_path_EQ>;
 def libomptarget_nvptx_bc_path_EQ : Joined<["--"], "libomptarget-nvptx-bc-path=">, Group<i_Group>,
   HelpText<"Path to libomptarget-nvptx bitcode library">;
+def libomptarget_amdgpu_wrapper_bc_path_EQ : Joined<["--"], "libomptarget-amdgpu-wrapper-bc-path=">,
+  Group<i_Group>, HelpText<"Path to libomptarget-amdgpu math wrapper bitcode library">;
+def libomptarget_nvptx_wrapper_bc_path_EQ : Joined<["--"], "libomptarget-nvptx-wrapper-bc-path=">,
+  Group<i_Group>, HelpText<"Path to libomptarget-nvptx math wrapper bitcode library">;
+def libomptarget_amdgpu_math_bc_path_EQ : Joined<["--"], "libomptarget-amdgpu-math-bc-path=">,
+  Group<i_Group>, HelpText<"Path to libomptarget-amdgpu math bitcode library">;
+def libomptarget_nvptx_math_bc_path_EQ : Joined<["--"], "libomptarget-nvptx-math-bc-path=">,
+  Group<i_Group>, HelpText<"Path to libomptarget-nvptx math bitcode library">;
 def dD : Flag<["-"], "dD">, Group<d_Group>, Flags<[CC1Option]>,
   HelpText<"Print macro definitions in -E mode in addition to normal output">;
 def dI : Flag<["-"], "dI">, Group<d_Group>, Flags<[CC1Option]>,
   HelpText<"Print include directives in -E mode in addition to normal output">,
   MarshallingInfoFlag<PreprocessorOutputOpts<"ShowIncludeDirectives">>;
 def dM : Flag<["-"], "dM">, Group<d_Group>, Flags<[CC1Option]>,
   HelpText<"Print macro definitions in -E mode instead of normal output">;
 def dead__strip : Flag<["-"], "dead_strip">;

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

void AMDGPUOpenMPToolChain::addClangTargetOptions(
 
   CC1Args.push_back("-target-cpu");
   CC1Args.push_back(DriverArgs.MakeArgStringRef(GPUArch));
   CC1Args.push_back("-fcuda-is-device");
 
   if (DriverArgs.hasArg(options::OPT_nogpulib))
     return;
 
+  if (DriverArgs.hasArg(options::OPT_fopenmp_device_libm))
+    addOpenMPMathRTL(getDriver(), DriverArgs, CC1Args, getTriple());
+
   // Link the bitcode library late if we're using device LTO.
   if (getDriver().isUsingLTO(/* IsOffload */ true))
     return;
 
   addOpenMPDeviceRTL(getDriver(), DriverArgs, CC1Args, GPUArch, getTriple());
 }
 
 llvm::opt::DerivedArgList *AMDGPUOpenMPToolChain::TranslateArgs(

clang/lib/Driver/ToolChains/Clang.cpp

void LinkerWrapper::ConstructJob(Compilation &C, const JobAction &JA,
   // Pass the CUDA path to the linker wrapper tool.
   for (auto &I : llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
     const ToolChain *TC = I.second;
     if (TC->getTriple().isNVPTX()) {
       CudaInstallationDetector CudaInstallation(D, TheTriple, Args);
       if (CudaInstallation.isValid())
         CmdArgs.push_back(Args.MakeArgString(
             "--cuda-path=" + CudaInstallation.getInstallPath()));
-      break;
+
+      if (D.isUsingLTO(/* IsOffload */ true)) {
+        const ArgList &TCArgs =
+            C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
+        StringRef Arch = TCArgs.getLastArgValue(options::OPT_march_EQ);
+        std::string LibDeviceFile = CudaInstallation.getLibDeviceFile(Arch);
+      }
     }
   }
 
   // Get the AMDGPU math libraries.
   // FIXME: This method is bad, remove once AMDGPU has a proper math library
   // (see AMDGCN::OpenMPLinker::constructLLVMLinkCommand).
   for (auto &I : llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
     const ToolChain *TC = I.second;
for (auto &I :
       ArgStringList BitcodeLibrary;
       addOpenMPDeviceRTL(TCDriver, TCArgs, BitcodeLibrary, Arch,
                          TC->getTriple());
 
       if (!BitcodeLibrary.empty())
         CmdArgs.push_back(
             Args.MakeArgString("-target-library=" + TC->getTripleString() +
                                "-" + Arch + "=" + BitcodeLibrary.back()));
+
+      ArgStringList MathLibrary;
+      addOpenMPMathRTL(TCDriver, TCArgs, MathLibrary, TC->getTriple(), true);
+
+      if (!MathLibrary.empty())
+        CmdArgs.push_back(
+            Args.MakeArgString("-target-library=" + TC->getTripleString() +
+                               "-" + Arch + "=" + MathLibrary.back()));
+
+      CudaInstallationDetector CudaInstallation(D, TheTriple, Args);
+      std::string LibDeviceFile = CudaInstallation.getLibDeviceFile(Arch);
+      if (!LibDeviceFile.empty())
+        CmdArgs.push_back(
+            Args.MakeArgString("-target-library=" + TC->getTripleString() +
+                               "-" + Arch + "=" + LibDeviceFile));
     }
 
     // Pass in the optimization level to use for LTO.
     if (const Arg *A = Args.getLastArg(options::OPT_O_Group)) {
       StringRef OOpt;
       if (A->getOption().matches(options::OPT_O4) ||
           A->getOption().matches(options::OPT_Ofast))
         OOpt = "3";

clang/lib/Driver/ToolChains/CommonArgs.h

 
 bool haveAMDGPUCodeObjectVersionArgument(const Driver &D,
                                          const llvm::opt::ArgList &Args);
 
 void addMachineOutlinerArgs(const Driver &D, const llvm::opt::ArgList &Args,
                             llvm::opt::ArgStringList &CmdArgs,
                             const llvm::Triple &Triple, bool IsLTO);
 
+void addOpenMPMathRTL(const Driver &D, const llvm::opt::ArgList &DriverArgs,
+                      llvm::opt::ArgStringList &CC1Args,
+                      const llvm::Triple &Triple, bool IncludeLibm = false);
+
 void addOpenMPDeviceRTL(const Driver &D, const llvm::opt::ArgList &DriverArgs,
                         llvm::opt::ArgStringList &CC1Args,
                         StringRef BitcodeSuffix, const llvm::Triple &Triple);
 } // end namespace tools
 } // end namespace driver
 } // end namespace clang
 
 #endif // LLVM_CLANG_LIB_DRIVER_TOOLCHAINS_COMMONARGS_H

clang/lib/Driver/ToolChains/CommonArgs.cpp

if (A->getOption().matches(options::OPT_moutline)) {
       }
     } else {
       // Disable all outlining behaviour.
       addArg(Twine("-enable-machine-outliner=never"));
     }
   }
 }
 
+static void
+addBitcodeLibrary(const Driver &D, const llvm::opt::ArgList &DriverArgs,
+                  llvm::opt::ArgStringList &CC1Args, const llvm::Triple &Triple,
+                  ArrayRef<StringRef> LibraryPaths,
+                  StringRef BitcodeLibraryName, OptSpecifier BCPathOpt) {
+  // First check whether user specified the bc library.
+  if (const Arg *A = DriverArgs.getLastArg(BCPathOpt)) {
+    SmallString<128> LibOmpTargetFile(A->getValue());
+    if (llvm::sys::fs::exists(LibOmpTargetFile) &&
+        llvm::sys::fs::is_directory(LibOmpTargetFile)) {
+      llvm::sys::path::append(LibOmpTargetFile, BitcodeLibraryName);
+    }
+
+    if (llvm::sys::fs::exists(LibOmpTargetFile)) {
+      CC1Args.push_back("-mlink-builtin-bitcode");
+      CC1Args.push_back(DriverArgs.MakeArgString(LibOmpTargetFile));
+    } else {
+      D.Diag(diag::err_drv_omp_offload_target_bcruntime_not_found)
+          << LibOmpTargetFile;
+    }
+    return;
+  }
+
+  bool FoundBCLibrary = false;
+
+  for (StringRef LibraryPath : LibraryPaths) {
+    SmallString<128> LibOmpTargetFile(LibraryPath);
+    llvm::sys::path::append(LibOmpTargetFile, BitcodeLibraryName);
+    if (llvm::sys::fs::exists(LibOmpTargetFile)) {
+      CC1Args.push_back("-mlink-builtin-bitcode");
+      CC1Args.push_back(DriverArgs.MakeArgString(LibOmpTargetFile));
+      FoundBCLibrary = true;
+      break;
+    }
+  }
+
+  if (!FoundBCLibrary)
+    D.Diag(diag::err_drv_omp_offload_target_missingbcruntime)
+        << BitcodeLibraryName << (Triple.isAMDGCN() ? "amdgpu" : "nvptx");
+}
+
+void tools::addOpenMPMathRTL(const Driver &D,
+                             const llvm::opt::ArgList &DriverArgs,
+                             llvm::opt::ArgStringList &CC1Args,
+                             const llvm::Triple &Triple, bool IncludeLibm) {
+  SmallVector<StringRef, 8> LibraryPaths;
+
+  // Add path to clang lib / lib64 folder.
+  SmallString<256> DefaultLibPath = llvm::sys::path::parent_path(D.Dir);
+  llvm::sys::path::append(DefaultLibPath, Twine("lib") + CLANG_LIBDIR_SUFFIX);
+  LibraryPaths.emplace_back(DefaultLibPath.c_str());
+
+  // Add user defined library paths from LIBRARY_PATH.
+  llvm::Optional<std::string> LibPath =
+      llvm::sys::Process::GetEnv("LIBRARY_PATH");
+  if (LibPath) {
+    SmallVector<StringRef, 8> Frags;
+    const char EnvPathSeparatorStr[] = {llvm::sys::EnvPathSeparator, '\0'};
+    llvm::SplitString(*LibPath, Frags, EnvPathSeparatorStr);
+    for (StringRef Path : Frags)
+      LibraryPaths.emplace_back(Path.trim());
+  }
+
+  StringRef ArchPrefix = Triple.isAMDGCN() ? "amdgpu" : "nvptx";
+
+  OptSpecifier MathWrapperBCPathOpt =
+      Triple.isAMDGCN() ? options::OPT_libomptarget_amdgpu_wrapper_bc_path_EQ
+                        : options::OPT_libomptarget_nvptx_wrapper_bc_path_EQ;
+  std::string MathWrapperName =
+      ("libomptarget-" + ArchPrefix + "-math-wrappers.bc").str();
+
+  addBitcodeLibrary(D, DriverArgs, CC1Args, Triple, LibraryPaths,
+                    MathWrapperName, MathWrapperBCPathOpt);
+
+  // If we are doing LTO only link the OpenMP math wrappers.
+  if (D.isUsingLTO(/* IsOffload */ true) && !IncludeLibm)
+    return;
+
+  OptSpecifier MathBCPathOpt =
+      Triple.isAMDGCN() ? options::OPT_libomptarget_amdgpu_math_bc_path_EQ
+                        : options::OPT_libomptarget_nvptx_math_bc_path_EQ;
+  std::string DeviceMathName =
+      ("libomptarget-" + ArchPrefix + "-libm.bc").str();
+
+  addBitcodeLibrary(D, DriverArgs, CC1Args, Triple, LibraryPaths,
+                    DeviceMathName, MathBCPathOpt);
+}
+
 void tools::addOpenMPDeviceRTL(const Driver &D,
                                const llvm::opt::ArgList &DriverArgs,
                                llvm::opt::ArgStringList &CC1Args,
                                StringRef BitcodeSuffix,
                                const llvm::Triple &Triple) {
   SmallVector<StringRef, 8> LibraryPaths;
 
   // Add path to clang lib / lib64 folder.
void tools::addOpenMPDeviceRTL(const Driver &D,
   OptSpecifier LibomptargetBCPathOpt =
       Triple.isAMDGCN() ? options::OPT_libomptarget_amdgpu_bc_path_EQ
                         : options::OPT_libomptarget_nvptx_bc_path_EQ;
 
   StringRef ArchPrefix = Triple.isAMDGCN() ? "amdgpu" : "nvptx";
   std::string LibOmpTargetName =
       ("libomptarget-" + ArchPrefix + "-" + BitcodeSuffix + ".bc").str();
 
-  // First check whether user specifies bc library
-  if (const Arg *A = DriverArgs.getLastArg(LibomptargetBCPathOpt)) {
-    SmallString<128> LibOmpTargetFile(A->getValue());
-    if (llvm::sys::fs::exists(LibOmpTargetFile) &&
-        llvm::sys::fs::is_directory(LibOmpTargetFile)) {
-      llvm::sys::path::append(LibOmpTargetFile, LibOmpTargetName);
-    }
-
-    if (llvm::sys::fs::exists(LibOmpTargetFile)) {
-      CC1Args.push_back("-mlink-builtin-bitcode");
-      CC1Args.push_back(DriverArgs.MakeArgString(LibOmpTargetFile));
-    } else {
-      D.Diag(diag::err_drv_omp_offload_target_bcruntime_not_found)
-          << LibOmpTargetFile;
-    }
-  } else {
-    bool FoundBCLibrary = false;
-
-    for (StringRef LibraryPath : LibraryPaths) {
-      SmallString<128> LibOmpTargetFile(LibraryPath);
-      llvm::sys::path::append(LibOmpTargetFile, LibOmpTargetName);
-      if (llvm::sys::fs::exists(LibOmpTargetFile)) {
-        CC1Args.push_back("-mlink-builtin-bitcode");
-        CC1Args.push_back(DriverArgs.MakeArgString(LibOmpTargetFile));
-        FoundBCLibrary = true;
-        break;
-      }
-    }
-
-    if (!FoundBCLibrary)
-      D.Diag(diag::err_drv_omp_offload_target_missingbcruntime)
-          << LibOmpTargetName << ArchPrefix;
-  }
+  addBitcodeLibrary(D, DriverArgs, CC1Args, Triple, LibraryPaths,
+                    LibOmpTargetName, LibomptargetBCPathOpt);
 }

clang/lib/Driver/ToolChains/Cuda.cpp

if (DeviceOffloadingKind == Action::OFK_OpenMP &&
     return;
 
   std::string LibDeviceFile = CudaInstallation.getLibDeviceFile(GpuArch);
   if (LibDeviceFile.empty()) {
     getDriver().Diag(diag::err_drv_no_cuda_libdevice) << GpuArch;
     return;
   }
 
-  CC1Args.push_back("-mlink-builtin-bitcode");
-  CC1Args.push_back(DriverArgs.MakeArgString(LibDeviceFile));
-
   clang::CudaVersion CudaInstallationVersion = CudaInstallation.version();
 
   // New CUDA versions often introduce new instructions that are only supported
   // by new PTX version, so we need to raise PTX level to enable them in NVPTX
   // back-end.
   const char *PtxFeature = nullptr;
   switch (CudaInstallationVersion) {
 #define CASE_CUDA_VERSION(CUDA_VER, PTX_VER)                                   \
#undef CASE_CUDA_VERSION
   if (DeviceOffloadingKind == Action::OFK_OpenMP) {
     if (CudaInstallationVersion < CudaVersion::CUDA_92) {
       getDriver().Diag(
           diag::err_drv_omp_offload_target_cuda_version_not_support)
           << CudaVersionToString(CudaInstallationVersion);
       return;
     }
 
+    if (DriverArgs.hasArg(options::OPT_fopenmp_device_libm))
+      addOpenMPMathRTL(getDriver(), DriverArgs, CC1Args, getTriple());
+
     // Link the bitcode library late if we're using device LTO.
     if (getDriver().isUsingLTO(/* IsOffload */ true))
       return;
 
     addOpenMPDeviceRTL(getDriver(), DriverArgs, CC1Args, GpuArch.str(),
                        getTriple());
+
+    CC1Args.push_back("-mlink-builtin-bitcode");
+    CC1Args.push_back(DriverArgs.MakeArgString(LibDeviceFile));
+
     AddStaticDeviceLibsPostLinking(getDriver(), DriverArgs, CC1Args, "nvptx",
                                    GpuArch, /*isBitCodeSDL=*/true,
                                    /*postClangLink=*/true);
   }
 }
 
 llvm::DenormalMode CudaToolChain::getDefaultDenormalModeForType(
     const llvm::opt::ArgList &DriverArgs, const JobAction &JA,

clang/test/Driver/openmp-offload-gpu.c

 // NEW_DRIVER: "nvptx64-nvidia-cuda" - "NVPTX::Assembler", inputs: ["[[DEVICE_ASM]]"], output: "[[DEVICE_OBJ:.+]]" 
 
 // RUN:   %clang -### -fopenmp=libomp -fopenmp-targets=nvptx64-nvidia-cuda -Xopenmp-target=nvptx64-nvida-cuda -march=sm_70 \
 // RUN:          --libomptarget-nvptx-bc-path=%S/Inputs/libomptarget/libomptarget-new-nvptx-test.bc \
 // RUN:          -fopenmp-new-driver -no-canonical-prefixes -nogpulib %s -o openmp-offload-gpu 2>&1 \
 // RUN:   | FileCheck -check-prefix=NEW_DRIVER_EMBEDDING %s
 
 // NEW_DRIVER_EMBEDDING: -fembed-offload-object=[[CUBIN:.*\.cubin]],nvptx64-nvidia-cuda.sm_70
+
+// RUN:   %clang -### -fopenmp=libomp -fopenmp-targets=nvptx64-nvidia-cuda -Xopenmp-target=nvptx64-nvida-cuda -march=sm_70 \
+// RUN:          --libomptarget-nvptx-wrapper-bc-path=%S/Inputs/libomptarget/libomptarget-nvptx-wrapper.bc \
+// RUN:          --libomptarget-nvptx-math-bc-path=%S/Inputs/libomptarget/libomptarget-nvptx-math.bc \
+// RUN:          --libomptarget-nvptx-bc-path=%S/Inputs/libomptarget/libomptarget-new-nvptx-test.bc \
+// RUN:          -foffload-lto -fopenmp-device-libm -fopenmp-new-driver -no-canonical-prefixes %s -o openmp-offload-gpu 2>&1 \
+// RUN:   | FileCheck -check-prefix=DEVICE_LIBM %s
+
+// DEVICE_LIBM: clang{{.*}}"-mlink-builtin-bitcode" "{{.*}}/Inputs/libomptarget/libomptarget-nvptx-wrapper.bc"
+// DEVICE_LIBM: clang-linker-wrapper{{.*}}"-target-library=nvptx64-nvidia-cuda-sm_70={{.*}}/Inputs/libomptarget/libomptarget-new-nvptx-test.bc" "-target-library=nvptx64-nvidia-cuda-sm_70={{.*}}/Inputs/libomptarget/libomptarget-nvptx-math.bc"