[PATCH] D98867: [HIP] Fix ROCm detection

[Yaxun Liu via Phabricator via cfe-commits]

yaxunl updated this revision to Diff 332104.
yaxunl marked 2 inline comments as done.
yaxunl added a comment.

revised by Artem's comments


CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D98867/new/

https://reviews.llvm.org/D98867

Files:
  clang/lib/Driver/ToolChains/AMDGPU.cpp
  clang/test/Driver/rocm-detect.hip


Index: clang/test/Driver/rocm-detect.hip
===================================================================
--- clang/test/Driver/rocm-detect.hip
+++ clang/test/Driver/rocm-detect.hip
@@ -21,6 +21,20 @@
 // RUN:   --rocm-path=%S/Inputs/rocm %s 2>&1 \
 // RUN:   | FileCheck -check-prefixes=COMMON,NODEFAULTLIBS %s
 
+// Test environment variable ROCM_PATH.
+// RUN: env ROCM_PATH=%S/Inputs/rocm %clang -### -target x86_64-linux-gnu \
+// RUN:   --print-rocm-search-dirs %s 2>&1 \
+// RUN:   | FileCheck -check-prefixes=ROCM-ENV %s
+
+// Test detecting latest /opt/rocm-{release} directory.
+// RUN: rm -rf %T/opt
+// RUN: mkdir -p %T/opt
+// RUN: cp -r %S/Inputs/rocm %T/opt/rocm-3.9.0-1234
+// RUN: cp -r %S/Inputs/rocm %T/opt/rocm-3.10.0
+// RUN: %clang -### -target x86_64-linux-gnu --sysroot=%T \
+// RUN:   --print-rocm-search-dirs %s 2>&1 \
+// RUN:   | FileCheck -check-prefixes=ROCM-REL %s
+
 // Test ROCm installation built by SPACK by invoke clang at %T/rocm-spack/llvm-amdgpu-*
 // directory through a soft link.
 
@@ -60,6 +74,11 @@
 
 // COMMON: "-triple" "amdgcn-amd-amdhsa"
 
+// ROCM-ENV: ROCm installation search path: {{.*}}/Inputs/rocm
+
+// ROCM-REL: ROCm installation search path: {{.*}}/opt/rocm
+// ROCM-REL: ROCm installation search path: {{.*}}/opt/rocm-3.10.0
+
 // SPACK: ROCm installation search path (Spack 4.0.0): [[DIR:.*]]
 // SPACK: ROCm installation search path: [[CLANG:.*]]
 // SPACK: ROCm installation search path: [[CLANG]]/lib/clang/{{[0-9.]+}}
Index: clang/lib/Driver/ToolChains/AMDGPU.cpp
===================================================================
--- clang/lib/Driver/ToolChains/AMDGPU.cpp
+++ clang/lib/Driver/ToolChains/AMDGPU.cpp
@@ -186,6 +186,12 @@
     ROCmSearchDirs.emplace_back(RocmPathArg.str());
     DoPrintROCmSearchDirs();
     return ROCmSearchDirs;
+  } else if (const char *RocmPathEnv = ::getenv("ROCM_PATH")) {
+    if (!StringRef(RocmPathEnv).empty()) {
+      ROCmSearchDirs.emplace_back(RocmPathEnv);
+      DoPrintROCmSearchDirs();
+      return ROCmSearchDirs;
+    }
   }
 
   // Try to find relative to the compiler binary.
@@ -247,6 +253,43 @@
 
   ROCmSearchDirs.emplace_back(D.SysRoot + "/opt/rocm",
                               /*StrictChecking=*/true);
+
+  // Find the latest /opt/rocm-{release} directory.
+  std::error_code EC;
+  std::string LatestROCm;
+  llvm::VersionTuple LatestVer;
+  // Get ROCm version from ROCm directory name.
+  auto GetROCmVersion = [](StringRef DirName) {
+    llvm::VersionTuple V;
+    std::string VerStr = DirName.drop_front(strlen("rocm-")).str();
+    // The ROCm directory name follows the format of
+    // rocm-{major}.{minor}.{subMinor}[-{build}]
+    std::replace(VerStr.begin(), VerStr.end(), '-', '.');
+    V.tryParse(VerStr);
+    return V;
+  };
+  for (llvm::vfs::directory_iterator
+           File = D.getVFS().dir_begin(D.SysRoot + "/opt", EC),
+           FileEnd;
+       File != FileEnd && !EC; File.increment(EC)) {
+    llvm::StringRef FileName = llvm::sys::path::filename(File->path());
+    if (!FileName.startswith("rocm-"))
+      continue;
+    if (LatestROCm.empty()) {
+      LatestROCm = FileName.str();
+      LatestVer = GetROCmVersion(LatestROCm);
+      continue;
+    }
+    auto Ver = GetROCmVersion(FileName);
+    if (LatestVer < Ver) {
+      LatestROCm = FileName.str();
+      LatestVer = Ver;
+    }
+  }
+  if (!LatestROCm.empty())
+    ROCmSearchDirs.emplace_back(D.SysRoot + "/opt/" + LatestROCm,
+                                /*StrictChecking=*/true);
+
   DoPrintROCmSearchDirs();
   return ROCmSearchDirs;
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D98867.332104.patch
Type: text/x-patch
Size: 3580 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/cfe-commits/attachments/20210320/0fea82b6/attachment-0001.bin>