[WIP][CUDA] Use shared MangleContext for CUDA and CXX CG

[psalz]

Summary

NOTE: This is a work in progress and mainly intended to highlight the issue - i.e., I'm not certain the provided solution is appropriate.

Given this CUDA program

template<typename Lambda>
__global__ void run_this(Lambda lambda) {
    lambda();
}

template<typename T>
struct remove_reference {
    using type = T;
};

template<typename T>
struct remove_reference<T&> {
    using type = T;
};

template<typename T>
constexpr typename remove_reference<T>::type&& move(T&& t) {
    return static_cast<typename remove_reference<T>::type&&>(t);
}

int main() {
    auto foo = move([](){});
    run_this<<<1, 1, 1>>>([]() __device__ { printf("Hello World\n"); }); 
    return 0;
}

the assertion at the top of CGNVCUDARuntime::emitDeviceStub will fail. For release builds the effect is simply a cudaErrorInvalidDeviceFunction error at run time. The reason for this is that the mangled names of the device stub and the actual device side function differ: The stub is called _Z8run_thisIZ4mainE3$_1EvT_, while the device function is _Z8run_thisIZ4mainE3$_0EvT_. The difference comes down to the anonymous struct ID that is maintained and assigned by the ManglerContext. It appears that for the latter getAnonymousStructId is never called for the moved no-op lambda, resulting in an ID of 0 for the kernel.

My proposed solution would be to simply share the ManglerContext used by the CGNVCUDARuntime and CGCXXABI code generators. For this I've added a new ASTContext::getSharedMangleContext function that memoizes created manglers for the given target ABI. From looking at ManglerContext to me at least it doesn't look like that could cause any issues, but then again, I really don't know much about Clang's internals.

Of course an alternative solution could be to make sure that getAnonymousStructId is always called for both lambdas (and in the correct order), but again I don't really know why that is not happening in the first place.

[erichkeane]

I don't know of any problems that the shared mangle context would cause, though I'm not sure about using the shared_ptr. It seems to me that SOMEONE should own this, and the other should use a reference. IMO, the SharedMangleContexts should be a unique_ptr and CGCXXABI/CGCUDANV should have a reference to the proper implementation.

[psalz]

Move ownership of shared MangleContexts to ASTContext, return references from getSharedMangleContext.

[aheejin]

Sorry, I don't think I know enough about this code to review this.

include/clang/AST/ASTContext.h

 #include "clang/AST/CanonicalType.h"
 #include "clang/AST/CommentCommandTraits.h"
 #include "clang/AST/ComparisonCategories.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclarationName.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExternalASTSource.h"
+#include "clang/AST/Mangle.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/PrettyPrinter.h"
 #include "clang/AST/RawCommentList.h"
 #include "clang/AST/TemplateBase.h"
 #include "clang/AST/TemplateName.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/AddressSpaces.h"
 #include "clang/Basic/AttrKinds.h"

 class CharUnits;
 class CXXABI;
 class CXXConstructorDecl;
 class CXXMethodDecl;
 class CXXRecordDecl;
 class DiagnosticsEngine;
 class Expr;
 class FixedPointSemantics;
-class MangleContext;
 class MangleNumberingContext;
 class MaterializeTemporaryExpr;
 class MemberSpecializationInfo;
 class Module;
 class ObjCCategoryDecl;
 class ObjCCategoryImplDecl;
 class ObjCContainerDecl;
 class ObjCImplDecl;
public:
 
   bool isNearlyEmpty(const CXXRecordDecl *RD) const;
 
   VTableContextBase *getVTableContext();
 
   /// If \p T is null pointer, assume the target in ASTContext.
   MangleContext *createMangleContext(const TargetInfo *T = nullptr);
 
+private:
+  llvm::DenseMap<std::underlying_type<MangleContext::ManglerKind>::type,
+    std::unique_ptr<MangleContext>> SharedMangleContexts;
+
+public:
+  /// If \p T is null pointer, assume the target in ASTContext.
+  MangleContext &getSharedMangleContext(const TargetInfo *T = nullptr);
+
   void DeepCollectObjCIvars(const ObjCInterfaceDecl *OI, bool leafClass,
                             SmallVectorImpl<const ObjCIvarDecl*> &Ivars) const;
 
   unsigned CountNonClassIvars(const ObjCInterfaceDecl *OI) const;
   void CollectInheritedProtocols(const Decl *CDecl,
                           llvm::SmallPtrSet<ObjCProtocolDecl*, 8> &Protocols);
 
   /// Return true if the specified type has unique object representations

lib/AST/ASTContext.cpp

if (!VTContext.get()) {
     if (Target->getCXXABI().isMicrosoft())
       VTContext.reset(new MicrosoftVTableContext(*this));
     else
       VTContext.reset(new ItaniumVTableContext(*this));
   }
   return VTContext.get();
 }
 
-MangleContext *ASTContext::createMangleContext(const TargetInfo *T) {
-  if (!T)
-    T = Target;
-  switch (T->getCXXABI().getKind()) {
+MangleContext::ManglerKind getManglerKindForABI(TargetCXXABI::Kind K) {
+  switch (K) {
   case TargetCXXABI::GenericAArch64:
   case TargetCXXABI::GenericItanium:
   case TargetCXXABI::GenericARM:
   case TargetCXXABI::GenericMIPS:
   case TargetCXXABI::iOS:
   case TargetCXXABI::iOS64:
   case TargetCXXABI::WebAssembly:
   case TargetCXXABI::WatchOS:
-    return ItaniumMangleContext::create(*this, getDiagnostics());
+    return MangleContext::MK_Itanium;
   case TargetCXXABI::Microsoft:
-    return MicrosoftMangleContext::create(*this, getDiagnostics());
+    return MangleContext::MK_Microsoft;
   }
   llvm_unreachable("Unsupported ABI");
 }
 
+MangleContext *ASTContext::createMangleContext(const TargetInfo *T) {
+  if (!T)
+    T = Target;
+  switch (getManglerKindForABI(T->getCXXABI().getKind())) {
+  case MangleContext::MK_Itanium:
+    return ItaniumMangleContext::create(*this, getDiagnostics());
+  case MangleContext::MK_Microsoft:
+    return MicrosoftMangleContext::create(*this, getDiagnostics());
+  }
+  llvm_unreachable("Unsupported MangleContext");
+}
+
+MangleContext &ASTContext::getSharedMangleContext(const TargetInfo *T) {
+  if (!T)
+    T = Target;
+  auto Kind = getManglerKindForABI(T->getCXXABI().getKind());
+  auto I = SharedMangleContexts.find(Kind);
+  if (I == SharedMangleContexts.end())
+    I = SharedMangleContexts.insert({Kind,
+        std::unique_ptr<MangleContext>(createMangleContext(T))}).first;
+  return *I->second;
+}
+
 CXXABI::~CXXABI() = default;
 
 size_t ASTContext::getSideTableAllocatedMemory() const {
   return ASTRecordLayouts.getMemorySize() +
          llvm::capacity_in_bytes(ObjCLayouts) +
          llvm::capacity_in_bytes(KeyFunctions) +
          llvm::capacity_in_bytes(ObjCImpls) +
          llvm::capacity_in_bytes(BlockVarCopyInits) +

lib/CodeGen/CGCUDANV.cpp

private:
   llvm::SmallVector<VarInfo, 16> DeviceVars;
   /// Keeps track of variable containing handle of GPU binary. Populated by
   /// ModuleCtorFunction() and used to create corresponding cleanup calls in
   /// ModuleDtorFunction()
   llvm::GlobalVariable *GpuBinaryHandle = nullptr;
   /// Whether we generate relocatable device code.
   bool RelocatableDeviceCode;
   /// Mangle context for device.
-  std::unique_ptr<MangleContext> DeviceMC;
+  MangleContext &DeviceMC;
 
   llvm::FunctionCallee getSetupArgumentFn() const;
   llvm::FunctionCallee getLaunchFn() const;
 
   llvm::FunctionType *getRegisterGlobalsFnTy() const;
   llvm::FunctionType *getCallbackFnTy() const;
   llvm::FunctionType *getRegisterLinkedBinaryFnTy() const;
   std::string addPrefixToName(StringRef FuncName) const;
if (CGM.getLangOpts().HIP)
     return ((Twine("__hip") + Twine(FuncName)).str());
   return ((Twine("__cuda") + Twine(FuncName)).str());
 }
 
 CGNVCUDARuntime::CGNVCUDARuntime(CodeGenModule &CGM)
     : CGCUDARuntime(CGM), Context(CGM.getLLVMContext()),
       TheModule(CGM.getModule()),
       RelocatableDeviceCode(CGM.getLangOpts().GPURelocatableDeviceCode),
-      DeviceMC(CGM.getContext().createMangleContext(
+      DeviceMC(CGM.getContext().getSharedMangleContext(
           CGM.getContext().getAuxTargetInfo())) {
   CodeGen::CodeGenTypes &Types = CGM.getTypes();
   ASTContext &Ctx = CGM.getContext();
 
   IntTy = CGM.IntTy;
   SizeTy = CGM.SizeTy;
   VoidTy = CGM.VoidTy;
 
llvm::FunctionType *CGNVCUDARuntime::getRegisterLinkedBinaryFnTy() const {
   llvm::Type *Params[] = {RegisterGlobalsFnTy->getPointerTo(), VoidPtrTy,
                           VoidPtrTy, CallbackFnTy->getPointerTo()};
   return llvm::FunctionType::get(VoidTy, Params, false);
 }
 
 std::string CGNVCUDARuntime::getDeviceSideName(const Decl *D) {
   auto *ND = cast<const NamedDecl>(D);
   std::string DeviceSideName;
-  if (DeviceMC->shouldMangleDeclName(ND)) {
+  if (DeviceMC.shouldMangleDeclName(ND)) {
     SmallString<256> Buffer;
     llvm::raw_svector_ostream Out(Buffer);
-    DeviceMC->mangleName(ND, Out);
+    DeviceMC.mangleName(ND, Out);
     DeviceSideName = Out.str();
   } else
     DeviceSideName = ND->getIdentifier()->getName();
   return DeviceSideName;
 }
 
 void CGNVCUDARuntime::emitDeviceStub(CodeGenFunction &CGF,
                                      FunctionArgList &Args) {

lib/CodeGen/CGCXXABI.h

 class CodeGenFunction;
 class CodeGenModule;
 struct CatchTypeInfo;
 
 /// Implements C++ ABI-specific code generation functions.
 class CGCXXABI {
 protected:
   CodeGenModule &CGM;
-  std::unique_ptr<MangleContext> MangleCtx;
+  MangleContext &MangleCtx;
 
   CGCXXABI(CodeGenModule &CGM)
-    : CGM(CGM), MangleCtx(CGM.getContext().createMangleContext()) {}
+    : CGM(CGM), MangleCtx(CGM.getContext().getSharedMangleContext()) {}
 
 protected:
   ImplicitParamDecl *getThisDecl(CodeGenFunction &CGF) {
     return CGF.CXXABIThisDecl;
   }
   llvm::Value *getThisValue(CodeGenFunction &CGF) {
     return CGF.CXXABIThisValue;
   }
protected:
   virtual bool isThisCompleteObject(GlobalDecl GD) const = 0;
 
 public:
 
   virtual ~CGCXXABI();
 
   /// Gets the mangle context.
   MangleContext &getMangleContext() {
-    return *MangleCtx;
+    return MangleCtx;
   }
 
   /// Returns true if the given constructor or destructor is one of the
   /// kinds that the ABI says returns 'this' (only applies when called
   /// non-virtually for destructors).
   ///
   /// There currently is no way to indicate if a destructor returns 'this'
   /// when called virtually, and code generation does not support the case.