Support stripping indirectly referenced DILocations from !llvm.loop metadata in stripDebugInfo()

[aprantl]

Summary

This patch fixes an oversight in https://reviews.llvm.org/D96181 and also takes into account loop metadata pointing to other MDNodes that point into the debug info.

rdar://78487175

[JDevlieghere]

llvm/lib/IR/DebugInfo.cpp

389–390

Do we need to remember that the MDNode was visited when it's a DILocation?

411–412

[shafik]

llvm/lib/IR/DebugInfo.cpp

391

This is why I dislike std::pair what does second mean? A comment like if it was already inserted exit now would be helpful.

411–412

I see in the other files the lambda explicitly list their captures instead of just using & to capture all by reference, I think that is a better style to be explicit.

llvm/lib/Transforms/Utils/InlineFunction.cpp

1523

I see that previously Loc was const can MD also be const?

[aprantl]

llvm/lib/Transforms/Utils/InlineFunction.cpp

1523

I don't think so because we need to return it as a Metadata * and the input may be nullptr.

[aprantl]

llvm/lib/IR/DebugInfo.cpp

389–390

In that case it will be in Reachable so there's no need to.

411–412

For a one-line lambda like this, I feel like it just adds visual noise and makes it harder to read:

if (!std::count_if(N->op_begin() + 1, N->op_end(),
 [&Visited, &DILocationReachable](const MDOperand &Op) {
   return isa<DILocation>(Op.get()) ||
             isDILocationReachable(Visited, DILocationReachable, Op.get());
    }))

[aprantl]

Address review feedback.

[aprantl]

Add another comment and explicitly mark lambda captures.

[aprantl]

Add more comments

[shafik]

LGTM

[vsk]

Thanks Adrian. + 1 from me as well, FTR.

llvm/lib/IR/DebugInfo.cpp

418

Is the isa<DILocation>(Op.get()) part redundant here, since it should be checked within isDILocationReachable?

[aprantl]

llvm/lib/IR/DebugInfo.cpp

418

Yes. I removed them in fcfaed4ae6d7. Thanks!

llvm/include/llvm/IR/DebugInfo.h

 ///   2) Delete all non-CU named metadata debug info nodes.
 ///   3) Create new DebugLocs for each instruction.
 ///   4) Create a new CU debug info, and similarly for every metadata node
 ///      that's reachable from the CU debug info.
 ///   All debug type metadata nodes are unreachable and garbage collected.
 bool stripNonLineTableDebugInfo(Module &M);
 
 /// Update the debug locations contained within the MD_loop metadata attached
-/// to the instruction \p I, if one exists. \p Updater is applied to each debug
-/// location in the MD_loop metadata: the returned value is included in the
+/// to the instruction \p I, if one exists. \p Updater is applied to Metadata
+/// operand in the MD_loop metadata: the returned value is included in the
 /// updated loop metadata node if it is non-null.
 void updateLoopMetadataDebugLocations(
-    Instruction &I, function_ref<DILocation *(const DILocation &)> Updater);
+    Instruction &I, function_ref<Metadata *(Metadata *)> Updater);
 
 /// Return Debug Info Metadata Version by checking module flags.
 unsigned getDebugMetadataVersionFromModule(const Module &M);
 
 /// Utility to find all debug info in a module.
 ///
 /// DebugInfoFinder tries to list all debug info MDNodes used in a module. To
 /// list debug info MDNodes used by an instruction, DebugInfoFinder uses

llvm/lib/IR/DebugInfo.cpp

if (Scope->getNumOperands() == 0)
     return false;
   if (!NodesSeen.insert(Scope).second)
     return false;
   Scopes.push_back(Scope);
   return true;
 }
 
 static MDNode *updateLoopMetadataDebugLocationsImpl(
-    MDNode *OrigLoopID,
-    function_ref<DILocation *(const DILocation &)> Updater) {
+    MDNode *OrigLoopID, function_ref<Metadata *(Metadata *)> Updater) {
   assert(OrigLoopID && OrigLoopID->getNumOperands() > 0 &&
          "Loop ID needs at least one operand");
   assert(OrigLoopID && OrigLoopID->getOperand(0).get() == OrigLoopID &&
          "Loop ID should refer to itself");
 
   // Save space for the self-referential LoopID.
   SmallVector<Metadata *, 4> MDs = {nullptr};
 
   for (unsigned i = 1; i < OrigLoopID->getNumOperands(); ++i) {
     Metadata *MD = OrigLoopID->getOperand(i);
-    if (DILocation *DL = dyn_cast<DILocation>(MD)) {
-      if (DILocation *NewDL = Updater(*DL))
-        MDs.push_back(NewDL);
-    } else
-      MDs.push_back(MD);
+    if (!MD)
+      MDs.push_back(nullptr);
+    else if (Metadata *NewMD = Updater(MD))
+      MDs.push_back(NewMD);
   }
 
   MDNode *NewLoopID = MDNode::getDistinct(OrigLoopID->getContext(), MDs);
   // Insert the self-referential LoopID.
   NewLoopID->replaceOperandWith(0, NewLoopID);
   return NewLoopID;
 }
 
 void llvm::updateLoopMetadataDebugLocations(
-    Instruction &I, function_ref<DILocation *(const DILocation &)> Updater) {
+    Instruction &I, function_ref<Metadata *(Metadata *)> Updater) {
   MDNode *OrigLoopID = I.getMetadata(LLVMContext::MD_loop);
   if (!OrigLoopID)
     return;
   MDNode *NewLoopID = updateLoopMetadataDebugLocationsImpl(OrigLoopID, Updater);
   I.setMetadata(LLVMContext::MD_loop, NewLoopID);
 }
 
+/// Return true if a node is a DILocation or if a DILocation is
+/// indirectly referenced by one of the node's children.
+static bool isDILocationReachable(SmallPtrSetImpl<Metadata *> &Visited,
+                                  SmallPtrSetImpl<Metadata *> &Reachable,
+                                  Metadata *MD) {
+  MDNode *N = dyn_cast_or_null<MDNode>(MD);
+  if (!N)
+    return false;
+  if (Reachable.count(N) || isa<DILocation>(N))
+    return true;
+  if (!Visited.insert(N).second)
+    return false;
+  for (auto &OpIt : N->operands()) {
+    Metadata *Op = OpIt.get();
+    if (isDILocationReachable(Visited, Reachable, Op)) {
+      Reachable.insert(N);
+      return true;
+    }
+  }
+  return false;
+}
+
 static MDNode *stripDebugLocFromLoopID(MDNode *N) {
   assert(!N->operands().empty() && "Missing self reference?");
+  SmallPtrSet<Metadata *, 8> Visited, DILocationReachable;
+  // If we already visited N, there is nothing to do.
+  if (!Visited.insert(N).second)
+    return N;
 
-  // if there is no debug location, we do not have to rewrite this MDNode.
-  if (std::none_of(N->op_begin() + 1, N->op_end(), [](const MDOperand &Op) {
-        return isa<DILocation>(Op.get());
+  // If there is no debug location, we do not have to rewrite this
+  // MDNode. This loop also initializes DILocationReachable, later
+  // needed by updateLoopMetadataDebugLocationsImpl; the use of
+  // count_if avoids an early exit.
+  if (!std::count_if(N->op_begin() + 1, N->op_end(),
+                     [&Visited, &DILocationReachable](const MDOperand &Op) {
+                       return isa<DILocation>(Op.get()) ||
+                              isDILocationReachable(
+                                  Visited, DILocationReachable, Op.get());
                      }))
     return N;
 
   // If there is only the debug location without any actual loop metadata, we
   // can remove the metadata.
-  if (std::none_of(N->op_begin() + 1, N->op_end(), [](const MDOperand &Op) {
-        return !isa<DILocation>(Op.get());
+  if (std::all_of(
+          N->op_begin() + 1, N->op_end(),
+          [&Visited, &DILocationReachable](const MDOperand &Op) {
+            return isa<DILocation>(Op.get()) ||
+                   isDILocationReachable(Visited, DILocationReachable,
+                                         Op.get());
           }))
     return nullptr;
 
-  auto dropDebugLoc = [](const DILocation &) -> DILocation * {
+  return updateLoopMetadataDebugLocationsImpl(
+      N, [&DILocationReachable](Metadata *MD) -> Metadata * {
+        if (isa<DILocation>(MD) || DILocationReachable.count(MD))
           return nullptr;
-  };
-  return updateLoopMetadataDebugLocationsImpl(N, dropDebugLoc);
+        return MD;
+      });
 }
 
 bool llvm::stripDebugInfo(Function &F) {
   bool Changed = false;
   if (F.hasMetadata(LLVMContext::MD_dbg)) {
     Changed = true;
     F.setSubprogram(nullptr);
   }
 
-  DenseMap<MDNode*, MDNode*> LoopIDsMap;
+  DenseMap<MDNode *, MDNode *> LoopIDsMap;
   for (BasicBlock &BB : F) {
     for (auto II = BB.begin(), End = BB.end(); II != End;) {
       Instruction &I = *II++; // We may delete the instruction, increment now.
       if (isa<DbgInfoIntrinsic>(&I)) {
         I.eraseFromParent();
         Changed = true;
         continue;
       }
for (auto &BB : F) {
           return DILocation::get(M.getContext(), DL.getLine(), DL.getCol(),
                                  Scope, InlinedAt);
         };
 
         if (I.getDebugLoc() != DebugLoc())
           I.setDebugLoc(remapDebugLoc(I.getDebugLoc()));
 
         // Remap DILocations in llvm.loop attachments.
-        updateLoopMetadataDebugLocations(I, [&](const DILocation &Loc) {
-          return remapDebugLoc(&Loc).get();
+        updateLoopMetadataDebugLocations(I, [&](Metadata *MD) -> Metadata * {
+          if (auto *Loc = dyn_cast_or_null<DILocation>(MD))
+            return remapDebugLoc(Loc).get();
+          return MD;
         });
 
         // Strip heapallocsite attachments, they point into the DIType system.
         if (I.hasMetadataOtherThanDebugLoc())
           I.setMetadata("heapallocsite", nullptr);
       }
     }
   }

llvm/lib/Transforms/Utils/CodeExtractor.cpp

static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
 
   // Fix up the scope information attached to the line locations in the new
   // function.
   for (Instruction &I : instructions(NewFunc)) {
     if (const DebugLoc &DL = I.getDebugLoc())
       I.setDebugLoc(DILocation::get(Ctx, DL.getLine(), DL.getCol(), NewSP));
 
     // Loop info metadata may contain line locations. Fix them up.
-    auto updateLoopInfoLoc = [&Ctx,
-                              NewSP](const DILocation &Loc) -> DILocation * {
-      return DILocation::get(Ctx, Loc.getLine(), Loc.getColumn(), NewSP,
+    auto updateLoopInfoLoc = [&Ctx, NewSP](Metadata *MD) -> Metadata * {
+      if (auto *Loc = dyn_cast_or_null<DILocation>(MD))
+        return DILocation::get(Ctx, Loc->getLine(), Loc->getColumn(), NewSP,
                                nullptr);
+      return MD;
     };
     updateLoopMetadataDebugLocations(I, updateLoopInfoLoc);
   }
   if (!TheCall.getDebugLoc())
     TheCall.setDebugLoc(DILocation::get(Ctx, 0, 0, OldSP));
 
   eraseDebugIntrinsicsWithNonLocalRefs(NewFunc);
 }

llvm/lib/Transforms/Utils/InlineFunction.cpp

static void fixupLineNumbers(Function *Fn, Function::iterator FI,
   // the call site location instead.
   bool NoInlineLineTables = Fn->hasFnAttribute("no-inline-line-tables");
 
   for (; FI != Fn->end(); ++FI) {
     for (BasicBlock::iterator BI = FI->begin(), BE = FI->end();
          BI != BE; ++BI) {
       // Loop metadata needs to be updated so that the start and end locs
       // reference inlined-at locations.
-      auto updateLoopInfoLoc = [&Ctx, &InlinedAtNode, &IANodes](
-                                   const DILocation &Loc) -> DILocation * {
-        return inlineDebugLoc(&Loc, InlinedAtNode, Ctx, IANodes).get();
+      auto updateLoopInfoLoc = [&Ctx, &InlinedAtNode,
+                                &IANodes](Metadata *MD) -> Metadata * {
+        if (auto *Loc = dyn_cast_or_null<DILocation>(MD))
+          return inlineDebugLoc(Loc, InlinedAtNode, Ctx, IANodes).get();
+        return MD;
       };
       updateLoopMetadataDebugLocations(*BI, updateLoopInfoLoc);
 
       if (!NoInlineLineTables)
         if (DebugLoc DL = BI->getDebugLoc()) {
           DebugLoc IDL =
               inlineDebugLoc(DL, InlinedAtNode, BI->getContext(), IANodes);
           BI->setDebugLoc(IDL);

llvm/test/Verifier/llvm.loop-cu-strip-indirect.ll

+; RUN: llvm-as -disable-output < %s -o /dev/null 2>&1 | FileCheck %s
+; RUN: llvm-as < %s -o - | llvm-dis - | FileCheck %s --check-prefix=CHECK-STRIP
+; CHECK: DICompileUnit not listed in llvm.dbg.cu
+; CHECK: ignoring invalid debug info in
+; CHECK-NOT: DICompileUnit not listed in llvm.dbg.cu
+declare hidden void @g() local_unnamed_addr #1 align 2
+define hidden void @f() {
+  tail call void @g() #2, !llvm.loop !5
+  ret void
+}
+!llvm.module.flags = !{!0, !1}
+!0 = !{i32 2, !"Dwarf Version", i32 4}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!5 = distinct !{!5, !14, !15, !"fake loop metadata"}
+; CHECK-STRIP: ![[MD:.*]] = distinct !{![[MD]], !{{[0-9]+}}, !"fake loop metadata"}
+!7 = distinct !DISubprogram(name: "f", scope: !8, file: !8, line: 1324, type: !9, scopeLine: 1324, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !11)
+!8 = !DIFile(filename: "/", directory: "f.cpp")
+!9 = !DISubroutineType(types: !10)
+!10 = !{}
+!11 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: !8)
+!14 = !{!"metadata 1", i1 true}
+!15 = !{!"metadata 1", !16}
+!16 = distinct !{!16, !17}
+!17 = !DILocation(line: 105, column: 3, scope: !18)
+!18 = distinct !DILexicalBlock(scope: !7, file: !8, line: 105, column: 3)