[DebugInfo][RemoveDIs] Extract DPValues in CodeExtractor like dbg.values (#73252)

CodeExtractor shifts dbg.value intrinsics out of the region being
extracted and updates them to be appropriate in the extracted function.
With new non-intrinsic variable locations, we need to manually do this
too, with DPValues.

Most of this patch shifts and refactors some utilities in
fixupDebugInfoPostExtraction so that we can add a single extra helper
lambda that iterates over DPValues and applies update-utilities. We also
have to assign the IsNewDbgInfoFormat flag in a bunch of places -- this
normally gets set the moment you insert a block into a function (or
function into a module), however a few blocks are constructed here
before being inserted, thus we have to do some manual setup.

Tested via LoopExtractor_alloca.ll, which invokes debugify.

1 files changed, 57 insertions, 24 deletions

diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index 405753158e9c..9c1186232e02 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -768,6 +768,7 @@ void CodeExtractor::severSplitPHINodesOfExits(
         NewBB = BasicBlock::Create(ExitBB->getContext(),
                                    ExitBB->getName() + ".split",
                                    ExitBB->getParent(), ExitBB);
+        NewBB->IsNewDbgInfoFormat = ExitBB->IsNewDbgInfoFormat;
         SmallVector<BasicBlock *, 4> Preds(predecessors(ExitBB));
         for (BasicBlock *PredBB : Preds)
           if (Blocks.count(PredBB))
@@ -889,6 +890,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   Function *newFunction = Function::Create(
       funcType, GlobalValue::InternalLinkage, oldFunction->getAddressSpace(),
       oldFunction->getName() + "." + SuffixToUse, M);
+  newFunction->IsNewDbgInfoFormat = oldFunction->IsNewDbgInfoFormat;
 
   // Inherit all of the target dependent attributes and white-listed
   // target independent attributes.
@@ -1505,10 +1507,14 @@ void CodeExtractor::calculateNewCallTerminatorWeights(
 static void eraseDebugIntrinsicsWithNonLocalRefs(Function &F) {
   for (Instruction &I : instructions(F)) {
     SmallVector<DbgVariableIntrinsic *, 4> DbgUsers;
-    findDbgUsers(DbgUsers, &I);
+    SmallVector<DPValue *, 4> DPValues;
+    findDbgUsers(DbgUsers, &I, &DPValues);
     for (DbgVariableIntrinsic *DVI : DbgUsers)
       if (DVI->getFunction() != &F)
         DVI->eraseFromParent();
+    for (DPValue *DPV : DPValues)
+      if (DPV->getFunction() != &F)
+        DPV->eraseFromParent();
   }
 }
 
@@ -1544,6 +1550,16 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
       /*LineNo=*/0, SPType, /*ScopeLine=*/0, DINode::FlagZero, SPFlags);
   NewFunc.setSubprogram(NewSP);
 
+  auto IsInvalidLocation = [&NewFunc](Value *Location) {
+    // Location is invalid if it isn't a constant or an instruction, or is an
+    // instruction but isn't in the new function.
+    if (!Location ||
+        (!isa<Constant>(Location) && !isa<Instruction>(Location)))
+      return true;
+    Instruction *LocationInst = dyn_cast<Instruction>(Location);
+    return LocationInst && LocationInst->getFunction() != &NewFunc;
+  };
+
   // Debug intrinsics in the new function need to be updated in one of two
   // ways:
   //  1) They need to be deleted, because they describe a value in the old
@@ -1552,8 +1568,41 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
   //     point to a variable in the wrong scope.
   SmallDenseMap<DINode *, DINode *> RemappedMetadata;
   SmallVector<Instruction *, 4> DebugIntrinsicsToDelete;
+  SmallVector<DPValue *, 4> DPVsToDelete;
   DenseMap<const MDNode *, MDNode *> Cache;
+
+  auto GetUpdatedDIVariable = [&](DILocalVariable *OldVar) {
+    DINode *&NewVar = RemappedMetadata[OldVar];
+    if (!NewVar) {
+      DILocalScope *NewScope = DILocalScope::cloneScopeForSubprogram(
+          *OldVar->getScope(), *NewSP, Ctx, Cache);
+      NewVar = DIB.createAutoVariable(
+          NewScope, OldVar->getName(), OldVar->getFile(), OldVar->getLine(),
+          OldVar->getType(), /*AlwaysPreserve=*/false, DINode::FlagZero,
+          OldVar->getAlignInBits());
+    }
+    return cast<DILocalVariable>(NewVar);
+  };
+
+  auto UpdateDPValuesOnInst = [&](Instruction &I) -> void {
+    for (auto &DPV : I.getDbgValueRange()) {
+      // Apply the two updates that dbg.values get: invalid operands, and
+      // variable metadata fixup.
+      // FIXME: support dbg.assign form of DPValues.
+      if (any_of(DPV.location_ops(), IsInvalidLocation)) {
+        DPVsToDelete.push_back(&DPV);
+        continue;
+      }
+      if (!DPV.getDebugLoc().getInlinedAt())
+        DPV.setVariable(GetUpdatedDIVariable(DPV.getVariable()));
+      DPV.setDebugLoc(DebugLoc::replaceInlinedAtSubprogram(DPV.getDebugLoc(),
+                                                           *NewSP, Ctx, Cache));
+    }
+  };
+
   for (Instruction &I : instructions(NewFunc)) {
+    UpdateDPValuesOnInst(I);
+
     auto *DII = dyn_cast<DbgInfoIntrinsic>(&I);
     if (!DII)
       continue;
@@ -1575,16 +1624,6 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
       continue;
     }
 
-    auto IsInvalidLocation = [&NewFunc](Value *Location) {
-      // Location is invalid if it isn't a constant or an instruction, or is an
-      // instruction but isn't in the new function.
-      if (!Location ||
-          (!isa<Constant>(Location) && !isa<Instruction>(Location)))
-        return true;
-      Instruction *LocationInst = dyn_cast<Instruction>(Location);
-      return LocationInst && LocationInst->getFunction() != &NewFunc;
-    };
-
     auto *DVI = cast<DbgVariableIntrinsic>(DII);
     // If any of the used locations are invalid, delete the intrinsic.
     if (any_of(DVI->location_ops(), IsInvalidLocation)) {
@@ -1599,23 +1638,14 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
     }
     // If the variable was in the scope of the old function, i.e. it was not
     // inlined, point the intrinsic to a fresh variable within the new function.
-    if (!DVI->getDebugLoc().getInlinedAt()) {
-      DILocalVariable *OldVar = DVI->getVariable();
-      DINode *&NewVar = RemappedMetadata[OldVar];
-      if (!NewVar) {
-        DILocalScope *NewScope = DILocalScope::cloneScopeForSubprogram(
-            *OldVar->getScope(), *NewSP, Ctx, Cache);
-        NewVar = DIB.createAutoVariable(
-            NewScope, OldVar->getName(), OldVar->getFile(), OldVar->getLine(),
-            OldVar->getType(), /*AlwaysPreserve=*/false, DINode::FlagZero,
-            OldVar->getAlignInBits());
-      }
-      DVI->setVariable(cast<DILocalVariable>(NewVar));
-    }
+    if (!DVI->getDebugLoc().getInlinedAt())
+      DVI->setVariable(GetUpdatedDIVariable(DVI->getVariable()));
   }
 
   for (auto *DII : DebugIntrinsicsToDelete)
     DII->eraseFromParent();
+  for (auto *DPV : DPVsToDelete)
+    DPV->getMarker()->MarkedInstr->dropOneDbgValue(DPV);
   DIB.finalizeSubprogram(NewSP);
 
   // Fix up the scope information attached to the line locations in the new
@@ -1721,11 +1751,14 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
   BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(),
                                                 "codeRepl", oldFunction,
                                                 header);
+  codeReplacer->IsNewDbgInfoFormat = oldFunction->IsNewDbgInfoFormat;
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(),
                                                "newFuncRoot");
+  newFuncRoot->IsNewDbgInfoFormat = oldFunction->IsNewDbgInfoFormat;
+
   auto *BranchI = BranchInst::Create(header);
   // If the original function has debug info, we have to add a debug location
   // to the new branch instruction from the artificial entry block.