1 files changed, 39 insertions, 0 deletions

diff --git a/llvm/lib/IR/DebugInfo.cpp b/llvm/lib/IR/DebugInfo.cpp
index 6b9bc689a44..3bdd5a6bd95 100644
--- a/llvm/lib/IR/DebugInfo.cpp
+++ b/llvm/lib/IR/DebugInfo.cpp
@@ -239,6 +239,35 @@ bool DebugInfoFinder::addScope(DIScope *Scope) {
   return true;
 }
 
+static llvm::MDNode *stripDebugLocFromLoopID(llvm::MDNode *N) {
+  assert(N->op_begin() != N->op_end() && "Missing self reference?");
+  auto DebugLocOp =
+      std::find_if(N->op_begin() + 1, N->op_end(), [](const MDOperand &Op) {
+        return isa<DILocation>(Op.get());
+      });
+
+  // No debug location, we do not have to rewrite this MDNode.
+  if (DebugLocOp == N->op_end())
+    return N;
+
+  // There is only the debug location without any actual loop metadata, hence we
+  // can remove the metadata.
+  if (N->getNumOperands() == 2)
+    return nullptr;
+
+  SmallVector<Metadata *, 4> Args;
+  // Reserve operand 0 for loop id self reference.
+  auto TempNode = MDNode::getTemporary(N->getContext(), None);
+  Args.push_back(TempNode.get());
+  Args.append(N->op_begin() + 1, DebugLocOp);
+  Args.append(DebugLocOp + 1, N->op_end());
+
+  // Set the first operand to itself.
+  MDNode *LoopID = MDNode::get(N->getContext(), Args);
+  LoopID->replaceOperandWith(0, LoopID);
+  return LoopID;
+}
+
 bool llvm::stripDebugInfo(Function &F) {
   bool Changed = false;
   if (F.getSubprogram()) {
@@ -246,6 +275,7 @@ bool llvm::stripDebugInfo(Function &F) {
     F.setSubprogram(nullptr);
   }
 
+  llvm::DenseMap<llvm::MDNode*, llvm::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.
@@ -259,6 +289,15 @@ bool llvm::stripDebugInfo(Function &F) {
         I.setDebugLoc(DebugLoc());
       }
     }
+
+    auto *TermInst = BB.getTerminator();
+    if (auto *LoopID = TermInst->getMetadata(LLVMContext::MD_loop)) {
+      auto *NewLoopID = LoopIDsMap.lookup(LoopID);
+      if (!NewLoopID)
+        NewLoopID = LoopIDsMap[LoopID] = stripDebugLocFromLoopID(LoopID);
+      if (NewLoopID != LoopID)
+        TermInst->setMetadata(LLVMContext::MD_loop, NewLoopID);
+    }
   }
   return Changed;
 }