[CIR][CUDA] Decorate global CUDA shadow variables with appropriate CIR attribute. (llvm#1467)

Started decorating CUDA shadow variables with the shadow_name CIR
attribute which will be used for registering the globals.

Author: Sharp-Edged

clang/lib/CIR/CodeGen/CIRGenCUDARuntime.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "CIRGenCUDARuntime.h"
+#include "CIRGenCXXABI.h"
 #include "CIRGenFunction.h"
 #include "mlir/IR/Operation.h"
 #include "clang/Basic/Cuda.h"
@@ -23,9 +24,26 @@
 using namespace clang;
 using namespace clang::CIRGen;
 
+static std::unique_ptr<MangleContext> initDeviceMC(CIRGenModule &cgm) {
+  // If the host and device have different C++ ABIs, mark it as the device
+  // mangle context so that the mangling needs to retrieve the additional
+  // device lambda mangling number instead of the regular host one.
+  if (cgm.getASTContext().getAuxTargetInfo() &&
+      cgm.getASTContext().getTargetInfo().getCXXABI().isMicrosoft() &&
+      cgm.getASTContext().getAuxTargetInfo()->getCXXABI().isItaniumFamily()) {
+    return std::unique_ptr<MangleContext>(
+        cgm.getASTContext().createDeviceMangleContext(
+            *cgm.getASTContext().getAuxTargetInfo()));
+  }
+
+  return std::unique_ptr<MangleContext>(cgm.getASTContext().createMangleContext(
+      cgm.getASTContext().getAuxTargetInfo()));
+}
+
 CIRGenCUDARuntime::~CIRGenCUDARuntime() {}
 
-CIRGenCUDARuntime::CIRGenCUDARuntime(CIRGenModule &cgm) : cgm(cgm) {
+CIRGenCUDARuntime::CIRGenCUDARuntime(CIRGenModule &cgm)
+    : cgm(cgm), deviceMC(initDeviceMC(cgm)) {
   if (cgm.getLangOpts().OffloadViaLLVM)
     llvm_unreachable("NYI");
   else if (cgm.getLangOpts().HIP)
@@ -289,6 +307,39 @@ mlir::Operation *CIRGenCUDARuntime::getKernelHandle(cir::FuncOp fn,
   return globalOp;
 }
 
+std::string CIRGenCUDARuntime::getDeviceSideName(const NamedDecl *nd) {
+  GlobalDecl gd;
+  // nd could be either a kernel or a variable.
+  if (auto *fd = dyn_cast<FunctionDecl>(nd))
+    gd = GlobalDecl(fd, KernelReferenceKind::Kernel);
+  else
+    gd = GlobalDecl(nd);
+  std::string deviceSideName;
+  MangleContext *mc;
+  if (cgm.getLangOpts().CUDAIsDevice)
+    mc = &cgm.getCXXABI().getMangleContext();
+  else
+    mc = deviceMC.get();
+  if (mc->shouldMangleDeclName(nd)) {
+    SmallString<256> buffer;
+    llvm::raw_svector_ostream out(buffer);
+    mc->mangleName(gd, out);
+    deviceSideName = std::string(out.str());
+  } else
+    deviceSideName = std::string(nd->getIdentifier()->getName());
+
+  // Make unique name for device side static file-scope variable for HIP.
+  if (cgm.getASTContext().shouldExternalize(nd) &&
+      cgm.getLangOpts().GPURelocatableDeviceCode) {
+    SmallString<256> buffer;
+    llvm::raw_svector_ostream out(buffer);
+    out << deviceSideName;
+    cgm.printPostfixForExternalizedDecl(out, nd);
+    deviceSideName = std::string(out.str());
+  }
+  return deviceSideName;
+}
+
 void CIRGenCUDARuntime::internalizeDeviceSideVar(
     const VarDecl *d, cir::GlobalLinkageKind &linkage) {
   if (cgm.getLangOpts().GPURelocatableDeviceCode)

clang/lib/CIR/CodeGen/CIRGenCUDARuntime.h

@@ -47,6 +47,9 @@ class CIRGenCUDARuntime {
   std::string addPrefixToName(StringRef FuncName) const;
   std::string addUnderscoredPrefixToName(StringRef FuncName) const;
 
+  // Mangle context for device.
+  std::unique_ptr<MangleContext> deviceMC;
+
 public:
   CIRGenCUDARuntime(CIRGenModule &cgm);
   virtual ~CIRGenCUDARuntime();
@@ -60,6 +63,9 @@ class CIRGenCUDARuntime {
   virtual mlir::Operation *getKernelHandle(cir::FuncOp fn, GlobalDecl GD);
   virtual void internalizeDeviceSideVar(const VarDecl *d,
                                         cir::GlobalLinkageKind &linkage);
+  /// Returns function or variable name on device side even if the current
+  /// compilation is for host.
+  virtual std::string getDeviceSideName(const NamedDecl *nd);
 };
 
 } // namespace clang::CIRGen

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -576,6 +576,25 @@ bool CIRGenModule::shouldEmitCUDAGlobalVar(const VarDecl *global) const {
          global->getType()->isCUDADeviceBuiltinTextureType();
 }
 
+void CIRGenModule::printPostfixForExternalizedDecl(llvm::raw_ostream &os,
+                                                   const Decl *d) const {
+  // ptxas does not allow '.' in symbol names. On the other hand, HIP prefers
+  // postfix beginning with '.' since the symbol name can be demangled.
+  if (langOpts.HIP)
+    os << (isa<VarDecl>(d) ? ".static." : ".intern.");
+  else
+    os << (isa<VarDecl>(d) ? "__static__" : "__intern__");
+
+  // If the CUID is not specified we try to generate a unique postfix.
+  if (getLangOpts().CUID.empty()) {
+    // TODO: Once we add 'PreprocessorOpts' into CIRGenModule this part can be
+    // brought in from OG.
+    llvm_unreachable("NYI");
+  } else {
+    os << getASTContext().getCUIDHash();
+  }
+}
+
 void CIRGenModule::emitGlobal(GlobalDecl gd) {
   llvm::TimeTraceScope scope("build CIR Global", [&]() -> std::string {
     auto *nd = dyn_cast<NamedDecl>(gd.getDecl());
@@ -1496,6 +1515,32 @@ void CIRGenModule::emitGlobalVarDefinition(const clang::VarDecl *d,
     }
   }
 
+  // Decorate CUDA shadow variables with the cu.shadow_name attribute so we know
+  // how to register them when lowering.
+  if (langOpts.CUDA && !langOpts.CUDAIsDevice &&
+      (d->hasAttr<CUDAConstantAttr>() || d->hasAttr<CUDADeviceAttr>())) {
+    // Shadow variables and their properties must be registered with CUDA
+    // runtime. Skip Extern global variables, which will be registered in
+    // the TU where they are defined.
+    //
+    // Don't register a C++17 inline variable. The local symbol can be
+    // discarded and referencing a discarded local symbol from outside the
+    // comdat (__cuda_register_globals) is disallowed by the ELF spec.
+    //
+    // HIP managed variables need to be always recorded in device and host
+    // compilations for transformation.
+    //
+    // HIP managed variables and variables in CUDADeviceVarODRUsedByHost are
+    // added to llvm.compiler-used, therefore they are safe to be registered.
+    if ((!d->hasExternalStorage() && !d->isInline()) ||
+        getASTContext().CUDADeviceVarODRUsedByHost.contains(d) ||
+        d->hasAttr<HIPManagedAttr>()) {
+      auto shadowName = cudaRuntime->getDeviceSideName(cast<NamedDecl>(d));
+      auto attr = CUDAShadowNameAttr::get(&getMLIRContext(), shadowName);
+      gv->setAttr(CUDAShadowNameAttr::getMnemonic(), attr);
+    }
+  }
+
   // Set initializer and finalize emission
   CIRGenModule::setInitializer(gv, init);
   if (emitter)

clang/lib/CIR/CodeGen/CIRGenModule.h

@@ -274,6 +274,14 @@ class CIRGenModule : public CIRGenTypeCache {
 
   bool shouldEmitCUDAGlobalVar(const VarDecl *global) const;
 
+  /// Print the postfix for externalized static variable or kernels for single
+  /// source offloading languages CUDA and HIP. The unique postfix is created
+  /// using either the CUID argument, or the file's UniqueID and active macros.
+  /// The fallback method without a CUID requires that the offloading toolchain
+  /// does not define separate macros via the -cc1 options.
+  void printPostfixForExternalizedDecl(llvm::raw_ostream &OS,
+                                       const Decl *D) const;
+
   bool shouldEmitConvergenceTokens() const {
     // TODO: this shuld probably become unconditional once the controlled
     // convergence becomes the norm.

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

@@ -10,9 +10,15 @@
 // RUN:            %s -o %t.cir
 // RUN: FileCheck --check-prefix=LLVM-DEVICE --input-file=%t.cir %s
 
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir \
+// RUN:            -x cuda -emit-cir -target-sdk-version=12.3 \
+// RUN:            %s -o %t.cir
+// RUN: FileCheck --check-prefix=CIR-HOST --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-HOST: {{.*}}cir.global external @a = #cir.undef : !s32i {alignment = 4 : i64, cu.shadow_name = #cir.cu.shadow_name<a>}{{.*}}
 
 __shared__ int shared;
 // CIR-DEVICE: cir.global external addrspace(offload_local) @shared = #cir.undef
@@ -21,3 +27,4 @@ __shared__ int shared;
 __constant__ int b;
 // CIR-DEVICE: cir.global constant external addrspace(offload_constant) @b = #cir.int<0> : !s32i {alignment = 4 : i64, cu.externally_initialized = #cir.cu.externally_initialized}
 // LLVM-DEVICE: @b = addrspace(4) externally_initialized constant i32 0, align 4
+// CIR-HOST: {{.*}}cir.global external @b = #cir.undef : !s32i {alignment = 4 : i64, cu.shadow_name = #cir.cu.shadow_name<b>}{{.*}}