6 files changed, 67 insertions, 6 deletions

diff --git a/include/clang/Basic/AddressSpaces.h b/include/clang/Basic/AddressSpaces.h
index d44a9c3b03..56c0f956f2 100644
--- a/include/clang/Basic/AddressSpaces.h
+++ b/include/clang/Basic/AddressSpaces.h
@@ -29,6 +29,10 @@ enum ID {
   opencl_local,
   opencl_constant,
 
+  cuda_device,
+  cuda_constant,
+  cuda_shared,
+
   Last,
   Count = Last-Offset
 };
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index a5028338cb..27057883e5 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -206,7 +206,10 @@ static const LangAS::Map *getAddressSpaceMap(const TargetInfo &T,
     static const unsigned FakeAddrSpaceMap[] = {
       1, // opencl_global
       2, // opencl_local
-      3  // opencl_constant
+      3, // opencl_constant
+      4, // cuda_device
+      5, // cuda_constant
+      6  // cuda_shared
     };
     return &FakeAddrSpaceMap;
   } else {
diff --git a/lib/Basic/Targets.cpp b/lib/Basic/Targets.cpp
index 045229e9fb..8e7cdf934d 100644
--- a/lib/Basic/Targets.cpp
+++ b/lib/Basic/Targets.cpp
@@ -949,7 +949,10 @@ namespace {
   static const unsigned PTXAddrSpaceMap[] = {
     0,    // opencl_global
     4,    // opencl_local
-    1     // opencl_constant
+    1,    // opencl_constant
+    0,    // cuda_device
+    1,    // cuda_constant
+    4,    // cuda_shared
   };
   class PTXTargetInfo : public TargetInfo {
     static const char * const GCCRegNames[];
@@ -3384,7 +3387,10 @@ namespace {
   static const unsigned TCEOpenCLAddrSpaceMap[] = {
       3, // opencl_global
       4, // opencl_local
-      5  // opencl_constant
+      5, // opencl_constant
+      0, // cuda_device
+      0, // cuda_constant
+      0  // cuda_shared
   };
 
   class TCETargetInfo : public TargetInfo{
diff --git a/lib/CodeGen/CodeGenModule.cpp b/lib/CodeGen/CodeGenModule.cpp
index f2dda5d658..0da493ffb5 100644
--- a/lib/CodeGen/CodeGenModule.cpp
+++ b/lib/CodeGen/CodeGenModule.cpp
@@ -1172,11 +1172,12 @@ CodeGenModule::GetOrCreateLLVMGlobal(StringRef MangledName,
     DeferredDecls.erase(DDI);
   }
 
+  unsigned AddrSpace = GetGlobalVarAddressSpace(D, Ty->getAddressSpace());
   llvm::GlobalVariable *GV =
     new llvm::GlobalVariable(getModule(), Ty->getElementType(), false,
                              llvm::GlobalValue::ExternalLinkage,
                              0, MangledName, 0,
-                             false, Ty->getAddressSpace());
+                             false, AddrSpace);
 
   // Handle things which are present even on external declarations.
   if (D) {
@@ -1202,7 +1203,10 @@ CodeGenModule::GetOrCreateLLVMGlobal(StringRef MangledName,
     GV->setThreadLocal(D->isThreadSpecified());
   }
 
-  return GV;
+  if (AddrSpace != Ty->getAddressSpace())
+    return llvm::ConstantExpr::getBitCast(GV, Ty);
+  else
+    return GV;
 }
 
 
@@ -1487,6 +1491,20 @@ CodeGenModule::MaybeEmitGlobalStdInitializerListInitializer(const VarDecl *D,
   return llvmInit;
 }
 
+unsigned CodeGenModule::GetGlobalVarAddressSpace(const VarDecl *D,
+                                                 unsigned AddrSpace) {
+  if (LangOpts.CUDA && CodeGenOpts.CUDAIsDevice) {
+    if (D->hasAttr<CUDAConstantAttr>())
+      AddrSpace = getContext().getTargetAddressSpace(LangAS::cuda_constant);
+    else if (D->hasAttr<CUDASharedAttr>())
+      AddrSpace = getContext().getTargetAddressSpace(LangAS::cuda_shared);
+    else
+      AddrSpace = getContext().getTargetAddressSpace(LangAS::cuda_device);
+  }
+
+  return AddrSpace;
+}
+
 void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D) {
   llvm::Constant *Init = 0;
   QualType ASTTy = D->getType();
@@ -1566,7 +1584,7 @@ void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D) {
   if (GV == 0 ||
       GV->getType()->getElementType() != InitType ||
       GV->getType()->getAddressSpace() !=
-        getContext().getTargetAddressSpace(ASTTy)) {
+       GetGlobalVarAddressSpace(D, getContext().getTargetAddressSpace(ASTTy))) {
 
     // Move the old entry aside so that we'll create a new one.
     Entry->setName(StringRef());
diff --git a/lib/CodeGen/CodeGenModule.h b/lib/CodeGen/CodeGenModule.h
index bde03a7ced..8f3bd78ac4 100644
--- a/lib/CodeGen/CodeGenModule.h
+++ b/lib/CodeGen/CodeGenModule.h
@@ -517,6 +517,12 @@ public:
   CreateOrReplaceCXXRuntimeVariable(StringRef Name, llvm::Type *Ty,
                                     llvm::GlobalValue::LinkageTypes Linkage);
 
+  /// GetGlobalVarAddressSpace - Return the address space of the underlying
+  /// global variable for D, as determined by its declaration.  Normally this
+  /// is the same as the address space of D's type, but in CUDA, address spaces
+  /// are associated with declarations, not types.
+  unsigned GetGlobalVarAddressSpace(const VarDecl *D, unsigned AddrSpace);
+
   /// GetAddrOfGlobalVar - Return the llvm::Constant for the address of the
   /// given global variable.  If Ty is non-null and if the global doesn't exist,
   /// then it will be greated with the specified type instead of whatever the
diff --git a/test/CodeGenCUDA/address-spaces.cu b/test/CodeGenCUDA/address-spaces.cu
new file mode 100644
index 0000000000..2da61ec95a
--- /dev/null
+++ b/test/CodeGenCUDA/address-spaces.cu
@@ -0,0 +1,24 @@
+// RUN: %clang_cc1 -emit-llvm %s -o - -fcuda-is-device -triple ptx32-unknown-unknown | FileCheck %s
+
+#include "../SemaCUDA/cuda.h"
+
+// CHECK: @i = global
+__device__ int i;
+
+// CHECK: @j = addrspace(1) global
+__constant__ int j;
+
+// CHECK: @k = addrspace(4) global
+__shared__ int k;
+
+__device__ void foo() {
+  // CHECK: load i32* @i
+  i++;
+
+  // CHECK: load i32* bitcast (i32 addrspace(1)* @j to i32*)
+  j++;
+
+  // CHECK: load i32* bitcast (i32 addrspace(4)* @k to i32*)
+  k++;
+}
+