[CIR][CUDA] Handle shared and local variables (#1368)

CUDA shared variables are device-only, accessible from all threads in a
block of some kernel. It's similar to `local` variables in OpenCL which
all threads in a work-group can access. Hence they are realized as
`static` variables in addrspace(local).

On the other hand, the local variables inside a kernel (without special
attributes) are just regular variables, typically emitted by
`CreateTempAlloca`. They are in the default address space.

OG checks if the expected address space, denoted by the type, is the
same as the actual address space indicated by attributes. If they aren't
the same, a `addrspacecast` is emitted when a global variable is
accessed. In CIR however, `cir.get_global` alreadys carries that
information in `!cir.ptr` type, so we don't need a cast.

Author: AdUhTkJm

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -479,9 +479,10 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
   // OpenCL variables in local address space and CUDA shared
   // variables cannot have an initializer.
   mlir::Attribute Init = nullptr;
-  if (D.hasAttr<CUDASharedAttr>() || D.hasAttr<LoaderUninitializedAttr>())
+  if (D.hasAttr<LoaderUninitializedAttr>())
     llvm_unreachable("CUDA is NYI");
-  else if (Ty.getAddressSpace() != LangAS::opencl_local)
+  else if (Ty.getAddressSpace() != LangAS::opencl_local &&
+           !D.hasAttr<CUDASharedAttr>())
     Init = builder.getZeroInitAttr(convertType(Ty));
 
   cir::GlobalOp GV = builder.createVersionedGlobal(
@@ -499,9 +500,14 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
 
   setGVProperties(GV, &D);
 
-  // Make sure the result is of the correct type.
-  if (AS != builder.getAddrSpaceAttr(Ty.getAddressSpace()))
-    llvm_unreachable("address space cast NYI");
+  // OG checks if the expected address space, denoted by the type, is the
+  // same as the actual address space indicated by attributes. If they aren't
+  // the same, an addrspacecast is emitted when this variable is accessed.
+  // In CIR however, cir.get_global alreadys carries that information in
+  // !cir.ptr type - if this global is in OpenCL local address space, then its
+  // type would be !cir.ptr<..., addrspace(offload_local)>. Therefore we don't
+  // need an explicit address space cast in CIR: they will get emitted when
+  // lowering to LLVM IR.
 
   // Ensure that the static local gets initialized by making sure the parent
   // function gets emitted eventually.

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -1604,7 +1604,14 @@ LangAS CIRGenModule::getGlobalConstantAddressSpace() const {
 LangAS CIRGenModule::getLangTempAllocaAddressSpace() const {
   if (getLangOpts().OpenCL)
     return LangAS::opencl_private;
-  if (getLangOpts().SYCLIsDevice || getLangOpts().CUDAIsDevice ||
+
+  // For temporaries inside functions, CUDA treats them as normal variables.
+  // LangAS::cuda_device, on the other hand, is reserved for those variables
+  // explicitly marked with __device__.
+  if (getLangOpts().CUDAIsDevice)
+    return LangAS::Default;
+
+  if (getLangOpts().SYCLIsDevice ||
       (getLangOpts().OpenMP && getLangOpts().OpenMPIsTargetDevice))
     llvm_unreachable("NYI");
   return LangAS::Default;
@@ -3983,8 +3990,11 @@ LangAS CIRGenModule::getGlobalVarAddressSpace(const VarDecl *D) {
       (!D || D->getType().getAddressSpace() == LangAS::Default))
     llvm_unreachable("NYI");
 
-  if (langOpts.CUDA && langOpts.CUDAIsDevice)
+  if (langOpts.CUDA && langOpts.CUDAIsDevice) {
+    if (D && D->hasAttr<CUDASharedAttr>())
+      return LangAS::cuda_shared;
     llvm_unreachable("NYI");
+  }
 
   if (langOpts.OpenMP)
     llvm_unreachable("NYI");

clang/lib/CIR/Dialect/IR/CIRAttrs.cpp

@@ -645,6 +645,9 @@ AddressSpaceAttr::getValueFromLangAS(clang::LangAS langAS) {
   case LangAS::opencl_global:
     return Kind::offload_global;
   case LangAS::opencl_local:
+  case LangAS::cuda_shared:
+    // Local means local among the work-group (OpenCL) or block (CUDA).
+    // All threads inside the kernel can access local memory.
     return Kind::offload_local;
   case LangAS::opencl_constant:
     return Kind::offload_constant;
@@ -657,7 +660,6 @@ AddressSpaceAttr::getValueFromLangAS(clang::LangAS langAS) {
   case LangAS::opencl_global_host:
   case LangAS::cuda_device:
   case LangAS::cuda_constant:
-  case LangAS::cuda_shared:
   case LangAS::sycl_global:
   case LangAS::sycl_global_device:
   case LangAS::sycl_global_host:

clang/test/CIR/CodeGen/CUDA/address-spaces.cu

@@ -0,0 +1,19 @@
+#include "../Inputs/cuda.h"
+
+// RUN: %clang_cc1 -triple nvptx64-nvidia-cuda -fclangir \
+// RUN:            -fcuda-is-device -emit-cir -target-sdk-version=12.3 \
+// RUN:            %s -o %t.cir
+// RUN: FileCheck --check-prefix=CIR --input-file=%t.cir %s
+
+__global__ void fn() {
+  int i = 0;
+  __shared__ int j;
+  j = i;
+}
+
+// CIR: cir.global "private" internal dsolocal addrspace(offload_local) @_ZZ2fnvE1j : !s32i
+// CIR: cir.func @_Z2fnv
+// CIR: [[Local:%[0-9]+]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["i", init]
+// CIR: [[Shared:%[0-9]+]] = cir.get_global @_ZZ2fnvE1j : !cir.ptr<!s32i, addrspace(offload_local)>
+// CIR: [[Tmp:%[0-9]+]] = cir.load [[Local]] : !cir.ptr<!s32i>, !s32i
+// CIR: cir.store [[Tmp]], [[Shared]] : !s32i, !cir.ptr<!s32i, addrspace(offload_local)>