From 4071763bb8816d36de7f2d6485b0afe9f716280b Mon Sep 17 00:00:00 2001
From: Artem Belevich <tra@google.com>
Date: Tue, 5 Feb 2019 22:38:58 +0000
Subject: [PATCH] Basic CUDA-10 support.

Differential Revision: https://reviews.llvm.org/D57771

llvm-svn: 353232
---
 clang/include/clang/Basic/Cuda.h                 |  5 ++++-
 clang/lib/Basic/Cuda.cpp                         |  9 ++++++++-
 clang/lib/CodeGen/CGCUDANV.cpp                   | 10 ++++++++++
 clang/lib/Driver/ToolChains/Cuda.cpp             |  2 ++
 clang/lib/Headers/__clang_cuda_runtime_wrapper.h |  2 +-
 5 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/clang/include/clang/Basic/Cuda.h b/clang/include/clang/Basic/Cuda.h
index 532c23098eef4..d96c7e0972f99 100644
--- a/clang/include/clang/Basic/Cuda.h
+++ b/clang/include/clang/Basic/Cuda.h
@@ -25,7 +25,8 @@ enum class CudaVersion {
   CUDA_91,
   CUDA_92,
   CUDA_100,
-  LATEST = CUDA_100,
+  CUDA_101,
+  LATEST = CUDA_101,
 };
 const char *CudaVersionToString(CudaVersion V);
 // Input is "Major.Minor"
@@ -107,6 +108,8 @@ CudaVersion MaxVersionForCudaArch(CudaArch A);
 enum class CudaFeature {
   // CUDA-9.2+ uses a new API for launching kernels.
   CUDA_USES_NEW_LAUNCH,
+  // CUDA-10.1+ needs explicit end of GPU binary registration.
+  CUDA_USES_FATBIN_REGISTER_END,
 };
 
 bool CudaFeatureEnabled(llvm::VersionTuple, CudaFeature);
diff --git a/clang/lib/Basic/Cuda.cpp b/clang/lib/Basic/Cuda.cpp
index 6ba9a8437e2c9..4bc8d8c5b42e5 100644
--- a/clang/lib/Basic/Cuda.cpp
+++ b/clang/lib/Basic/Cuda.cpp
@@ -25,6 +25,8 @@ const char *CudaVersionToString(CudaVersion V) {
     return "9.2";
   case CudaVersion::CUDA_100:
     return "10.0";
+  case CudaVersion::CUDA_101:
+    return "10.1";
   }
   llvm_unreachable("invalid enum");
 }
@@ -37,7 +39,8 @@ CudaVersion CudaStringToVersion(llvm::StringRef S) {
       .Case("9.0", CudaVersion::CUDA_90)
       .Case("9.1", CudaVersion::CUDA_91)
       .Case("9.2", CudaVersion::CUDA_92)
-      .Case("10.0", CudaVersion::CUDA_100);
+      .Case("10.0", CudaVersion::CUDA_100)
+      .Case("10.1", CudaVersion::CUDA_101);
 }
 
 const char *CudaArchToString(CudaArch A) {
@@ -352,6 +355,8 @@ static CudaVersion ToCudaVersion(llvm::VersionTuple Version) {
     return CudaVersion::CUDA_92;
   case 100:
     return CudaVersion::CUDA_100;
+  case 101:
+    return CudaVersion::CUDA_101;
   default:
     return CudaVersion::UNKNOWN;
   }
@@ -365,6 +370,8 @@ bool CudaFeatureEnabled(CudaVersion Version, CudaFeature Feature) {
   switch (Feature) {
   case CudaFeature::CUDA_USES_NEW_LAUNCH:
     return Version >= CudaVersion::CUDA_92;
+  case CudaFeature::CUDA_USES_FATBIN_REGISTER_END:
+    return Version >= CudaVersion::CUDA_101;
   }
   llvm_unreachable("Unknown CUDA feature.");
 }
diff --git a/clang/lib/CodeGen/CGCUDANV.cpp b/clang/lib/CodeGen/CGCUDANV.cpp
index 93ec94ca89bc3..68e83b939ae4f 100644
--- a/clang/lib/CodeGen/CGCUDANV.cpp
+++ b/clang/lib/CodeGen/CGCUDANV.cpp
@@ -616,6 +616,16 @@ llvm::Function *CGNVCUDARuntime::makeModuleCtorFunction() {
     // Call __cuda_register_globals(GpuBinaryHandle);
     if (RegisterGlobalsFunc)
       CtorBuilder.CreateCall(RegisterGlobalsFunc, RegisterFatbinCall);
+
+    // Call __cudaRegisterFatBinaryEnd(Handle) if this CUDA version needs it.
+    if (CudaFeatureEnabled(CGM.getTarget().getSDKVersion(),
+                           CudaFeature::CUDA_USES_FATBIN_REGISTER_END)) {
+      // void __cudaRegisterFatBinaryEnd(void **);
+      llvm::FunctionCallee RegisterFatbinEndFunc = CGM.CreateRuntimeFunction(
+          llvm::FunctionType::get(VoidTy, VoidPtrPtrTy, false),
+          "__cudaRegisterFatBinaryEnd");
+      CtorBuilder.CreateCall(RegisterFatbinEndFunc, RegisterFatbinCall);
+    }
   } else {
     // Generate a unique module ID.
     SmallString<64> ModuleID;
diff --git a/clang/lib/Driver/ToolChains/Cuda.cpp b/clang/lib/Driver/ToolChains/Cuda.cpp
index e3fe84b10181b..7d01f4b48d258 100644
--- a/clang/lib/Driver/ToolChains/Cuda.cpp
+++ b/clang/lib/Driver/ToolChains/Cuda.cpp
@@ -60,6 +60,8 @@ static CudaVersion ParseCudaVersionFile(llvm::StringRef V) {
     return CudaVersion::CUDA_92;
   if (Major == 10 && Minor == 0)
     return CudaVersion::CUDA_100;
+  if (Major == 10 && Minor == 1)
+    return CudaVersion::CUDA_101;
   return CudaVersion::UNKNOWN;
 }
 
diff --git a/clang/lib/Headers/__clang_cuda_runtime_wrapper.h b/clang/lib/Headers/__clang_cuda_runtime_wrapper.h
index 4b3672463e614..6b41e1741cb1a 100644
--- a/clang/lib/Headers/__clang_cuda_runtime_wrapper.h
+++ b/clang/lib/Headers/__clang_cuda_runtime_wrapper.h
@@ -62,7 +62,7 @@
 #include "cuda.h"
 #if !defined(CUDA_VERSION)
 #error "cuda.h did not define CUDA_VERSION"
-#elif CUDA_VERSION < 7000 || CUDA_VERSION > 10000
+#elif CUDA_VERSION < 7000 || CUDA_VERSION > 10010
 #error "Unsupported CUDA version!"
 #endif