[HeterogeneousDwarf] Keep track of GlobalVariables DIFragments using WeakVH

In DwarfDebug::beginModule, the compiler stores for each fragment its
associated GlobalVariable in a DenseMap.
However, nothing prevents the GlobalVariable from being removed during a later pass,
invalidating the reference in the map.
In the attached test case, AMDGPULowerLDS removes LDS variables without any use,
which resulted in an use-after-free access.

The WeakVH works as a reference to the GlobalVariable, however, when the
global is erased a callback sets the WeakVH to invalid.

Change-Id: If1627f548441f07ab32454ece2beefb518ab2aa2

Author: jmmartinez

llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp

@@ -204,7 +204,7 @@ DIE *DwarfCompileUnit::getOrCreateGlobalVariableDIE(
 
 DIE *DwarfCompileUnit::getOrCreateGlobalVariableDIE(
     const DILifetime &Lifetime,
-    const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap) {
+    const DwarfDebug::GVFragmentMapTy &GVFragmentMap) {
 
   const DIGlobalVariable *GV = dyn_cast<DIGlobalVariable>(Lifetime.getObject());
 
@@ -462,7 +462,7 @@ void DwarfCompileUnit::addLocationAttribute(
 
 void DwarfCompileUnit::addLocationAttribute(
     DIE *VariableDIE, const DIGlobalVariable *GV, const DILifetime &Lifetime,
-    const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap) {
+    const DwarfDebug::GVFragmentMapTy &GVFragmentMap) {
   // FIXME: Determine when this is appropriate, considering the existing
   // implementation.
   bool AddToAccelTable = true;
@@ -505,7 +505,7 @@ DIE *DwarfCompileUnit::getOrCreateCommonBlock(
 
 DIE *DwarfCompileUnit::getOrCreateCommonBlock(
     const DICommonBlock *CB, const DILifetime &Lifetime,
-    const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap) {
+    const DwarfDebug::GVFragmentMapTy &GVFragmentMap) {
   // Check for pre-existence.
   if (DIE *NDie = getDIE(CB))
     return NDie;

llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.h

@@ -165,21 +165,21 @@ class DwarfCompileUnit final : public DwarfUnit {
 
   DIE *getOrCreateGlobalVariableDIE(
       const DILifetime &Lifetime,
-      const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap);
+      const DwarfDebug::GVFragmentMapTy &GVFragmentMap);
 
   DIE *getOrCreateCommonBlock(const DICommonBlock *CB,
                               ArrayRef<GlobalExpr> GlobalExprs);
 
-  DIE *getOrCreateCommonBlock(
-      const DICommonBlock *CB, const DILifetime &Lifetime,
-      const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap);
+  DIE *getOrCreateCommonBlock(const DICommonBlock *CB,
+                              const DILifetime &Lifetime,
+                              const DwarfDebug::GVFragmentMapTy &GVFragmentMap);
 
   void addLocationAttribute(DIE *ToDIE, const DIGlobalVariable *GV,
                             ArrayRef<GlobalExpr> GlobalExprs);
 
-  void addLocationAttribute(
-      DIE *ToDIE, const DIGlobalVariable *GV, const DILifetime &Lifetime,
-      const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap);
+  void addLocationAttribute(DIE *ToDIE, const DIGlobalVariable *GV,
+                            const DILifetime &Lifetime,
+                            const DwarfDebug::GVFragmentMapTy &GVFragmentMap);
 
   /// addLabelAddress - Add a dwarf label attribute data and value using
   /// either DW_FORM_addr or DW_FORM_GNU_addr_index.

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -1177,7 +1177,7 @@ void DwarfDebug::beginModule(Module *M) {
   SingleCU = NumDebugCUs == 1;
   DenseMap<DIGlobalVariable *, SmallVector<DwarfCompileUnit::GlobalExpr, 1>>
       GVMap;
-  for (const GlobalVariable &Global : M->globals()) {
+  for (GlobalVariable &Global : M->globals()) {
     // To support the "inlining" of GV-fragments as an optimization, we record
     // the referrer for each such fragment.
     if (llvm::isHeterogeneousDebug(*M)) {
@@ -2003,10 +2003,10 @@ void DwarfDebug::collectEntityInfo(DwarfCompileUnit &TheCU,
       Entry.finalize(*Asm, List, BT, TheCU);
   }
 
-  const Module &M = *Asm->MF->getFunction().getParent();
-      DenseMap<DICompileUnit *, SmallVector<DILifetime *>> AddCULifetimeMap;
+  Module &M = *Asm->MF->getFunction().getParent();
+  DenseMap<DICompileUnit *, SmallVector<DILifetime *>> AddCULifetimeMap;
   if (isHeterogeneousDebug(M)) {
-    for (const GlobalVariable &Global : M.globals()) {
+    for (GlobalVariable &Global : M.globals()) {
       if (DIFragment *F = Global.getDbgDef()) {
         GVFragmentMap[F] = &Global;
       }

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.h

@@ -492,7 +492,11 @@ class DwarfDebug : public DebugHandlerBase {
   /// Avoid using DW_OP_convert due to consumer incompatibilities.
   bool EnableOpConvert;
 
-  DenseMap<DIFragment *, const GlobalVariable *> GVFragmentMap;
+public:
+  using GVFragmentMapTy = DenseMap<DIFragment *, WeakVH>;
+
+private:
+  GVFragmentMapTy GVFragmentMap;
   DenseMap<DISubprogram *, SmallVector<DILifetime *>> SPLifetimeMap;
   DenseSet<DILifetime *> ProcessedLifetimes;
 

llvm/lib/CodeGen/AsmPrinter/DwarfExpression.cpp

@@ -842,10 +842,10 @@ bool DwarfExprAST::tryInlineArgObject(DIObject *ArgObject) {
   auto *Fragment = dyn_cast<DIFragment>(ArgObject);
   if (!Fragment)
     return false;
-  const auto GV = GVFragmentMap->find(Fragment);
-  if (GV == GVFragmentMap->end())
+  const auto GVRef = GVFragmentMap->find(Fragment);
+  if (GVRef == GVFragmentMap->end() || !GVRef->getSecond())
     return false;
-  const GlobalVariable *Global = GV->getSecond();
+  const GlobalVariable *Global = cast<GlobalVariable>(GVRef->getSecond());
   // FIXME(KZHURAVL): This depends on the target and address space
   // semantics. For AMDGPU, address space 3 is lds/local/shared.
   // Need to replace this with a target hook!

llvm/lib/CodeGen/AsmPrinter/DwarfExpression.h

@@ -14,6 +14,7 @@
 #define LLVM_LIB_CODEGEN_ASMPRINTER_DWARFEXPRESSION_H
 
 #include "ByteStreamer.h"
+#include "DwarfDebug.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/DebugInfoMetadata.h"
@@ -513,7 +514,7 @@ class DwarfExprAST {
   // An `std::optional<const DenseMap<_, _>&>` where `nullptr` represents
   // `None`. Only present and applicable as part of an optimization for
   // DIFragments which refer to global variable fragments.
-  const DenseMap<DIFragment *, const GlobalVariable *> *GVFragmentMap;
+  const DwarfDebug::GVFragmentMapTy *GVFragmentMap;
   std::unique_ptr<DwarfExprAST::Node> Root;
   // FIXME(KZHURAVL): This is a temporary boolean variable that indicates
   // whether the lowering of this expression is supported or not. If the
@@ -586,10 +587,10 @@ class DwarfExprAST {
   virtual void emitDwarfLabelDelta(const MCSymbol *Hi, const MCSymbol *Lo) = 0;
 
 public:
-  DwarfExprAST(
-      const AsmPrinter &AP, const TargetRegisterInfo *TRI, DwarfCompileUnit &CU,
-      const DILifetime &Lifetime, const MachineOperand *Referrer,
-      const DenseMap<DIFragment *, const GlobalVariable *> *GVFragmentMap)
+  DwarfExprAST(const AsmPrinter &AP, const TargetRegisterInfo *TRI,
+               DwarfCompileUnit &CU, const DILifetime &Lifetime,
+               const MachineOperand *Referrer,
+               const DwarfDebug::GVFragmentMapTy *GVFragmentMap)
       : AP(AP), TRI(TRI), CU(CU), Lifetime(Lifetime), Referrer(Referrer),
         GVFragmentMap(GVFragmentMap) {
     buildDIExprAST();
@@ -610,24 +611,23 @@ class DebugLocDwarfExprAST final : DwarfExprAST {
   void emitDwarfOpAddrx(unsigned Index) override;
   void emitDwarfLabelDelta(const MCSymbol *Hi, const MCSymbol *Lo) override;
 
-  DebugLocDwarfExprAST(
-      const AsmPrinter &AP, const TargetRegisterInfo *TRI, DwarfCompileUnit &CU,
-      BufferByteStreamer &BS, const DILifetime &Lifetime,
-      const MachineOperand *Referrer,
-      const DenseMap<DIFragment *, const GlobalVariable *> *GVFragmentMap)
+  DebugLocDwarfExprAST(const AsmPrinter &AP, const TargetRegisterInfo *TRI,
+                       DwarfCompileUnit &CU, BufferByteStreamer &BS,
+                       const DILifetime &Lifetime,
+                       const MachineOperand *Referrer,
+                       const DwarfDebug::GVFragmentMapTy *GVFragmentMap)
       : DwarfExprAST(AP, TRI, CU, Lifetime, Referrer, GVFragmentMap),
         OutBS(BS) {}
 
-  public:
+public:
   DebugLocDwarfExprAST(const AsmPrinter &AP, const TargetRegisterInfo &TRI,
                        DwarfCompileUnit &CU, BufferByteStreamer &BS,
                        const DILifetime &Lifetime,
                        const MachineOperand &Referrer)
       : DebugLocDwarfExprAST(AP, &TRI, CU, BS, Lifetime, &Referrer, nullptr) {}
-  DebugLocDwarfExprAST(
-      const AsmPrinter &AP, DwarfCompileUnit &CU, BufferByteStreamer &BS,
-      const DILifetime &Lifetime,
-      const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap)
+  DebugLocDwarfExprAST(const AsmPrinter &AP, DwarfCompileUnit &CU,
+                       BufferByteStreamer &BS, const DILifetime &Lifetime,
+                       const DwarfDebug::GVFragmentMapTy &GVFragmentMap)
       : DebugLocDwarfExprAST(AP, nullptr, CU, BS, Lifetime, nullptr,
                              &GVFragmentMap) {}
   DebugLocDwarfExprAST(const DebugLocDwarfExprAST &) = delete;
@@ -653,10 +653,10 @@ class DIEDwarfExprAST final : DwarfExprAST {
   void emitDwarfOpAddrx(unsigned Index) override;
   void emitDwarfLabelDelta(const MCSymbol *Hi, const MCSymbol *Lo) override;
 
-  DIEDwarfExprAST(
-      const AsmPrinter &AP, const TargetRegisterInfo *TRI, DwarfCompileUnit &CU,
-      DIELoc &DIE, const DILifetime &Lifetime, const MachineOperand *Referrer,
-      const DenseMap<DIFragment *, const GlobalVariable *> *GVFragmentMap)
+  DIEDwarfExprAST(const AsmPrinter &AP, const TargetRegisterInfo *TRI,
+                  DwarfCompileUnit &CU, DIELoc &DIE, const DILifetime &Lifetime,
+                  const MachineOperand *Referrer,
+                  const DwarfDebug::GVFragmentMapTy *GVFragmentMap)
       : DwarfExprAST(AP, TRI, CU, Lifetime, Referrer, GVFragmentMap),
         OutDIE(DIE) {}
 
@@ -665,10 +665,9 @@ class DIEDwarfExprAST final : DwarfExprAST {
                   DwarfCompileUnit &CU, DIELoc &DIE, const DILifetime &Lifetime,
                   const MachineOperand &Referrer)
       : DIEDwarfExprAST(AP, &TRI, CU, DIE, Lifetime, &Referrer, nullptr) {}
-  DIEDwarfExprAST(
-      const AsmPrinter &AP, DwarfCompileUnit &CU, DIELoc &DIE,
-      const DILifetime &Lifetime,
-      const DenseMap<DIFragment *, const GlobalVariable *> &GVFragmentMap)
+  DIEDwarfExprAST(const AsmPrinter &AP, DwarfCompileUnit &CU, DIELoc &DIE,
+                  const DILifetime &Lifetime,
+                  const DwarfDebug::GVFragmentMapTy &GVFragmentMap)
       : DIEDwarfExprAST(AP, nullptr, CU, DIE, Lifetime, nullptr,
                         &GVFragmentMap) {}
   DIEDwarfExprAST(const DIEDwarfExprAST &) = delete;

llvm/test/CodeGen/AMDGPU/heterogeneous-debug-info-lower-module-lds-unused-variable.ll

@@ -0,0 +1,35 @@
+; RUN: llc -O0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx906 -verify-machineinstrs -filetype=obj %s -o - | llvm-dwarfdump -v -debug-info - | FileCheck %s
+; CHECK: DW_TAG_compile_unit
+; CHECK:   DW_TAG_subprogram
+; CHECK:     DW_TAG_variable
+; CHECK:       DW_AT_name [DW_FORM_strx1] (indexed ({{[0-9]+}}) string = "unused")
+; CHECK:       DW_AT_LLVM_memory_space [DW_FORM_data4] (DW_MSPACE_LLVM_group)
+; CHECK:       DW_AT_location [DW_FORM_exprloc] (<empty>)
+; CHECK-NOT: DW_TAG_subprogram
+; CHECK-NOT: DW_TAG_variable
+
+@unused = external addrspace(3) global [4 x float], !dbg.def !0
+
+define amdgpu_kernel void @foo() !dbg !1 {
+  ret void
+}
+
+!llvm.module.flags = !{!11, !12}
+!llvm.dbg.cu = !{!2}
+!llvm.dbg.retainedNodes = !{!5}
+
+!0 = distinct !DIFragment()
+!1 = distinct !DISubprogram(name: "foo", linkageName: "foo", scope: !2, file: !3, line: 1, type: !13, scopeLine: 1, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !2)
+!2 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !3, producer: "clang", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !4, retainedTypes: !4, imports: !4, splitDebugInlining: false, nameTableKind: None)
+!3 = !DIFile(filename: "file.cpp", directory: "/dir")
+!4 = !{}
+!5 = distinct !DILifetime(object: !6, location: !DIExpr(DIOpArg(0, ptr addrspace(3)), DIOpDeref([4 x float])), argObjects: {!0})
+!6 = distinct !DIGlobalVariable(name: "unused", scope: !1, file: !3, line: 1, type: !7, isLocal: false, isDefinition: true, memorySpace: DW_MSPACE_LLVM_group)
+!7 = !DICompositeType(tag: DW_TAG_array_type, baseType: !8, size: 128, elements: !9)
+!8 = !DIBasicType(name: "float", size: 32, encoding: DW_ATE_float)
+!9 = !{!10}
+!10 = !DISubrange(count: 4)
+!11 = !{i32 7, !"Dwarf Version", i32 5}
+!12 = !{i32 2, !"Debug Info Version", i32 4}
+!13 = distinct !DISubroutineType(types: !14)
+!14 = !{null}