From 94255511a6d59b14159544e2489905c62dab9fca Mon Sep 17 00:00:00 2001
From: Brad King <brad.king@kitware.com>
Date: Tue, 7 Mar 2017 16:21:29 -0500
Subject: VS: Select CUDA code generation architectures

Parse the `-gencode=`, `-arch`, and `-code` flags and generate a
`CodeGeneration` field in the project file.
---
 Source/cmVS10CudaFlagTable.h               | 22 +++++++++++++
 Source/cmVisualStudio10TargetGenerator.cxx |  2 ++
 Source/cmVisualStudioGeneratorOptions.cxx  | 51 ++++++++++++++++++++++++++++++
 Source/cmVisualStudioGeneratorOptions.h    |  2 ++
 4 files changed, 77 insertions(+)

diff --git a/Source/cmVS10CudaFlagTable.h b/Source/cmVS10CudaFlagTable.h
index 268553d..d8c27d7 100644
--- a/Source/cmVS10CudaFlagTable.h
+++ b/Source/cmVS10CudaFlagTable.h
@@ -14,5 +14,27 @@ static cmVS7FlagTable cmVS10CudaFlagTable[] = {
   { "CudaRuntime", "cudart", "CUDA runtime library", "",
     cmVS7FlagTable::UserFollowing },
 
+  // Capture arch/code arguments into temporaries for post-processing.
+  { "cmake-temp-gencode", "gencode=", "", "",
+    cmVS7FlagTable::UserValue | cmVS7FlagTable::SemicolonAppendable },
+  { "cmake-temp-gencode", "gencode", "", "",
+    cmVS7FlagTable::UserFollowing | cmVS7FlagTable::SemicolonAppendable },
+  { "cmake-temp-gencode", "-generate-code=", "", "",
+    cmVS7FlagTable::UserValue | cmVS7FlagTable::SemicolonAppendable },
+  { "cmake-temp-gencode", "-generate-code", "", "",
+    cmVS7FlagTable::UserFollowing | cmVS7FlagTable::SemicolonAppendable },
+
+  { "cmake-temp-code", "code=", "", "", cmVS7FlagTable::UserValue },
+  { "cmake-temp-code", "code", "", "", cmVS7FlagTable::UserFollowing },
+  { "cmake-temp-code", "-gpu-code=", "", "", cmVS7FlagTable::UserValue },
+  { "cmake-temp-code", "-gpu-code", "", "", cmVS7FlagTable::UserFollowing },
+
+  { "cmake-temp-arch", "arch=", "", "", cmVS7FlagTable::UserValue },
+  { "cmake-temp-arch", "arch", "", "", cmVS7FlagTable::UserFollowing },
+  { "cmake-temp-arch", "-gpu-architecture=", "", "",
+    cmVS7FlagTable::UserValue },
+  { "cmake-temp-arch", "-gpu-architecture", "", "",
+    cmVS7FlagTable::UserFollowing },
+
   { 0, 0, 0, 0, 0 }
 };
diff --git a/Source/cmVisualStudio10TargetGenerator.cxx b/Source/cmVisualStudio10TargetGenerator.cxx
index 8de9fe1..5ffa55c 100644
--- a/Source/cmVisualStudio10TargetGenerator.cxx
+++ b/Source/cmVisualStudio10TargetGenerator.cxx
@@ -2480,6 +2480,8 @@ bool cmVisualStudio10TargetGenerator::ComputeCudaOptions(
   // did not parse and hope it works.
   cudaOptions.RemoveFlag("AdditionalCompilerOptions");
 
+  cudaOptions.FixCudaCodeGeneration();
+
   std::vector<std::string> targetDefines;
   this->GeneratorTarget->GetCompileDefinitions(targetDefines,
                                                configName.c_str(), "CUDA");
diff --git a/Source/cmVisualStudioGeneratorOptions.cxx b/Source/cmVisualStudioGeneratorOptions.cxx
index c79dc11..1ca6b9c 100644
--- a/Source/cmVisualStudioGeneratorOptions.cxx
+++ b/Source/cmVisualStudioGeneratorOptions.cxx
@@ -208,6 +208,57 @@ cmVisualStudioGeneratorOptions::GetCudaRuntime() const
   return CudaRuntimeStatic;
 }
 
+void cmVisualStudioGeneratorOptions::FixCudaCodeGeneration()
+{
+  // Extract temporary values stored by our flag table.
+  FlagValue arch = this->TakeFlag("cmake-temp-arch");
+  FlagValue code = this->TakeFlag("cmake-temp-code");
+  FlagValue gencode = this->TakeFlag("cmake-temp-gencode");
+
+  // No -code allowed without -arch.
+  if (arch.empty()) {
+    code.clear();
+  }
+
+  if (arch.empty() && gencode.empty()) {
+    return;
+  }
+
+  // Create a CodeGeneration field with [arch],[code] syntax in each entry.
+  // CUDA will convert it to `-gencode=arch=[arch],code="[code],[arch]"`.
+  FlagValue& result = this->FlagMap["CodeGeneration"];
+
+  // First entries for the -arch=<arch> [-code=<code>,...] pair.
+  if (!arch.empty()) {
+    std::string arch_name = arch[0];
+    std::vector<std::string> codes;
+    if (!code.empty()) {
+      codes = cmSystemTools::tokenize(code[0], ",");
+    }
+    if (codes.empty()) {
+      codes.push_back(arch_name);
+      // nvcc -arch=<arch> has a special case that allows a real
+      // architecture to be specified instead of a virtual arch.
+      // It translates to -arch=<virtual> -code=<real>.
+      cmSystemTools::ReplaceString(arch_name, "sm_", "compute_");
+    }
+    for (std::vector<std::string>::iterator ci = codes.begin();
+         ci != codes.end(); ++ci) {
+      std::string entry = arch_name + "," + *ci;
+      result.push_back(entry);
+    }
+  }
+
+  // Now add entries for the -gencode=<arch>,<code> pairs.
+  for (std::vector<std::string>::iterator ei = gencode.begin();
+       ei != gencode.end(); ++ei) {
+    std::string entry = *ei;
+    cmSystemTools::ReplaceString(entry, "arch=", "");
+    cmSystemTools::ReplaceString(entry, "code=", "");
+    result.push_back(entry);
+  }
+}
+
 void cmVisualStudioGeneratorOptions::Parse(const char* flags)
 {
   // Parse the input string as a windows command line since the string
diff --git a/Source/cmVisualStudioGeneratorOptions.h b/Source/cmVisualStudioGeneratorOptions.h
index 6722503..52689e0 100644
--- a/Source/cmVisualStudioGeneratorOptions.h
+++ b/Source/cmVisualStudioGeneratorOptions.h
@@ -75,6 +75,8 @@ public:
   };
   CudaRuntime GetCudaRuntime() const;
 
+  void FixCudaCodeGeneration();
+
   bool IsDebug() const;
   bool IsWinRt() const;
   bool IsManaged() const;
-- 
cgit v0.12