[NFC][CUDA] Refactor registering device variable

[yaxunl]

Summary

Extract registering device variable to CUDA runtime codegen function since it
will be called in multiple places.

[tra]

clang/lib/CodeGen/CGCUDANV.cpp

158

clang-format it? void hanging all by itself looks odd.

925

The name and the functionality do not match -- despite what the name says, the function will set internal linkage on any var with device-side attributes.

Rename to internalizeDeviceSideVars ?

clang/lib/CodeGen/CGCUDARuntime.h

85

handleVarRegistration to make it distinct from the implementation of specific registration kinds.

Maybe, move registerDevice* functions under private: as they will only be used from the top-level handler called by CGM.

88

Just registerDeviceVar would do.

[yaxunl]

clang/lib/CodeGen/CGCUDANV.cpp

158

Although it looks strange, it is a result of clang-format.

925

done

clang/lib/CodeGen/CGCUDARuntime.h

85

done

88

done

[yaxunl]

Revised by Artem's comments

[tra]

Couple of minor nits LGTM otherwise,

clang/lib/CodeGen/CGCUDANV.cpp

925

My suggestion could be improved, too. Vars -> Var as we're dealing with only one VarDecl here.

clang/lib/CodeGen/CodeGenModule.cpp

4300–4302

Should we fold internalizeDeviceSideVar into handleVarRegistration or call it from there?
I don't think we'll have any independent use for it in CGM. It seems to be an implementation detail for handleVarRegistration and may not even need to be virtual.

[yaxunl]

clang/lib/CodeGen/CGCUDANV.cpp

925

will do

clang/lib/CodeGen/CodeGenModule.cpp

4300–4302

For function scope static variable, I only need to call handleVarRegistration since the variable already has internal linkage. If internalizeDeviceSideVars is absorbed into handleVarRegistration, there be useless work and I also need to define a useless automatic variable Linkage and pass it to handleVarRegistration.

[tra]

clang/lib/CodeGen/CodeGenModule.cpp

4300–4302

internalizeDeviceSideVars is nearly trivial. I doubt that it will have any measureable impact, either way.

At the moment handleVarRegistration() is called only from here. Are you saying that separate internalization and registration will be needed in the future changes? If that's the case I'm fine keeping them separate.

As things are in the patch, I do not see why internalizeDeviceSideVars needs to be a base class interface. It's always used along with handleVarRegistration and we have no need to be able to override it independently of it. Folding it into or calling it from handleVarRegistration looks like a natural fit.

[yaxunl]

clang/lib/CodeGen/CodeGenModule.cpp

4300–4302

It is used differently in https://reviews.llvm.org/D95560, where only handleVarRegistration is called.

[tra]

LGTM.

clang/lib/CodeGen/CodeGenModule.cpp

4300–4302

Got it.

clang/lib/CodeGen/CGCUDANV.cpp

llvm::Function *makeDummyFunction(llvm::FunctionType *FnTy) {
 
     return DummyFunc;
   }
 
   void emitDeviceStubBodyLegacy(CodeGenFunction &CGF, FunctionArgList &Args);
   void emitDeviceStubBodyNew(CodeGenFunction &CGF, FunctionArgList &Args);
   std::string getDeviceSideName(const NamedDecl *ND) override;
 
-public:
-  CGNVCUDARuntime(CodeGenModule &CGM);
-
-  void emitDeviceStub(CodeGenFunction &CGF, FunctionArgList &Args) override;
   void registerDeviceVar(const VarDecl *VD, llvm::GlobalVariable &Var,
-                         bool Extern, bool Constant) override {
+                         bool Extern, bool Constant) {
     DeviceVars.push_back({&Var,
                           VD,
                           {DeviceVarFlags::Variable, Extern, Constant,
                            VD->hasAttr<HIPManagedAttr>(),
                            /*Normalized*/ false, 0}});
   }
   void registerDeviceSurf(const VarDecl *VD, llvm::GlobalVariable &Var,
-                          bool Extern, int Type) override {
+                          bool Extern, int Type) {
     DeviceVars.push_back({&Var,
                           VD,
                           {DeviceVarFlags::Surface, Extern, /*Constant*/ false,
                            /*Managed*/ false,
                            /*Normalized*/ false, Type}});
   }
   void registerDeviceTex(const VarDecl *VD, llvm::GlobalVariable &Var,
-                         bool Extern, int Type, bool Normalized) override {
+                         bool Extern, int Type, bool Normalized) {
     DeviceVars.push_back({&Var,
                           VD,
                           {DeviceVarFlags::Texture, Extern, /*Constant*/ false,
                            /*Managed*/ false, Normalized, Type}});
   }
 
+public:
+  CGNVCUDARuntime(CodeGenModule &CGM);
+
+  void emitDeviceStub(CodeGenFunction &CGF, FunctionArgList &Args) override;
+  void handleVarRegistration(const VarDecl *VD,
+                             llvm::GlobalVariable &Var) override;
+
   /// Creates module constructor function
   llvm::Function *makeModuleCtorFunction() override;
   /// Creates module destructor function
   llvm::Function *makeModuleDtorFunction() override;
+  void
+  internalizeDeviceSideVar(const VarDecl *D,
+                           llvm::GlobalValue::LinkageTypes &Linkage) override;
 };
 
 }
 
 std::string CGNVCUDARuntime::addPrefixToName(StringRef FuncName) const {
   if (CGM.getLangOpts().HIP)
     return ((Twine("hip") + Twine(FuncName)).str());
   return ((Twine("cuda") + Twine(FuncName)).str());
llvm::Function *CGNVCUDARuntime::makeModuleDtorFunction() {
   }
   DtorBuilder.CreateRetVoid();
   return ModuleDtorFunc;
 }
 
 CGCUDARuntime *CodeGen::CreateNVCUDARuntime(CodeGenModule &CGM) {
   return new CGNVCUDARuntime(CGM);
 }
+
+void CGNVCUDARuntime::internalizeDeviceSideVar(
+    const VarDecl *D, llvm::GlobalValue::LinkageTypes &Linkage) {
+  // Host-side shadows of external declarations of device-side
+  // global variables become internal definitions. These have to
+  // be internal in order to prevent name conflicts with global
+  // host variables with the same name in a different TUs.
+  //
+  // __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>() ||
+      D->getType()->isCUDADeviceBuiltinSurfaceType() ||
+      D->getType()->isCUDADeviceBuiltinTextureType()) {
+    Linkage = llvm::GlobalValue::InternalLinkage;
+  }
+}
+
+void CGNVCUDARuntime::handleVarRegistration(const VarDecl *D,
+                                            llvm::GlobalVariable &GV) {
+  if (D->hasAttr<CUDADeviceAttr>() || D->hasAttr<CUDAConstantAttr>()) {
+    // 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.
+    // TODO: Reject __device__ constexpr and __device__ inline in Sema.
+    if (!D->hasExternalStorage() && !D->isInline())
+      registerDeviceVar(D, GV, !D->hasDefinition(),
+                        D->hasAttr<CUDAConstantAttr>());
+  } else if (D->getType()->isCUDADeviceBuiltinSurfaceType() ||
+             D->getType()->isCUDADeviceBuiltinTextureType()) {
+    // Builtin surfaces and textures and their template arguments are
+    // also registered with CUDA runtime.
+    const ClassTemplateSpecializationDecl *TD =
+        cast<ClassTemplateSpecializationDecl>(
+            D->getType()->getAs<RecordType>()->getDecl());
+    const TemplateArgumentList &Args = TD->getTemplateArgs();
+    if (TD->hasAttr<CUDADeviceBuiltinSurfaceTypeAttr>()) {
+      assert(Args.size() == 2 &&
+             "Unexpected number of template arguments of CUDA device "
+             "builtin surface type.");
+      auto SurfType = Args[1].getAsIntegral();
+      if (!D->hasExternalStorage())
+        registerDeviceSurf(D, GV, !D->hasDefinition(), SurfType.getSExtValue());
+    } else {
+      assert(Args.size() == 3 &&
+             "Unexpected number of template arguments of CUDA device "
+             "builtin texture type.");
+      auto TexType = Args[1].getAsIntegral();
+      auto Normalized = Args[2].getAsIntegral();
+      if (!D->hasExternalStorage())
+        registerDeviceTex(D, GV, !D->hasDefinition(), TexType.getSExtValue(),
+                          Normalized.getZExtValue());
+    }
+  }
+}

clang/lib/CodeGen/CGCUDARuntime.h

 // runtime libraries.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_LIB_CODEGEN_CGCUDARUNTIME_H
 #define LLVM_CLANG_LIB_CODEGEN_CGCUDARUNTIME_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/GlobalValue.h"
 
 namespace llvm {
 class Function;
 class GlobalVariable;
 }
 
 namespace clang {
 
public:
   virtual ~CGCUDARuntime();
 
   virtual RValue EmitCUDAKernelCallExpr(CodeGenFunction &CGF,
                                         const CUDAKernelCallExpr *E,
                                         ReturnValueSlot ReturnValue);
 
   /// Emits a kernel launch stub.
   virtual void emitDeviceStub(CodeGenFunction &CGF, FunctionArgList &Args) = 0;
-  virtual void registerDeviceVar(const VarDecl *VD, llvm::GlobalVariable &Var,
-                                 bool Extern, bool Constant) = 0;
-  virtual void registerDeviceSurf(const VarDecl *VD, llvm::GlobalVariable &Var,
-                                  bool Extern, int Type) = 0;
-  virtual void registerDeviceTex(const VarDecl *VD, llvm::GlobalVariable &Var,
-                                 bool Extern, int Type, bool Normalized) = 0;
+
+  /// Check whether a variable is a device variable and register it if true.
+  virtual void handleVarRegistration(const VarDecl *VD,
+                                     llvm::GlobalVariable &Var) = 0;
 
   /// Constructs and returns a module initialization function or nullptr if it's
   /// not needed. Must be called after all kernels have been emitted.
   virtual llvm::Function *makeModuleCtorFunction() = 0;
 
   /// Returns a module cleanup function or nullptr if it's not needed.
   /// Must be called after ModuleCtorFunction
   virtual llvm::Function *makeModuleDtorFunction() = 0;
 
   /// Returns function or variable name on device side even if the current
   /// compilation is for host.
   virtual std::string getDeviceSideName(const NamedDecl *ND) = 0;
+
+  /// Adjust linkage of shadow variables in host compilation.
+  virtual void
+  internalizeDeviceSideVar(const VarDecl *D,
+                           llvm::GlobalValue::LinkageTypes &Linkage) = 0;
 };
 
 /// Creates an instance of a CUDA runtime class.
 CGCUDARuntime *CreateNVCUDARuntime(CodeGenModule &CGM);
 
 }
 }
 
 #endif

clang/lib/CodeGen/CodeGenModule.cpp

void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D,
   // through the runtime library (cudaGetSymbolAddress() / cudaGetSymbolSize()
   // / cudaMemcpyToSymbol() / cudaMemcpyFromSymbol())."
   if (GV && LangOpts.CUDA) {
     if (LangOpts.CUDAIsDevice) {
       if (Linkage != llvm::GlobalValue::InternalLinkage &&
           (D->hasAttr<CUDADeviceAttr>() || D->hasAttr<CUDAConstantAttr>()))
         GV->setExternallyInitialized(true);
     } else {
-      // Host-side shadows of external declarations of device-side
-      // global variables become internal definitions. These have to
-      // be internal in order to prevent name conflicts with global
-      // host variables with the same name in a different TUs.
-      if (D->hasAttr<CUDADeviceAttr>() || D->hasAttr<CUDAConstantAttr>()) {
-        Linkage = llvm::GlobalValue::InternalLinkage;
-        // 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.
-        // TODO: Reject __device__ constexpr and __device__ inline in Sema.
-        if (!D->hasExternalStorage() && !D->isInline())
-          getCUDARuntime().registerDeviceVar(D, *GV, !D->hasDefinition(),
-                                             D->hasAttr<CUDAConstantAttr>());
-      } else if (D->hasAttr<CUDASharedAttr>()) {
-        // __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.
-        Linkage = llvm::GlobalValue::InternalLinkage;
-      } else if (D->getType()->isCUDADeviceBuiltinSurfaceType() ||
-                 D->getType()->isCUDADeviceBuiltinTextureType()) {
-        // Builtin surfaces and textures and their template arguments are
-        // also registered with CUDA runtime.
-        Linkage = llvm::GlobalValue::InternalLinkage;
-        const ClassTemplateSpecializationDecl *TD =
-            cast<ClassTemplateSpecializationDecl>(
-                D->getType()->getAs<RecordType>()->getDecl());
-        const TemplateArgumentList &Args = TD->getTemplateArgs();
-        if (TD->hasAttr<CUDADeviceBuiltinSurfaceTypeAttr>()) {
-          assert(Args.size() == 2 &&
-                 "Unexpected number of template arguments of CUDA device "
-                 "builtin surface type.");
-          auto SurfType = Args[1].getAsIntegral();
-          if (!D->hasExternalStorage())
-            getCUDARuntime().registerDeviceSurf(D, *GV, !D->hasDefinition(),
-                                                SurfType.getSExtValue());
-        } else {
-          assert(Args.size() == 3 &&
-                 "Unexpected number of template arguments of CUDA device "
-                 "builtin texture type.");
-          auto TexType = Args[1].getAsIntegral();
-          auto Normalized = Args[2].getAsIntegral();
-          if (!D->hasExternalStorage())
-            getCUDARuntime().registerDeviceTex(D, *GV, !D->hasDefinition(),
-                                               TexType.getSExtValue(),
-                                               Normalized.getZExtValue());
-        }
-      }
+      getCUDARuntime().internalizeDeviceSideVar(D, Linkage);
+      getCUDARuntime().handleVarRegistration(D, *GV);
     }
   }
 
   GV->setInitializer(Init);
   if (emitter)
     emitter->finalize(GV);
 
   // If it is safe to mark the global 'constant', do so now.
   GV->setConstant(!NeedsGlobalCtor && !NeedsGlobalDtor &&