[NFC][CUDA][HIP] rework mangling number for aux target

[yaxunl]

Summary

CUDA/HIP needs to mangle for aux target. When mangling for aux target,
the mangler should use mangling number for aux target. Previously
in https://reviews.llvm.org/D122734 a state was introduced in
ASTContext to let the mangler get mangling number for aux target
from ASTContext. This patch removes that state from ASTConext
and add an IsAux member to MangleContext to indicate that
the mangle context is for aux target. This reflects the reality that
the mangle context is created for mangling aux target and makes
ASTContext cleaner.

[tra]

LGTM.

clang/lib/AST/ASTContext.cpp

11770–11778

Nit: I'd rephrase it as :

if (LangOpts.CUDA && !LangOpts.CUDAIsDevice) {
    Res = ForAuxTarget ? Res >> 16 : Res & 0xFFFF; 
} else {
    assert(!ForAuxTarget && "Only CUDA/HIP host compilation supports mangling number for aux target");
}
return Res > 1 ? Res : 1;

[rnk]

lgtm, thanks!

[yaxunl]

clang/lib/AST/ASTContext.cpp

11770–11778

will do when committing.

clang/include/clang/AST/ASTContext.h

struct CUDAConstantEvalContextRAII {
     CUDAConstantEvalContext SavedCtx;
     CUDAConstantEvalContextRAII(ASTContext &Ctx_, bool NoWrongSidedVars)
         : Ctx(Ctx_), SavedCtx(Ctx_.CUDAConstantEvalCtx) {
       Ctx_.CUDAConstantEvalCtx.NoWrongSidedVars = NoWrongSidedVars;
     }
     ~CUDAConstantEvalContextRAII() { Ctx.CUDAConstantEvalCtx = SavedCtx; }
   };
 
-  /// Current CUDA name mangling is for device name in host compilation.
-  bool CUDAMangleDeviceNameInHostCompilation = false;
-
   /// Returns the dynamic AST node parent map context.
   ParentMapContext &getParentMapContext();
 
   // A traversal scope limits the parts of the AST visible to certain analyses.
   // RecursiveASTVisitor only visits specified children of TranslationUnitDecl.
   // getParents() will only observe reachable parent edges.
   //
   // The scope is defined by a set of "top-level" declarations which will be
public:
 
   TypedefNameDecl *getTypedefNameForUnnamedTagDecl(const TagDecl *TD);
 
   void addDeclaratorForUnnamedTagDecl(TagDecl *TD, DeclaratorDecl *DD);
 
   DeclaratorDecl *getDeclaratorForUnnamedTagDecl(const TagDecl *TD);
 
   void setManglingNumber(const NamedDecl *ND, unsigned Number);
-  unsigned getManglingNumber(const NamedDecl *ND) const;
+  unsigned getManglingNumber(const NamedDecl *ND,
+                             bool ForAuxTarget = false) const;
 
   void setStaticLocalNumber(const VarDecl *VD, unsigned Number);
   unsigned getStaticLocalNumber(const VarDecl *VD) const;
 
   /// Retrieve the context for computing mangling numbers in the given
   /// DeclContext.
   MangleNumberingContext &getManglingNumberContext(const DeclContext *DC);
   enum NeedExtraManglingDecl_t { NeedExtraManglingDecl };

clang/include/clang/AST/Mangle.h

public:
   };
 
 private:
   virtual void anchor();
 
   ASTContext &Context;
   DiagnosticsEngine &Diags;
   const ManglerKind Kind;
+  /// For aux target. If true, uses mangling number for aux target from
+  /// ASTContext.
+  bool IsAux = false;
 
   llvm::DenseMap<const BlockDecl*, unsigned> GlobalBlockIds;
   llvm::DenseMap<const BlockDecl*, unsigned> LocalBlockIds;
   llvm::DenseMap<const NamedDecl*, uint64_t> AnonStructIds;
 
 public:
   ManglerKind getKind() const { return Kind; }
 
-  explicit MangleContext(ASTContext &Context,
-                         DiagnosticsEngine &Diags,
-                         ManglerKind Kind)
-      : Context(Context), Diags(Diags), Kind(Kind) {}
+  bool isAux() const { return IsAux; }
+
+  explicit MangleContext(ASTContext &Context, DiagnosticsEngine &Diags,
+                         ManglerKind Kind, bool IsAux = false)
+      : Context(Context), Diags(Diags), Kind(Kind), IsAux(IsAux) {}
 
   virtual ~MangleContext() { }
 
   ASTContext &getASTContext() const { return Context; }
 
   DiagnosticsEngine &getDiags() const { return Diags; }
 
   virtual void startNewFunction() { LocalBlockIds.clear(); }
public:
 
   /// @}
 };
 
 class ItaniumMangleContext : public MangleContext {
 public:
   using DiscriminatorOverrideTy =
       llvm::Optional<unsigned> (*)(ASTContext &, const NamedDecl *);
-  explicit ItaniumMangleContext(ASTContext &C, DiagnosticsEngine &D)
-      : MangleContext(C, D, MK_Itanium) {}
+  explicit ItaniumMangleContext(ASTContext &C, DiagnosticsEngine &D,
+                                bool IsAux = false)
+      : MangleContext(C, D, MK_Itanium, IsAux) {}
 
   virtual void mangleCXXVTable(const CXXRecordDecl *RD, raw_ostream &) = 0;
   virtual void mangleCXXVTT(const CXXRecordDecl *RD, raw_ostream &) = 0;
   virtual void mangleCXXCtorVTable(const CXXRecordDecl *RD, int64_t Offset,
                                    const CXXRecordDecl *Type,
                                    raw_ostream &) = 0;
   virtual void mangleItaniumThreadLocalInit(const VarDecl *D,
                                             raw_ostream &) = 0;
public:
 
   // This has to live here, otherwise the CXXNameMangler won't have access to
   // it.
   virtual DiscriminatorOverrideTy getDiscriminatorOverride() const = 0;
   static bool classof(const MangleContext *C) {
     return C->getKind() == MK_Itanium;
   }
 
-  static ItaniumMangleContext *create(ASTContext &Context,
-                                      DiagnosticsEngine &Diags);
+  static ItaniumMangleContext *
+  create(ASTContext &Context, DiagnosticsEngine &Diags, bool IsAux = false);
   static ItaniumMangleContext *create(ASTContext &Context,
                                       DiagnosticsEngine &Diags,
-                                      DiscriminatorOverrideTy Discriminator);
+                                      DiscriminatorOverrideTy Discriminator,
+                                      bool IsAux = false);
 };
 
 class MicrosoftMangleContext : public MangleContext {
 public:
-  explicit MicrosoftMangleContext(ASTContext &C, DiagnosticsEngine &D)
-      : MangleContext(C, D, MK_Microsoft) {}
+  explicit MicrosoftMangleContext(ASTContext &C, DiagnosticsEngine &D,
+                                  bool IsAux = false)
+      : MangleContext(C, D, MK_Microsoft, IsAux) {}
 
   /// Mangle vftable symbols.  Only a subset of the bases along the path
   /// to the vftable are included in the name.  It's up to the caller to pick
   /// them correctly.
   virtual void mangleCXXVFTable(const CXXRecordDecl *Derived,
                                 ArrayRef<const CXXRecordDecl *> BasePath,
                                 raw_ostream &Out) = 0;
 
public:
   mangleCXXRTTICompleteObjectLocator(const CXXRecordDecl *Derived,
                                      ArrayRef<const CXXRecordDecl *> BasePath,
                                      raw_ostream &Out) = 0;
 
   static bool classof(const MangleContext *C) {
     return C->getKind() == MK_Microsoft;
   }
 
-  static MicrosoftMangleContext *create(ASTContext &Context,
-                                        DiagnosticsEngine &Diags);
+  static MicrosoftMangleContext *
+  create(ASTContext &Context, DiagnosticsEngine &Diags, bool IsAux = false);
 };
 
 class ASTNameGenerator {
 public:
   explicit ASTNameGenerator(ASTContext &Ctx);
   ~ASTNameGenerator();
 
   /// Writes name for \p D to \p OS.

clang/lib/AST/ASTContext.cpp

MangleContext *ASTContext::createDeviceMangleContext(const TargetInfo &T) {
   case TargetCXXABI::WatchOS:
   case TargetCXXABI::XL:
     return ItaniumMangleContext::create(
         *this, getDiagnostics(),
         [](ASTContext &, const NamedDecl *ND) -> llvm::Optional<unsigned> {
           if (const auto *RD = dyn_cast<CXXRecordDecl>(ND))
             return RD->getDeviceLambdaManglingNumber();
           return llvm::None;
-        });
+        },
+        /*IsAux=*/true);
   case TargetCXXABI::Microsoft:
-    return MicrosoftMangleContext::create(*this, getDiagnostics());
+    return MicrosoftMangleContext::create(*this, getDiagnostics(),
+                                          /*IsAux=*/true);
   }
   llvm_unreachable("Unsupported ABI");
 }
 
 CXXABI::~CXXABI() = default;
 
 size_t ASTContext::getSideTableAllocatedMemory() const {
   return ASTRecordLayouts.getMemorySize() +
QualType ASTContext::getRealTypeForBitwidth(unsigned DestWidth,
   llvm_unreachable("Unhandled TargetInfo::RealType value");
 }
 
 void ASTContext::setManglingNumber(const NamedDecl *ND, unsigned Number) {
   if (Number > 1)
     MangleNumbers[ND] = Number;
 }
 
-unsigned ASTContext::getManglingNumber(const NamedDecl *ND) const {
+unsigned ASTContext::getManglingNumber(const NamedDecl *ND,
+                                       bool ForAuxTarget) const {
   auto I = MangleNumbers.find(ND);
   unsigned Res = I != MangleNumbers.end() ? I->second : 1;
-  if (!LangOpts.CUDA || LangOpts.CUDAIsDevice)
-    return Res;
-
   // CUDA/HIP host compilation encodes host and device mangling numbers
   // as lower and upper half of 32 bit integer.
-  Res = CUDAMangleDeviceNameInHostCompilation ? Res >> 16 : Res & 0xFFFF;
+  if (LangOpts.CUDA && !LangOpts.CUDAIsDevice) {
+    Res = ForAuxTarget ? Res >> 16 : Res & 0xFFFF;
+  } else {
+    assert(!ForAuxTarget && "Only CUDA/HIP host compilation supports mangling "
+                            "number for aux target");
+  }
   return Res > 1 ? Res : 1;
 }
 
 void ASTContext::setStaticLocalNumber(const VarDecl *VD, unsigned Number) {
   if (Number > 1)
     StaticLocalNumbers[VD] = Number;
 }
 
 unsigned ASTContext::getStaticLocalNumber(const VarDecl *VD) const {
   auto I = StaticLocalNumbers.find(VD);

clang/lib/AST/ItaniumMangle.cpp

class ItaniumMangleContextImpl : public ItaniumMangleContext {
   const DiscriminatorOverrideTy DiscriminatorOverride = nullptr;
   NamespaceDecl *StdNamespace = nullptr;
 
   bool NeedsUniqueInternalLinkageNames = false;
 
 public:
   explicit ItaniumMangleContextImpl(
       ASTContext &Context, DiagnosticsEngine &Diags,
-      DiscriminatorOverrideTy DiscriminatorOverride)
-      : ItaniumMangleContext(Context, Diags),
+      DiscriminatorOverrideTy DiscriminatorOverride, bool IsAux = false)
+      : ItaniumMangleContext(Context, Diags, IsAux),
         DiscriminatorOverride(DiscriminatorOverride) {}
 
   /// @name Mangler Entry Points
   /// @{
 
   bool shouldMangleCXXName(const NamedDecl *D) override;
   bool shouldMangleStringLiteral(const StringLiteral *) override {
     return false;
bool getNextDiscriminator(const NamedDecl *ND, unsigned &disc) {
     // Anonymous tags are already numbered.
     if (const TagDecl *Tag = dyn_cast<TagDecl>(ND)) {
       if (Tag->getName().empty() && !Tag->getTypedefNameForAnonDecl())
         return false;
     }
 
     // Use the canonical number for externally visible decls.
     if (ND->isExternallyVisible()) {
-      unsigned discriminator = getASTContext().getManglingNumber(ND);
+      unsigned discriminator = getASTContext().getManglingNumber(ND, isAux());
       if (discriminator == 1)
         return false;
       disc = discriminator - 2;
       return true;
     }
 
     // Make up a reasonable number for internal decls.
     unsigned &discriminator = Uniquifier[ND];
if (const CXXRecordDecl *Record = dyn_cast<CXXRecordDecl>(TD)) {
         assert(!AdditionalAbiTags &&
                "Lambda type cannot have additional abi tags");
         mangleLambda(Record);
         break;
       }
     }
 
     if (TD->isExternallyVisible()) {
-      unsigned UnnamedMangle = getASTContext().getManglingNumber(TD);
+      unsigned UnnamedMangle =
+          getASTContext().getManglingNumber(TD, Context.isAux());
       Out << "Ut";
       if (UnnamedMangle > 1)
         Out << UnnamedMangle - 2;
       Out << '_';
       writeAbiTags(TD, AdditionalAbiTags);
       break;
     }
 

 
 void ItaniumMangleContextImpl::mangleLambdaSig(const CXXRecordDecl *Lambda,
                                                raw_ostream &Out) {
   CXXNameMangler Mangler(*this, Out);
   Mangler.mangleLambdaSig(Lambda);
 }
 
 ItaniumMangleContext *ItaniumMangleContext::create(ASTContext &Context,
-                                                   DiagnosticsEngine &Diags) {
+                                                   DiagnosticsEngine &Diags,
+                                                   bool IsAux) {
   return new ItaniumMangleContextImpl(
       Context, Diags,
       [](ASTContext &, const NamedDecl *) -> llvm::Optional<unsigned> {
         return llvm::None;
-      });
+      },
+      IsAux);
 }
 
 ItaniumMangleContext *
 ItaniumMangleContext::create(ASTContext &Context, DiagnosticsEngine &Diags,
-                             DiscriminatorOverrideTy DiscriminatorOverride) {
-  return new ItaniumMangleContextImpl(Context, Diags, DiscriminatorOverride);
+                             DiscriminatorOverrideTy DiscriminatorOverride,
+                             bool IsAux) {
+  return new ItaniumMangleContextImpl(Context, Diags, DiscriminatorOverride,
+                                      IsAux);
 }

clang/lib/AST/MicrosoftMangle.cpp

class MicrosoftMangleContextImpl : public MicrosoftMangleContext {
   llvm::DenseMap<DiscriminatorKeyTy, unsigned> Discriminator;
   llvm::DenseMap<const NamedDecl *, unsigned> Uniquifier;
   llvm::DenseMap<const CXXRecordDecl *, unsigned> LambdaIds;
   llvm::DenseMap<const NamedDecl *, unsigned> SEHFilterIds;
   llvm::DenseMap<const NamedDecl *, unsigned> SEHFinallyIds;
   SmallString<16> AnonymousNamespaceHash;
 
 public:
-  MicrosoftMangleContextImpl(ASTContext &Context, DiagnosticsEngine &Diags);
+  MicrosoftMangleContextImpl(ASTContext &Context, DiagnosticsEngine &Diags,
+                             bool IsAux = false);
   bool shouldMangleCXXName(const NamedDecl *D) override;
   bool shouldMangleStringLiteral(const StringLiteral *SL) override;
   void mangleCXXName(GlobalDecl GD, raw_ostream &Out) override;
   void mangleVirtualMemPtrThunk(const CXXMethodDecl *MD,
                                 const MethodVFTableLocation &ML,
                                 raw_ostream &Out) override;
   void mangleThunk(const CXXMethodDecl *MD, const ThunkInfo &Thunk,
                    raw_ostream &) override;
if (const auto *RD = dyn_cast<CXXRecordDecl>(ND)) {
       if (RD->isLambda()) {
         disc = 1;
         return true;
       }
     }
 
     // Use the canonical number for externally visible decls.
     if (ND->isExternallyVisible()) {
-      disc = getASTContext().getManglingNumber(ND);
+      disc = getASTContext().getManglingNumber(ND, isAux());
       return true;
     }
 
     // Anonymous tags are already numbered.
     if (const TagDecl *Tag = dyn_cast<TagDecl>(ND)) {
       if (!Tag->hasNameForLinkage() &&
           !getASTContext().getDeclaratorForUnnamedTagDecl(Tag) &&
           !getASTContext().getTypedefNameForUnnamedTagDecl(Tag))
#undef TYPE
   void mangleObjCLifetime(const QualType T, Qualifiers Quals,
                           SourceRange Range);
   void mangleObjCKindOfType(const ObjCObjectType *T, Qualifiers Quals,
                             SourceRange Range);
 };
 }
 
 MicrosoftMangleContextImpl::MicrosoftMangleContextImpl(ASTContext &Context,
-                                                       DiagnosticsEngine &Diags)
-    : MicrosoftMangleContext(Context, Diags) {
+                                                       DiagnosticsEngine &Diags,
+                                                       bool IsAux)
+    : MicrosoftMangleContext(Context, Diags, IsAux) {
   // To mangle anonymous namespaces, hash the path to the main source file. The
   // path should be whatever (probably relative) path was passed on the command
   // line. The goal is for the compiler to produce the same output regardless of
   // working directory, so use the uncanonicalized relative path.
   //
   // It's important to make the mangled names unique because, when CodeView
   // debug info is in use, the debugger uses mangled type names to distinguish
   // between otherwise identically named types in anonymous namespaces.
if (SL->isWide())
       MangleByte(GetBigEndianByte(I));
     else
       MangleByte(GetLittleEndianByte(I));
   }
 
   Mangler.getStream() << '@';
 }
 
-MicrosoftMangleContext *
-MicrosoftMangleContext::create(ASTContext &Context, DiagnosticsEngine &Diags) {
-  return new MicrosoftMangleContextImpl(Context, Diags);
+MicrosoftMangleContext *MicrosoftMangleContext::create(ASTContext &Context,
+                                                       DiagnosticsEngine &Diags,
+                                                       bool IsAux) {
+  return new MicrosoftMangleContextImpl(Context, Diags, IsAux);
 }

clang/lib/CodeGen/CGCUDANV.cpp

 #include "clang/Basic/Cuda.h"
 #include "clang/CodeGen/CodeGenABITypes.h"
 #include "clang/CodeGen/ConstantInitBuilder.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/ReplaceConstant.h"
 #include "llvm/Support/Format.h"
-#include "llvm/Support/SaveAndRestore.h"
 
 using namespace clang;
 using namespace CodeGen;
 
 namespace {
 constexpr unsigned CudaFatMagic = 0x466243b1;
 constexpr unsigned HIPFatMagic = 0x48495046; // "HIPF"
 
llvm::FunctionType *CGNVCUDARuntime::getRegisterLinkedBinaryFnTy() const {
   auto *CallbackFnTy = getCallbackFnTy();
   auto *RegisterGlobalsFnTy = getRegisterGlobalsFnTy();
   llvm::Type *Params[] = {RegisterGlobalsFnTy->getPointerTo(), VoidPtrTy,
                           VoidPtrTy, CallbackFnTy->getPointerTo()};
   return llvm::FunctionType::get(VoidTy, Params, false);
 }
 
 std::string CGNVCUDARuntime::getDeviceSideName(const NamedDecl *ND) {
-  llvm::SaveAndRestore<bool> MangleAsDevice(
-      CGM.getContext().CUDAMangleDeviceNameInHostCompilation, true);
   GlobalDecl GD;
   // D 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;