land [Driver][ASan] Refactor Clang-Driver "Sanitizer Bitcode" linking. (llvm#123922)

Author: ampandey-1995

clang/lib/Driver/ToolChains/AMDGPU.cpp

@@ -721,9 +721,10 @@ void amdgpu::getAMDGPUTargetFeatures(const Driver &D,
                             options::OPT_m_amdgpu_Features_Group);
 }
 
-llvm::SmallVector<std::string, 12> amdgpu::dlr::getCommonDeviceLibNames(
-    const llvm::opt::ArgList &DriverArgs, const Driver &D,
-    const std::string &GPUArch, bool isOpenMP,
+llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12>
+amdgpu::dlr::getCommonDeviceLibNames(
+    const llvm::opt::ArgList &DriverArgs, const SanitizerArgs &SanArgs,
+    const Driver &D, const std::string &GPUArch, bool isOpenMP,
     const RocmInstallationDetector &RocmInstallation) {
   auto Kind = llvm::AMDGPU::parseArchAMDGCN(GPUArch);
   const StringRef CanonArch = llvm::AMDGPU::getArchNameAMDGCN(Kind);
@@ -732,13 +733,17 @@ llvm::SmallVector<std::string, 12> amdgpu::dlr::getCommonDeviceLibNames(
   auto ABIVer = DeviceLibABIVersion::fromCodeObjectVersion(
       getAMDGPUCodeObjectVersion(D, DriverArgs));
   bool noGPULib = DriverArgs.hasArg(options::OPT_nogpulib);
-  if (!RocmInstallation.checkCommonBitcodeLibs(CanonArch, LibDeviceFile,
-                                               ABIVer, noGPULib))
+  if (!RocmInstallation.checkCommonBitcodeLibs(CanonArch, LibDeviceFile, ABIVer,
+                                               noGPULib))
     return {};
 
   // If --hip-device-lib is not set, add the default bitcode libraries.
   // TODO: There are way too many flags that change this. Do we need to check
   // them all?
+  std::tuple<bool, const SanitizerArgs> GPUSan(
+      DriverArgs.hasFlag(options::OPT_fgpu_sanitize,
+                         options::OPT_fno_gpu_sanitize, true),
+      SanArgs);
   bool DAZ = DriverArgs.hasFlag(
       options::OPT_fgpu_flush_denormals_to_zero,
       options::OPT_fno_gpu_flush_denormals_to_zero,
@@ -757,7 +762,7 @@ llvm::SmallVector<std::string, 12> amdgpu::dlr::getCommonDeviceLibNames(
 
   return RocmInstallation.getCommonBitcodeLibs(
       DriverArgs, LibDeviceFile, Wave64, DAZ, FiniteOnly, UnsafeMathOpt,
-      FastRelaxedMath, CorrectSqrt, ABIVer, isOpenMP);
+      FastRelaxedMath, CorrectSqrt, ABIVer, GPUSan, isOpenMP);
 }
 
 /// AMDGPU Toolchain
@@ -1003,6 +1008,11 @@ void ROCMToolChain::addClangTargetOptions(
                                                 ABIVer, noGPULib))
     return;
 
+  std::tuple<bool, const SanitizerArgs> GPUSan(
+      DriverArgs.hasFlag(options::OPT_fgpu_sanitize,
+                         options::OPT_fno_gpu_sanitize, true),
+      getSanitizerArgs((DriverArgs)));
+
   bool Wave64 = isWave64(DriverArgs, Kind);
 
   // TODO: There are way too many flags that change this. Do we need to check
@@ -1018,21 +1028,19 @@ void ROCMToolChain::addClangTargetOptions(
       DriverArgs.hasArg(options::OPT_cl_fp32_correctly_rounded_divide_sqrt);
 
   // Add the OpenCL specific bitcode library.
-  llvm::SmallVector<std::string, 12> BCLibs;
-  BCLibs.push_back(RocmInstallation->getOpenCLPath().str());
+  llvm::SmallVector<BitCodeLibraryInfo, 12> BCLibs;
+  BCLibs.emplace_back(RocmInstallation->getOpenCLPath().str());
 
   // Add the generic set of libraries.
   BCLibs.append(RocmInstallation->getCommonBitcodeLibs(
       DriverArgs, LibDeviceFile, Wave64, DAZ, FiniteOnly, UnsafeMathOpt,
-      FastRelaxedMath, CorrectSqrt, ABIVer, false));
+      FastRelaxedMath, CorrectSqrt, ABIVer, GPUSan, false));
 
-  if (getSanitizerArgs(DriverArgs).needsAsanRt()) {
-    CC1Args.push_back("-mlink-bitcode-file");
-    CC1Args.push_back(
-        DriverArgs.MakeArgString(RocmInstallation->getAsanRTLPath()));
-  }
-  for (StringRef BCFile : BCLibs) {
-    CC1Args.push_back("-mlink-builtin-bitcode");
+  for (auto [BCFile, Internalize] : BCLibs) {
+    if (Internalize)
+      CC1Args.push_back("-mlink-builtin-bitcode");
+    else
+      CC1Args.push_back("-mlink-bitcode-file");
     CC1Args.push_back(DriverArgs.MakeArgString(BCFile));
   }
 }
@@ -1058,15 +1066,31 @@ bool RocmInstallationDetector::checkCommonBitcodeLibs(
   return true;
 }
 
-llvm::SmallVector<std::string, 12>
+llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12>
 RocmInstallationDetector::getCommonBitcodeLibs(
     const llvm::opt::ArgList &DriverArgs, StringRef LibDeviceFile, bool Wave64,
     bool DAZ, bool FiniteOnly, bool UnsafeMathOpt, bool FastRelaxedMath,
-    bool CorrectSqrt, DeviceLibABIVersion ABIVer, bool isOpenMP = false) const {
-  llvm::SmallVector<std::string, 12> BCLibs;
-
-  auto AddBCLib = [&](StringRef BCFile) { BCLibs.push_back(BCFile.str()); };
+    bool CorrectSqrt, DeviceLibABIVersion ABIVer,
+    const std::tuple<bool, const SanitizerArgs> &GPUSan,
+    bool isOpenMP = false) const {
+  llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12> BCLibs;
+
+  auto GPUSanEnabled = [GPUSan]() {
+    return std::get<bool>(GPUSan) &&
+           std::get<const SanitizerArgs>(GPUSan).needsAsanRt();
+  };
+  auto AddBCLib = [&](ToolChain::BitCodeLibraryInfo BCLib,
+                      bool Internalize = true) {
+    BCLib.ShouldInternalize = Internalize;
+    BCLibs.push_back(BCLib);
+  };
+  auto AddSanBCLibs = [&]() {
+    if (GPUSanEnabled()) {
+      AddBCLib(getAsanRTLPath(), false);
+    }
+  };
 
+  AddSanBCLibs();
   AddBCLib(getOCMLPath());
   // FIXME: OpenMP has ockl and ocml contained in libomptarget.bc. However,
   // we cannot exclude ocml here because of the crazy always-compile clang
@@ -1076,6 +1100,8 @@ RocmInstallationDetector::getCommonBitcodeLibs(
   // __BUILD_MATH_BUILTINS_LIB__ turning static libm functions to extern.
   if (!isOpenMP)
     AddBCLib(getOCKLPath());
+  else if (GPUSanEnabled() && isOpenMP)
+    AddBCLib(getOCKLPath(), false);
   AddBCLib(getDenormalsAreZeroPath(DAZ));
   AddBCLib(getUnsafeMathPath(UnsafeMathOpt || FastRelaxedMath));
   AddBCLib(getFiniteOnlyPath(FiniteOnly || FastRelaxedMath));
@@ -1096,14 +1122,15 @@ bool AMDGPUToolChain::shouldSkipArgument(const llvm::opt::Arg *A) const {
   return false;
 }
 
-llvm::SmallVector<std::string, 12>
+llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12>
 ROCMToolChain::getCommonDeviceLibNames(const llvm::opt::ArgList &DriverArgs,
                                        const std::string &GPUArch,
                                        bool isOpenMP) const {
   RocmInstallationDetector RocmInstallation(getDriver(), getTriple(),
                                             DriverArgs, true, true);
-  return amdgpu::dlr::getCommonDeviceLibNames(DriverArgs, getDriver(), GPUArch,
-                                              isOpenMP, RocmInstallation);
+  return amdgpu::dlr::getCommonDeviceLibNames(
+      DriverArgs, getSanitizerArgs(DriverArgs), getDriver(), GPUArch, isOpenMP,
+      RocmInstallation);
 }
 
 bool AMDGPUToolChain::shouldSkipSanitizeOption(
@@ -1117,7 +1144,7 @@ bool AMDGPUToolChain::shouldSkipSanitizeOption(
     return false;
 
   if (!DriverArgs.hasFlag(options::OPT_fgpu_sanitize,
-                          options::OPT_fno_gpu_sanitize, false))
+                          options::OPT_fno_gpu_sanitize, true))
     return true;
 
   auto &Diags = TC.getDriver().getDiags();

clang/lib/Driver/ToolChains/AMDGPU.h

@@ -43,8 +43,9 @@ void getAMDGPUTargetFeatures(const Driver &D, const llvm::Triple &Triple,
                              StringRef TcTargetID = StringRef());
 
 namespace dlr {
-llvm::SmallVector<std::string, 12>
-getCommonDeviceLibNames(const llvm::opt::ArgList &DriverArgs, const Driver &D,
+llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12>
+getCommonDeviceLibNames(const llvm::opt::ArgList &DriverArgs,
+                        const SanitizerArgs &SanArgs, const Driver &D,
                         const std::string &GPUArch, bool isOpenMP,
                         const RocmInstallationDetector &RocmInstallation);
 
@@ -188,7 +189,7 @@ class LLVM_LIBRARY_VISIBILITY ROCMToolChain : public AMDGPUToolChain {
                         Action::OffloadKind DeviceOffloadKind) const override;
 
   // Returns a list of device library names shared by different languages
-  llvm::SmallVector<std::string, 12>
+  llvm::SmallVector<BitCodeLibraryInfo, 12>
   getCommonDeviceLibNames(const llvm::opt::ArgList &DriverArgs,
                           const std::string &GPUArch,
                           bool isOpenMP = false) const;

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

@@ -10,8 +10,9 @@
 #include "AMDGPU.h"
 #include "CommonArgs.h"
 #include "HIPUtility.h"
-#include "ToolChains/ROCm.h"
+#include "ROCm.h"
 #include "clang/Basic/DiagnosticDriver.h"
+#include "clang/Config/config.h"
 #include "clang/Driver/Compilation.h"
 #include "clang/Driver/Driver.h"
 #include "clang/Driver/DriverDiagnostic.h"
@@ -20,14 +21,13 @@
 #include "clang/Driver/SanitizerArgs.h"
 #include "clang/Driver/Tool.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/StringExtras.h"    
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormatAdapters.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
 #include "llvm/TargetParser/TargetParser.h"
-#include "clang/Config/config.h"
 
 using namespace clang::driver;
 using namespace clang::driver::toolchains;
@@ -232,20 +232,7 @@ const char *amdgpu::dlr::getLinkCommandArgs(
       LibSuffix.append("/asan");
   }
 
-  llvm::SmallVector<std::string, 12> BCLibs;
-
-  std::string AsanRTL;
-  if (Args.hasFlag(options::OPT_fgpu_sanitize, options::OPT_fno_gpu_sanitize,
-                   true) &&
-      TC.getSanitizerArgs(Args).needsAsanRt()) {
-    if (!Args.hasArg(options::OPT_nogpulib)){
-      AsanRTL = RocmInstallation.getAsanRTLPath();
-      if(AsanRTL.empty())
-        TC.getDriver().Diag(diag::err_drv_no_asan_rt_lib);
-      else
-        BCLibs.push_back(AsanRTL);
-    }
-  }
+  llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12> BCLibs;
   StringRef GPUArch = getProcessorFromTargetID(Triple, TargetID);
 
   // When the base lib directory is called `lib` we enable
@@ -261,7 +248,7 @@ const char *amdgpu::dlr::getLinkCommandArgs(
       std::string EnvOmpLibDevice = EnvLibraryPath + LibDeviceName;
       if (llvm::sys::fs::exists(EnvOmpLibDevice)) {
         EnvOmpLibDeviceFound = true;
-        BCLibs.push_back(EnvOmpLibDevice);
+        BCLibs.emplace_back(EnvOmpLibDevice);
         break;
       }
     }
@@ -273,33 +260,23 @@ const char *amdgpu::dlr::getLinkCommandArgs(
       StringRef bc_file_lib =
           Args.MakeArgString(C.getDriver().Dir + "/../lib" + LibDeviceName);
       if (llvm::sys::fs::exists(bc_file_suf))
-        BCLibs.push_back(Args.MakeArgString(bc_file_suf));
+        BCLibs.emplace_back(Args.MakeArgString(bc_file_suf));
       else if (llvm::sys::fs::exists(bc_file_lib))
         // In case a LibSuffix version not found, use suffix "lib"
-        BCLibs.push_back(Args.MakeArgString(bc_file_lib));
+        BCLibs.emplace_back(Args.MakeArgString(bc_file_lib));
       else
         TC.getDriver().Diag(diag::err_drv_omp_offload_target_bcruntime_not_found)
           << "libomptarget-amdgpu.bc";
     }
 
-    if (!AsanRTL.empty()) {
-      // asanrtl is dependent on ockl so for every asanrtl bitcode linking
-      // requires ockl but viceversa is not true.
-      std::string OcklRTL(RocmInstallation.getOCKLPath());
-      if (OcklRTL.empty())
-        TC.getDriver().Diag(diag::err_drv_no_rocm_device_lib);
-      else
-        BCLibs.push_back(OcklRTL);
-    }
-
     // Add the generic set of libraries, OpenMP subset only
     BCLibs.append(amdgpu::dlr::getCommonDeviceLibNames(
-        C.getArgs(), C.getDriver(), GPUArch.str(), /* isOpenMP=*/true,
-        RocmInstallation));
+        C.getArgs(), TC.getSanitizerArgs(C.getArgs()), C.getDriver(),
+        GPUArch.str(), /* isOpenMP=*/true, RocmInstallation));
   }
 
-  llvm::for_each(BCLibs, [&](StringRef BCFile) {
-    LastLinkArgs.push_back(Args.MakeArgString(BCFile));
+  llvm::for_each(BCLibs, [&](auto BCLib) {
+    LastLinkArgs.push_back(Args.MakeArgString(BCLib.Path));
   });
 
   LastLinkArgs.push_back("-o");
@@ -481,7 +458,7 @@ llvm::opt::DerivedArgList *AMDGPUOpenMPToolChain::TranslateArgs(
 
   DerivedArgList *DAL =
       HostTC.TranslateArgs(Args, BoundArch, DeviceOffloadKind);
-  if (!DAL)
+  if (!DAL || Args.hasArg(options::OPT_fsanitize_EQ))
     DAL = new DerivedArgList(Args.getBaseArgs());
 
   const OptTable &Opts = getDriver().getOpts();
@@ -602,11 +579,6 @@ AMDGPUOpenMPToolChain::getDeviceLibs(const llvm::opt::ArgList &Args) const {
   if (Args.hasArg(options::OPT_nogpulib))
     return {};
 
-  if (!RocmInstallation->hasDeviceLibrary()) {
-    getDriver().Diag(diag::err_drv_no_rocm_device_lib) << 0;
-    return {};
-  }
-
   StringRef GpuArch = getProcessorFromTargetID(
       getTriple(), Args.getLastArgValue(options::OPT_march_EQ));
 

clang/lib/Driver/ToolChains/HIPAMD.cpp

@@ -387,7 +387,7 @@ HIPAMDToolChain::getDeviceLibs(const llvm::opt::ArgList &DriverArgs) const {
         llvm::sys::path::append(Path, BCName);
         FullName = Path;
         if (llvm::sys::fs::exists(FullName)) {
-          BCLibs.push_back(FullName);
+          BCLibs.emplace_back(FullName);
           return;
         }
       }
@@ -402,22 +402,11 @@ HIPAMDToolChain::getDeviceLibs(const llvm::opt::ArgList &DriverArgs) const {
     assert(!GpuArch.empty() && "Must have an explicit GPU arch.");
 
     // If --hip-device-lib is not set, add the default bitcode libraries.
-    if (DriverArgs.hasFlag(options::OPT_fgpu_sanitize,
-                           options::OPT_fno_gpu_sanitize, true) &&
-        getSanitizerArgs(DriverArgs).needsAsanRt()) {
-      auto AsanRTL = RocmInstallation->getAsanRTLPath();
-      if (AsanRTL.empty()) {
-        getDriver().Diag(diag::err_drv_no_asan_rt_lib);
-        return {};
-      } else
-        BCLibs.emplace_back(AsanRTL, /*ShouldInternalize=*/false);
-    }
-
     // Add the HIP specific bitcode library.
-    BCLibs.push_back(RocmInstallation->getHIPPath());
+    BCLibs.emplace_back(RocmInstallation->getHIPPath());
 
     // Add common device libraries like ocml etc.
-    for (StringRef N : getCommonDeviceLibNames(DriverArgs, GpuArch.str()))
+    for (auto N : getCommonDeviceLibNames(DriverArgs, GpuArch.str()))
       BCLibs.emplace_back(N);
 
     // Add instrument lib.

clang/lib/Driver/ToolChains/ROCm.h

@@ -13,6 +13,7 @@
 #include "clang/Basic/LLVM.h"
 #include "clang/Driver/Driver.h"
 #include "clang/Driver/Options.h"
+#include "clang/Driver/SanitizerArgs.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Option/ArgList.h"
@@ -173,12 +174,12 @@ class RocmInstallationDetector {
 
   /// Get file paths of default bitcode libraries common to AMDGPU based
   /// toolchains.
-  llvm::SmallVector<std::string, 12>
-  getCommonBitcodeLibs(const llvm::opt::ArgList &DriverArgs,
-                       StringRef LibDeviceFile, bool Wave64, bool DAZ,
-                       bool FiniteOnly, bool UnsafeMathOpt,
-                       bool FastRelaxedMath, bool CorrectSqrt,
-                       DeviceLibABIVersion ABIVer, bool isOpenMP) const;
+  llvm::SmallVector<ToolChain::BitCodeLibraryInfo, 12> getCommonBitcodeLibs(
+      const llvm::opt::ArgList &DriverArgs, StringRef LibDeviceFile,
+      bool Wave64, bool DAZ, bool FiniteOnly, bool UnsafeMathOpt,
+      bool FastRelaxedMath, bool CorrectSqrt, DeviceLibABIVersion ABIVer,
+      const std::tuple<bool, const SanitizerArgs> &GPUSan, bool isOpenMP) const;
+
   /// Check file paths of default bitcode libraries common to AMDGPU based
   /// toolchains. \returns false if there are invalid or missing files.
   bool checkCommonBitcodeLibs(StringRef GPUArch, StringRef LibDeviceFile,