1 files changed, 24 insertions, 2 deletions
diff --git a/clang/lib/Driver/Action.cpp b/clang/lib/Driver/Action.cpp
index 0117f8a..e9490e9 100644
--- a/clang/lib/Driver/Action.cpp
+++ b/clang/lib/Driver/Action.cpp
@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Driver/Action.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Regex.h"
 #include <cassert>
@@ -50,6 +51,24 @@ void BindArchAction::anchor() {}
 BindArchAction::BindArchAction(Action *Input, const char *_ArchName)
     : Action(BindArchClass, Input), ArchName(_ArchName) {}
 
+// Converts CUDA GPU architecture, e.g. "sm_21", to its corresponding virtual
+// compute arch, e.g. "compute_20".  Returns null if the input arch is null or
+// doesn't match an existing arch.
+static const char* GpuArchToComputeName(const char *ArchName) {
+  if (!ArchName)
+    return nullptr;
+  return llvm::StringSwitch<const char *>(ArchName)
+      .Cases("sm_20", "sm_21", "compute_20")
+      .Case("sm_30", "compute_30")
+      .Case("sm_32", "compute_32")
+      .Case("sm_35", "compute_35")
+      .Case("sm_37", "compute_37")
+      .Case("sm_50", "compute_50")
+      .Case("sm_52", "compute_52")
+      .Case("sm_53", "compute_53")
+      .Default(nullptr);
+}
+
 void CudaDeviceAction::anchor() {}
 
 CudaDeviceAction::CudaDeviceAction(Action *Input, const char *ArchName,
@@ -59,9 +78,12 @@ CudaDeviceAction::CudaDeviceAction(Action *Input, const char *ArchName,
   assert(IsValidGpuArchName(GpuArchName));
 }
 
+const char *CudaDeviceAction::getComputeArchName() const {
+  return GpuArchToComputeName(GpuArchName);
+}
+
 bool CudaDeviceAction::IsValidGpuArchName(llvm::StringRef ArchName) {
-  static llvm::Regex RE("^sm_[0-9]+$");
-  return RE.match(ArchName);
+  return GpuArchToComputeName(ArchName.data()) != nullptr;
 }
 
 void CudaHostAction::anchor() {}