[CIR][CUDA] Lowering device and shared variables (#1438)

Currently `__shared__` and `__constant__` variables are ignored by
CodeGen. This patch fixes this.
(It is also fixed in #1436 .)

Device and constant variables should be marked as
`externally_initialized`, as they might be initialized by host, rather
than on device. We can't identify which variables are device ones at
lowering stage, so this patch adds a new attribute for it in CodeGen.

Similar to `__global__` functions, global variables on device
corresponds to "shadow" variables on host, and they must be registered
to their counterpart. I added a `CUDAShadowNameAttr` in this patch for
later use, but I didn't insert code to actually generate it.

Author: AdUhTkJm

clang/include/clang/CIR/Dialect/IR/CIRCUDAAttrs.td

@@ -35,6 +35,24 @@ def CUDAKernelNameAttr : CIR_Attr<"CUDAKernelName",
   let assemblyFormat = "`<` $kernel_name `>`";
 }
 
+def CUDAShadowNameAttr : CIR_Attr<"CUDAShadowName",
+                                  "cu.shadow_name"> {
+  let summary = "Device-side global variable name for this shadow.";
+  let description =
+  [{
+    This attribute is attached to global variable definitions and records the
+    mangled name of the global variable used on the device.
+
+    In CUDA, __device__, __constant__ and __shared__ variables, as well as 
+    surface and texture variables, will generate a shadow symbol on host.
+    We must preserve the correspodence in order to generate registration
+    functions.
+  }];
+
+  let parameters = (ins "std::string":$device_side_name);
+  let assemblyFormat = "`<` $device_side_name `>`";
+}
+
 def CUDABinaryHandleAttr : CIR_Attr<"CUDABinaryHandle",
                                   "cu.binary_handle"> {
   let summary = "Fat binary handle for device code.";
@@ -52,4 +70,17 @@ def CUDABinaryHandleAttr : CIR_Attr<"CUDABinaryHandle",
   let assemblyFormat = "`<` $name `>`";
 }
 
+def CUDAExternallyInitializedAttr : CIR_Attr<"CUDAExternallyInitialized",
+                                             "cu.externally_initialized"> {
+  let summary = "The marked variable is externally initialized.";
+  let description =
+  [{
+    CUDA __device__ and __constant__ variables, along with surface and
+    textures, might be initialized by host, hence "externally initialized".
+    Therefore they must be emitted even if they are not referenced.
+
+    The attribute corresponds to the attribute on LLVM with the same name.
+  }];
+}
+
 #endif // MLIR_CIR_DIALECT_CIR_CUDA_ATTRS

clang/lib/CIR/CodeGen/CIRGenCUDARuntime.cpp

@@ -283,3 +283,23 @@ mlir::Operation *CIRGenCUDARuntime::getKernelHandle(cir::FuncOp fn,
 
   return globalOp;
 }
+
+void CIRGenCUDARuntime::internalizeDeviceSideVar(
+    const VarDecl *d, cir::GlobalLinkageKind &linkage) {
+  if (cgm.getLangOpts().GPURelocatableDeviceCode)
+    llvm_unreachable("NYI");
+
+  // __shared__ variables are odd. Shadows do get created, but
+  // they are not registered with the CUDA runtime, so they
+  // can't really be used to access their device-side
+  // counterparts. It's not clear yet whether it's nvcc's bug or
+  // a feature, but we've got to do the same for compatibility.
+  if (d->hasAttr<CUDADeviceAttr>() || d->hasAttr<CUDAConstantAttr>() ||
+      d->hasAttr<CUDASharedAttr>()) {
+    linkage = cir::GlobalLinkageKind::InternalLinkage;
+  }
+
+  if (d->getType()->isCUDADeviceBuiltinSurfaceType() ||
+      d->getType()->isCUDADeviceBuiltinTextureType())
+    llvm_unreachable("NYI");
+}

clang/lib/CIR/CodeGen/CIRGenCUDARuntime.h

@@ -58,6 +58,8 @@ class CIRGenCUDARuntime {
                                         const CUDAKernelCallExpr *expr,
                                         ReturnValueSlot retValue);
   virtual mlir::Operation *getKernelHandle(cir::FuncOp fn, GlobalDecl GD);
+  virtual void internalizeDeviceSideVar(const VarDecl *d,
+                                        cir::GlobalLinkageKind &linkage);
 };
 
 } // namespace clang::CIRGen

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -570,13 +570,13 @@ bool CIRGenModule::shouldEmitCUDAGlobalVar(const VarDecl *global) const {
   // their device-side incarnations.
 
   if (global->hasAttr<CUDAConstantAttr>() ||
-      global->hasAttr<CUDASharedAttr>() ||
       global->getType()->isCUDADeviceBuiltinSurfaceType() ||
       global->getType()->isCUDADeviceBuiltinTextureType()) {
     llvm_unreachable("NYI");
   }
 
-  return !langOpts.CUDAIsDevice || global->hasAttr<CUDADeviceAttr>();
+  return !langOpts.CUDAIsDevice || global->hasAttr<CUDADeviceAttr>() ||
+         global->hasAttr<CUDASharedAttr>();
 }
 
 void CIRGenModule::emitGlobal(GlobalDecl gd) {
@@ -599,8 +599,10 @@ void CIRGenModule::emitGlobal(GlobalDecl gd) {
   assert(!global->hasAttr<CPUDispatchAttr>() && "NYI");
 
   if (langOpts.CUDA || langOpts.HIP) {
-    // clang uses the same flag when building HIP code
-    if (langOpts.CUDAIsDevice) {
+    if (const auto *vd = dyn_cast<VarDecl>(global)) {
+      if (!shouldEmitCUDAGlobalVar(vd))
+        return;
+    } else if (langOpts.CUDAIsDevice) {
       // This will implicitly mark templates and their
       // specializations as __host__ __device__.
       if (langOpts.OffloadImplicitHostDeviceTemplates)
@@ -622,11 +624,6 @@ void CIRGenModule::emitGlobal(GlobalDecl gd) {
         return;
       }
     }
-
-    if (const auto *vd = dyn_cast<VarDecl>(global)) {
-      if (!shouldEmitCUDAGlobalVar(vd))
-        return;
-    }
   }
 
   if (langOpts.OpenMP) {
@@ -1395,7 +1392,7 @@ void CIRGenModule::emitGlobalVarDefinition(const clang::VarDecl *d,
        d->getType()->isCUDADeviceBuiltinTextureType());
   if (getLangOpts().CUDA &&
       (isCudaSharedVar || isCudaShadowVar || isCudaDeviceShadowVar))
-    assert(0 && "not implemented");
+    init = UndefAttr::get(&getMLIRContext(), convertType(d->getType()));
   else if (d->hasAttr<LoaderUninitializedAttr>())
     assert(0 && "not implemented");
   else if (!initExpr) {
@@ -1491,11 +1488,19 @@ void CIRGenModule::emitGlobalVarDefinition(const clang::VarDecl *d,
   cir::GlobalLinkageKind linkage =
       getCIRLinkageVarDefinition(d, /*IsConstant=*/false);
 
-  // TODO(cir):
   // CUDA B.2.1 "The __device__ qualifier declares a variable that resides on
   // the device. [...]"
   // CUDA B.2.2 "The __constant__ qualifier, optionally used together with
   // __device__, declares a variable that: [...]
+  if (langOpts.CUDA && langOpts.CUDAIsDevice) {
+    // __shared__ variables is not marked as externally initialized,
+    // because they must not be initialized.
+    if (linkage != cir::GlobalLinkageKind::InternalLinkage &&
+        (d->hasAttr<CUDADeviceAttr>())) {
+      gv->setAttr(CUDAExternallyInitializedAttr::getMnemonic(),
+                  CUDAExternallyInitializedAttr::get(&getMLIRContext()));
+    }
+  }
 
   // Set initializer and finalize emission
   CIRGenModule::setInitializer(gv, init);

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -2397,6 +2397,12 @@ mlir::LogicalResult CIRToLLVMGlobalOpLowering::matchAndRewrite(
 
   attributes.push_back(rewriter.getNamedAttr("visibility_", visibility));
 
+  if (auto extInit =
+          op->getAttr(CUDAExternallyInitializedAttr::getMnemonic())) {
+    attributes.push_back(rewriter.getNamedAttr("externally_initialized",
+                                               rewriter.getUnitAttr()));
+  }
+
   if (init.has_value()) {
     if (mlir::isa<cir::FPAttr, cir::IntAttr, cir::BoolAttr>(init.value())) {
       // If a directly equivalent attribute is available, use it.

clang/test/CIR/CodeGen/CUDA/global-vars.cu

@@ -5,7 +5,15 @@
 // RUN:            %s -o %t.cir
 // RUN: FileCheck --check-prefix=CIR-DEVICE --input-file=%t.cir %s
 
+// RUN: %clang_cc1 -triple nvptx64-nvidia-cuda -fclangir \
+// RUN:            -fcuda-is-device -emit-llvm -target-sdk-version=12.3 \
+// RUN:            %s -o %t.cir
+// RUN: FileCheck --check-prefix=LLVM-DEVICE --input-file=%t.cir %s
 
 __device__ int a;
+// CIR-DEVICE: cir.global external addrspace(offload_global) @a = #cir.int<0>
+// LLVM-DEVICE: @a = addrspace(1) externally_initialized global i32 0, align 4
 
-// CIR-DEVICE: cir.global external addrspace(offload_global) @a = #cir.int<0> : !s32i {alignment = 4 : i64} loc(#loc3)
+__shared__ int shared;
+// CIR-DEVICE: cir.global external addrspace(offload_local) @shared = #cir.undef
+// LLVM-DEVICE: @shared = addrspace(3) global i32 undef, align 4