Avoid removing useful loop metadata when striping debug info

[dongrui]

Summary

The stripDebugLocFromLoopID() may mistakenly remove useful metadata nodes
when they are operands of a child node, which also has DILocation operands.

It can be reproduced by the output of clang for code similar to the following:

#pragma clang loop unroll_count(4)
#pragma clang loop vectorize(assume_safety)
for(int i = 0; i < n; i++)
  x[i] = 10;

-strip-debug removes the child node of llvm.loop.vectorize.followup_all,
which contains llvm.loop.isvectorized and llvm.loop.unroll.count.

This patch fixes by checking all children nodes and only remove a metadata
node if all its children are DILocation.

[dongrui]

HI,

I have no commit access to llvm git repo. Could someone please commit the patch for me?

Thanks,
Dongrui

llvm/lib/IR/DebugInfo.cpp

if (!N)
     return false;
   if (isa<DILocation>(N) || Reachable.count(N))
     return true;
   if (!Visited.insert(N).second)
     return false;
   for (auto &OpIt : N->operands()) {
     Metadata *Op = OpIt.get();
     if (isDILocationReachable(Visited, Reachable, Op)) {
+      // Don't return just yet as we want to visit all MD's children to
+      // initialize DILocationReachable in stripDebugLocFromLoopID
       Reachable.insert(N);
-      return true;
     }
   }
+  return Reachable.count(N);
+}
+
+static bool isAllDILocation(SmallPtrSetImpl<Metadata *> &Visited,
+                            SmallPtrSetImpl<Metadata *> &AllDILocation,
+                            const SmallPtrSetImpl<Metadata *> &DIReachable,
+                            Metadata *MD) {
+  MDNode *N = dyn_cast_or_null<MDNode>(MD);
+  if (!N)
+    return false;
+  if (isa<DILocation>(N) || AllDILocation.count(N))
+    return true;
+  if (!DIReachable.count(N))
     return false;
+  if (!Visited.insert(N).second)
+    return false;
+  for (auto &OpIt : N->operands()) {
+    Metadata *Op = OpIt.get();
+    if (Op == MD)
+      continue;
+    if (!isAllDILocation(Visited, AllDILocation, DIReachable, Op)) {
+      return false;
+    }
+  }
+  AllDILocation.insert(N);
+  return true;
+}
+
+static Metadata *
+stripLoopMDLoc(const SmallPtrSetImpl<Metadata *> &AllDILocation,
+               const SmallPtrSetImpl<Metadata *> &DIReachable, Metadata *MD) {
+  if (isa<DILocation>(MD) || AllDILocation.count(MD))
+    return nullptr;
+
+  if (!DIReachable.count(MD))
+    return MD;
+
+  MDNode *N = dyn_cast_or_null<MDNode>(MD);
+  if (!N)
+    return MD;
+
+  SmallVector<Metadata *, 4> Args;
+  bool HasSelfRef = false;
+  for (unsigned i = 0; i < N->getNumOperands(); ++i) {
+    Metadata *A = N->getOperand(i);
+    if (!A) {
+      Args.push_back(nullptr);
+    } else if (A == MD) {
+      assert(i == 0 && "expected i==0 for self-reference");
+      HasSelfRef = true;
+      Args.push_back(nullptr);
+    } else if (Metadata *NewArg =
+                   stripLoopMDLoc(AllDILocation, DIReachable, A)) {
+      Args.push_back(NewArg);
+    }
+  }
+  if (Args.empty() || (HasSelfRef && Args.size() == 1))
+    return nullptr;
+
+  MDNode *NewMD = N->isDistinct() ? MDNode::getDistinct(N->getContext(), Args)
+                                  : MDNode::get(N->getContext(), Args);
+  if (HasSelfRef)
+    NewMD->replaceOperandWith(0, NewMD);
+  return NewMD;
 }
 
 static MDNode *stripDebugLocFromLoopID(MDNode *N) {
   assert(!N->operands().empty() && "Missing self reference?");
-  SmallPtrSet<Metadata *, 8> Visited, DILocationReachable;
+  SmallPtrSet<Metadata *, 8> Visited, DILocationReachable, AllDILocation;
   // 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. 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(),
+  if (!llvm::count_if(llvm::drop_begin(N->operands()),
                      [&Visited, &DILocationReachable](const MDOperand &Op) {
                        return isDILocationReachable(
                                   Visited, DILocationReachable, Op.get());
                      }))
     return N;
 
+  Visited.clear();
   // If there is only the debug location without any actual loop metadata, we
   // can remove the metadata.
   if (llvm::all_of(llvm::drop_begin(N->operands()),
-                   [&Visited, &DILocationReachable](const MDOperand &Op) {
-                     return isDILocationReachable(Visited, DILocationReachable,
-                                                  Op.get());
+                   [&Visited, &AllDILocation,
+                    &DILocationReachable](const MDOperand &Op) {
+                     return isAllDILocation(Visited, AllDILocation,
+                                            DILocationReachable, Op.get());
                    }))
     return nullptr;
 
   return updateLoopMetadataDebugLocationsImpl(
-      N, [&DILocationReachable](Metadata *MD) -> Metadata * {
-        if (isa<DILocation>(MD) || DILocationReachable.count(MD))
-          return nullptr;
-        return MD;
+      N, [&AllDILocation, &DILocationReachable](Metadata *MD) -> Metadata * {
+        return stripLoopMDLoc(AllDILocation, DILocationReachable, MD);
       });
 }
 
 bool llvm::stripDebugInfo(Function &F) {
   bool Changed = false;
   if (F.hasMetadata(LLVMContext::MD_dbg)) {
     Changed = true;
     F.setSubprogram(nullptr);

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

+; RUN: llvm-as < %s -o - | llvm-dis - | FileCheck %s
+
+; The loop metadata in this test is similar to the output from clang for
+; code like the this:
+; 
+;     #pragma clang loop unroll_count(4)
+;     #pragma clang loop vectorize(assume_safety)
+;     for(int i = 0; i < n; i++)
+;       ...
+;
+; strip debug should not remove the useful metadata in the followup part.
+
+declare hidden void @g() local_unnamed_addr align 2
+
+define hidden void @f() {
+; CHECK: tail call {{.*}} !llvm.loop [[LID:![0-9]+]]
+  tail call void @g(), !llvm.loop !2
+  ret void
+}
+
+!llvm.module.flags = !{!0, !1}
+
+; CHECK-NOT: DILocation
+; CHECK: [[LID]] = distinct !{[[LID]], [[VE:![0-9]+]], [[VF:![0-9]+]]}
+; CHECK-NOT: DILocation
+; CHECK: [[VE]] = !{!"llvm.loop.vectorize.enable", i1 true}
+; CHECK: [[VF]] = !{!"llvm.loop.vectorize.followup_all", [[FU:![0-9]+]]}
+; CHECK: [[FU]] = distinct !{[[FU]], [[VD:![0-9]+]], [[UC:![0-9]+]]}
+; CHECK: [[VD]] = !{!"llvm.loop.isvectorized"}
+; CHECK: [[UC]] = !{!"llvm.loop.unroll.count", i32 4}
+
+!0 = !{i32 2, !"Dwarf Version", i32 4}
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!2 = distinct !{!2, !3, !10, !11, !12}
+!3 = !DILocation(line: 5, column: 12, scope: !4)
+!4 = distinct !DILexicalBlock(scope: !6, file: !5, line: 105, column: 3)
+!5 = !DIFile(filename: "/", directory: "f.cpp")
+!6 = distinct !DISubprogram(name: "f", scope: !5, file: !5, line: 1324, type: !7, scopeLine: 1324, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !9)
+!7 = !DISubroutineType(types: !8)
+!8 = !{}
+!9 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: !5, isOptimized: false, runtimeVersion: 0, emissionKind: NoDebug)
+!10 = !DILocation(line: 6, column: 1, scope: !4)
+!11 = !{!"llvm.loop.vectorize.enable", i1 true}
+!12 = !{!"llvm.loop.vectorize.followup_all", !13}
+!13 = distinct !{!13, !3, !10, !14, !15}
+!14 = !{!"llvm.loop.isvectorized"}
+!15 = !{!"llvm.loop.unroll.count", i32 4}

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"}
+; CHECK-STRIP: ![[MD:.*]] = distinct !{![[MD]], !{{[0-9]+}}, !{{[0-9]+}}, !"fake loop metadata"}
+; CHECK-STRIP-NOT: DILocation
 !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)