From d700237f1aa1bc05d584a0f50fdad89370e17987 Mon Sep 17 00:00:00 2001
From: Artem Belevich <tra@google.com>
Date: Fri, 12 Jun 2020 15:28:02 -0700
Subject: [PATCH] [CUDA,HIP] Use VFS for SDK detection.

It's useful for using clang from tools that may need need to provide SDK files
from non-standard locations.

Clang CLI only provides a way to specify VFS for include files, so there's no
good way to test this yet.

Differential Revision: https://reviews.llvm.org/D81771
---
 clang/lib/Driver/ToolChains/AMDGPU.cpp | 11 ++++++-----
 clang/lib/Driver/ToolChains/Cuda.cpp   |  9 +++++----
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/clang/lib/Driver/ToolChains/AMDGPU.cpp b/clang/lib/Driver/ToolChains/AMDGPU.cpp
index 44a7e7fc3be0..adb659e3b229 100644
--- a/clang/lib/Driver/ToolChains/AMDGPU.cpp
+++ b/clang/lib/Driver/ToolChains/AMDGPU.cpp
@@ -27,7 +27,9 @@ void RocmInstallationDetector::scanLibDevicePath() {
   const StringRef Suffix(".bc");
 
   std::error_code EC;
-  for (llvm::sys::fs::directory_iterator LI(LibDevicePath, EC), LE;
+  for (llvm::vfs::directory_iterator
+           LI = D.getVFS().dir_begin(LibDevicePath, EC),
+           LE;
        !EC && LI != LE; LI = LI.increment(EC)) {
     StringRef FilePath = LI->path();
     StringRef FileName = llvm::sys::path::filename(FilePath);
@@ -137,11 +139,12 @@ RocmInstallationDetector::RocmInstallationDetector(
     LibDevicePath = LibPathEnv;
   }
 
+  auto &FS = D.getVFS();
   if (!LibDevicePath.empty()) {
     // Maintain compatability with HIP flag/envvar pointing directly at the
     // bitcode library directory. This points directly at the library path instead
     // of the rocm root installation.
-    if (!D.getVFS().exists(LibDevicePath))
+    if (!FS.exists(LibDevicePath))
       return;
 
     scanLibDevicePath();
@@ -151,7 +154,7 @@ RocmInstallationDetector::RocmInstallationDetector(
 
   for (const auto &Candidate : Candidates) {
     InstallPath = Candidate.Path;
-    if (InstallPath.empty() || !D.getVFS().exists(InstallPath))
+    if (InstallPath.empty() || !FS.exists(InstallPath))
       continue;
 
     // The install path situation in old versions of ROCm is a real mess, and
@@ -167,8 +170,6 @@ RocmInstallationDetector::RocmInstallationDetector(
     llvm::sys::path::append(IncludePath, InstallPath, "include");
     llvm::sys::path::append(LibDevicePath, InstallPath, "amdgcn", "bitcode");
 
-    auto &FS = D.getVFS();
-
     // We don't need the include path for OpenCL, since clang already ships with
     // the default header.
 
diff --git a/clang/lib/Driver/ToolChains/Cuda.cpp b/clang/lib/Driver/ToolChains/Cuda.cpp
index 55b205921e39..a25541cf31e9 100644
--- a/clang/lib/Driver/ToolChains/Cuda.cpp
+++ b/clang/lib/Driver/ToolChains/Cuda.cpp
@@ -78,6 +78,7 @@ CudaInstallationDetector::CudaInstallationDetector(
 
   // In decreasing order so we prefer newer versions to older versions.
   std::initializer_list<const char *> Versions = {"8.0", "7.5", "7.0"};
+  auto &FS = D.getVFS();
 
   if (Args.hasArg(clang::driver::options::OPT_cuda_path_EQ)) {
     Candidates.emplace_back(
@@ -114,7 +115,7 @@ CudaInstallationDetector::CudaInstallationDetector(
     for (const char *Ver : Versions)
       Candidates.emplace_back(D.SysRoot + "/usr/local/cuda-" + Ver);
 
-    Distro Dist(D.getVFS(), llvm::Triple(llvm::sys::getProcessTriple()));
+    Distro Dist(FS, llvm::Triple(llvm::sys::getProcessTriple()));
     if (Dist.IsDebian() || Dist.IsUbuntu())
       // Special case for Debian to have nvidia-cuda-toolkit work
       // out of the box. More info on http://bugs.debian.org/882505
@@ -125,14 +126,13 @@ CudaInstallationDetector::CudaInstallationDetector(
 
   for (const auto &Candidate : Candidates) {
     InstallPath = Candidate.Path;
-    if (InstallPath.empty() || !D.getVFS().exists(InstallPath))
+    if (InstallPath.empty() || !FS.exists(InstallPath))
       continue;
 
     BinPath = InstallPath + "/bin";
     IncludePath = InstallPath + "/include";
     LibDevicePath = InstallPath + "/nvvm/libdevice";
 
-    auto &FS = D.getVFS();
     if (!(FS.exists(IncludePath) && FS.exists(BinPath)))
       continue;
     bool CheckLibDevice = (!NoCudaLib || Candidate.StrictChecking);
@@ -177,7 +177,8 @@ CudaInstallationDetector::CudaInstallationDetector(
       }
     } else {
       std::error_code EC;
-      for (llvm::sys::fs::directory_iterator LI(LibDevicePath, EC), LE;
+      for (llvm::vfs::directory_iterator LI = FS.dir_begin(LibDevicePath, EC),
+                                         LE;
            !EC && LI != LE; LI = LI.increment(EC)) {
         StringRef FilePath = LI->path();
         StringRef FileName = llvm::sys::path::filename(FilePath);