[CodeExtractor] Correctly propagate scope information post extraction

[fdeazeve]

Summary

When a new function "NewF" is created with instructions extracted from
another function "OldF", the CodeExtractor only preserves debug
line/column of the extracted instructions. However:

1. Any inlinedAt nodes are dropped.
2. The scope chain is replaced with a single node, the Subprogram of NewF.

Both of these are incorrect: most of the debug metadata from the
original instructions should be preserved. We only need to update the
Subprogram found at the scope of the last node of the inline chain; this
Subprogram used to be OldF but now should be NewF.

[fdeazeve]

Simplified map update

[fdeazeve]

Update test

[fdeazeve]

Updated commit message

[fdeazeve]

Rename a variable in the test for clarity

[aprantl]

This is great! The old code was quite incorrect :-)

I have one possible suggestion for a cleaner/safer implementation of replaceOperand().
Did you test this on something large, like Clang?

llvm/lib/IR/DebugLoc.cpp

71

should this be static?

74

This seems to be very brute force. How much more code would it be to Switch over getKind() or put a replaceScope() method on DILexicalScopeBase?

Can you make an assertion that the Node is distinct? Otherwise this would be dangerous.

114

This might benefit from a comment.
IIUC, we collect a chain of not-yet processed inlined-at locations, which would be useful when processing increasingly deeper inlined instructions.

126

maybe also make clear this this means we are now in the DILocation belonging to the current subprogram and that that's why we need to reparent the scope into the new function.

[aprantl]

Some more context for the other reviewers: This affects hot-cold splitting, particularly the inlined frames of an assertion.

[dblaikie]

Yeah, this sounds like it was quite incorrect & the general vibe of the new code looks good to me (I'd wonder if it could be simpler to write recursively & be OK (admittedly that'd be a bit more recursive than the real code - since it'd recurse for every lexical scope, not just every inlined frame) - but, yeah, probably better the way it's written here, but more complicated to follow).

I'll leave the approval/review to you, @aprantl, but endorse the direction.

[fdeazeve]

In D139217#3967447, @aprantl wrote:

This is great! The old code was quite incorrect :-)

I have one possible suggestion for a cleaner/safer implementation of replaceOperand().
Did you test this on something large, like Clang?

I had only run the Clang/LLVM/LLDB test suites, but I will try compiling a larger project.

llvm/lib/IR/DebugLoc.cpp

71

Yes! Though this code is going to be removed in the new version

74

Can you make an assertion that the Node is distinct? Otherwise this would be dangerous.

The assert on line 73 is doing that (assert(!N.isUniqued())). This assert allo

This seems to be very brute force. How much more code would it be to Switch over getKind() or put a replaceScope() method on DILexicalScopeBase?

Agreed, this implementation is very similar to what the ValueMapper class does, so maybe it is more general than it needs to be.
Adding a method to DILexicalBlockBase (I'm assuming this is the class you meant?) seems straightforward.

114

Fixed in the next version

126

Will fix in the next version.

[fdeazeve]

Addressed review comments

[fdeazeve]

I've now used this patch to compile a RelWithDebugSymbols build of llvm-project, and had no issues. Also tried on an e2e test that used to be buggy without the fix.

[fdeazeve]

Can you make an assertion that the Node is distinct? Otherwise this would be dangerous.

The assert on line 73 is doing that (assert(!N.isUniqued())). This assert allo

Oops, I realized I never finished the sentence above! That assert (which is now in the setScope function), checks whether the node is either temporary or distinct, both conditions which would allow mutation.

[aprantl]

This looks good to me now. I have a few suggestions for better names and comments, but I'm fine with the mechanics!

llvm/include/llvm/IR/DebugInfoMetadata.h

2119 ↗

(On Diff #479748)

The other sibling classes would call this replaceScope().

llvm/lib/IR/DebugLoc.cpp

72

This function doesn't replace anything. What do you think of reparentScopeChain, cloneScopesForSubprogram()?

87

Maybe add a comment explaining what the entire loop does?

126

Top-level comment for this loop?

[fdeazeve]

llvm/include/llvm/IR/DebugInfoMetadata.h

2119 ↗

(On Diff #479748)

Good point!

llvm/lib/IR/DebugLoc.cpp

72

The irony here is that we've just created a "replaceScope" method... but you're right, this makes it sound like we're modifying the original scope. I'll go with the "clone" variant

[fdeazeve]

Addressed review comments

[fdeazeve]

Test failure due to the spurious timeout from libfuzzer, unrelated to this patch.

llvm/include/llvm/IR/DebugLoc.h

public:
     MDNode *getScope() const;
     DILocation *getInlinedAt() const;
 
     /// Get the fully inlined-at scope for a DebugLoc.
     ///
     /// Gets the inlined-at scope for a DebugLoc.
     MDNode *getInlinedAtScope() const;
 
+    /// Rebuild the entire inline-at chain by replacing the subprogram at the
+    /// end of the chain with NewSP.
+    static DebugLoc
+    replaceInlinedAtSubprogram(const DebugLoc &DL, DISubprogram &NewSP,
+                               LLVMContext &Ctx,
+                               DenseMap<const MDNode *, MDNode *> &Cache);
+
     /// Find the debug info location for the start of the function.
     ///
     /// Walk up the scope chain of given debug loc and find line number info
     /// for the function.
     ///
     /// FIXME: Remove this.  Users should use DILocation/DILocalScope API to
     /// find the subprogram, and then DILocation::get().
     DebugLoc getFnDebugLoc() const;

llvm/lib/IR/DebugLoc.cpp

 }
 
 void DebugLoc::setImplicitCode(bool ImplicitCode) {
   if (DILocation *Loc = get()) {
     Loc->setImplicitCode(ImplicitCode);
   }
 }
 
+/// Updates the operands of N, replacing OperandToReplace with Replacement.
+void replaceOperand(MDNode &N, Metadata *OperandToReplace,
+                    MDNode *Replacement) {
+  assert(!N.isUniqued() && "Expected distinct or temporary nodes");
+  for (auto Idx : seq(0u, N.getNumOperands())) {
+    Metadata *Old = N.getOperand(Idx);
+    if (Old == OperandToReplace)
+      N.replaceOperandWith(Idx, Replacement);
+  }
+}
+
+/// Traverses the scope chain rooted at RootScope until it hits a Subprogram,
+/// recreating the chain with "NewSP" instead.
+static DIScope *replaceSubprogram(DILocalScope &RootScope, DISubprogram &NewSP,
+                                  LLVMContext &Ctx,
+                                  DenseMap<const MDNode *, MDNode *> &Cache) {
+  SmallVector<DIScope *> ScopeChain;
+  DIScope *CachedResult = nullptr;
+
+  for (DIScope *Scope = &RootScope; !isa<DISubprogram>(Scope);
+       Scope = Scope->getScope()) {
+    if (auto It = Cache.find(Scope); It != Cache.end()) {
+      CachedResult = cast<DIScope>(It->second);
+      break;
+    }
+    ScopeChain.push_back(Scope);
+  }
+
+  DIScope *UpdatedScope = CachedResult ? CachedResult : &NewSP;
+  for (DIScope *ScopeToUpdate : reverse(ScopeChain)) {
+    TempMDNode ClonedScope = ScopeToUpdate->clone();
+    replaceOperand(*ClonedScope, ScopeToUpdate->getScope(), UpdatedScope);
+    UpdatedScope =
+        cast<DIScope>(MDNode::replaceWithUniqued(std::move(ClonedScope)));
+    Cache.try_emplace(ScopeToUpdate, UpdatedScope);
+  }
+  return UpdatedScope;
+}
+
+DebugLoc DebugLoc::replaceInlinedAtSubprogram(
+    const DebugLoc &RootLoc, DISubprogram &NewSP, LLVMContext &Ctx,
+    DenseMap<const MDNode *, MDNode *> &Cache) {
+  SmallVector<DILocation *> LocChain;
+  DILocation *CachedResult = nullptr;
+
+  for (DILocation *Loc = RootLoc; Loc; Loc = Loc->getInlinedAt()) {
+    if (auto It = Cache.find(Loc); It != Cache.end()) {
+      CachedResult = cast<DILocation>(It->second);
+      break;
+    }
+    LocChain.push_back(Loc);
+  }
+
+  DILocation *UpdatedLoc = CachedResult;
+  if (!UpdatedLoc) {
+    // If no cache hits, then back() is the end of the inline chain and we need
+    // to update its *scope* chain.
+    DILocation *LocToUpdate = LocChain.pop_back_val();
+    DIScope *NewScope =
+        replaceSubprogram(*LocToUpdate->getScope(), NewSP, Ctx, Cache);
+    UpdatedLoc = DILocation::get(Ctx, LocToUpdate->getLine(),
+                                 LocToUpdate->getColumn(), NewScope);
+    Cache.try_emplace(LocToUpdate, UpdatedLoc);
+  }
+
+  for (const DILocation *LocToUpdate : reverse(LocChain)) {
+    UpdatedLoc =
+        DILocation::get(Ctx, LocToUpdate->getLine(), LocToUpdate->getColumn(),
+                        LocToUpdate->getScope(), UpdatedLoc);
+    Cache.try_emplace(LocToUpdate, UpdatedLoc);
+  }
+
+  return UpdatedLoc;
+}
+
 DebugLoc DebugLoc::appendInlinedAt(const DebugLoc &DL, DILocation *InlinedAt,
                                    LLVMContext &Ctx,
                                    DenseMap<const MDNode *, MDNode *> &Cache) {
   SmallVector<DILocation *, 3> InlinedAtLocations;
   DILocation *Last = InlinedAt;
   DILocation *CurInlinedAt = DL;
 
   // Gather all the inlined-at nodes.

llvm/lib/Transforms/Utils/CodeExtractor.cpp

for (Instruction &I : instructions(NewFunc)) {
     DVI->setVariable(cast<DILocalVariable>(NewVar));
   }
   for (auto *DII : DebugIntrinsicsToDelete)
     DII->eraseFromParent();
   DIB.finalizeSubprogram(NewSP);
 
   // Fix up the scope information attached to the line locations in the new
   // function.
+  DenseMap<const MDNode *, MDNode *> Cache;
   for (Instruction &I : instructions(NewFunc)) {
     if (const DebugLoc &DL = I.getDebugLoc())
-      I.setDebugLoc(DILocation::get(Ctx, DL.getLine(), DL.getCol(), NewSP));
+      I.setDebugLoc(
+          DebugLoc::replaceInlinedAtSubprogram(DL, *NewSP, Ctx, Cache));
 
     // Loop info metadata may contain line locations. Fix them up.
     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;
     };

llvm/test/Transforms/HotColdSplit/transfer-debug-info.ll

 ; CHECK-NEXT: call void @sink(i32 [[ADD1]]), !dbg [[LINE1:![0-9]+]]
 
 ; - llvm.dbg.value intrinsics for values local to @foo.cold.1 are preserved
 ; CHECK-NEXT: llvm.dbg.value(metadata i32 [[ADD1]], metadata [[VAR1:![0-9]+]], metadata !DIExpression()), !dbg [[LINE1]]
 
 ; - Expressions inside of dbg.value intrinsics are preserved
 ; CHECK-NEXT: llvm.dbg.value(metadata i32 [[ADD1]], metadata [[VAR1]], metadata !DIExpression(DW_OP_constu, 1, DW_OP_plus, DW_OP_stack_value)
 
+; CHECK-NEXT: call void @sink(i32 [[ADD1]]), !dbg [[LINE2:![0-9]+]]
+
 ; - The DISubprogram for @foo.cold.1 has an empty DISubroutineType
 ; CHECK: [[FILE:![0-9]+]] = !DIFile(filename: "<stdin>"
 ; CHECK: [[EMPTY_MD:![0-9]+]] = !{}
 ; CHECK: [[EMPTY_TYPE:![0-9]+]] = !DISubroutineType(types: [[EMPTY_MD]])
+; CHECK: [[ORIGINAL_FOO_SCOPE:![0-9]+]] = distinct !DISubprogram(name: "inline_me"
 ; CHECK: [[NEWSCOPE:![0-9]+]] = distinct !DISubprogram(name: "foo.cold.1", linkageName: "foo.cold.1", scope: null, file: [[FILE]], type: [[EMPTY_TYPE]], spFlags: DISPFlagLocalToUnit | DISPFlagDefinition | DISPFlagOptimized
 
 ; - Line locations in @foo.cold.1 point to the new scope for @foo.cold.1
 ; CHECK: [[LINE1]] = !DILocation(line: 1, column: 1, scope: [[NEWSCOPE]])
 
+; CHECK: [[LINE2]] =         !DILocation(line: 2, column: 2, scope: [[ORIGINAL_FOO_SCOPE]]
+; CHECK-SAME:                            inlinedAt: [[INLINED_AT:![0-9]*]]
+; CHECK: [[INLINED_AT]] =    !DILocation(line: 3, column: 3, scope: [[INLINED_SCOPE1:![0-9]*]]
+; CHECK: [[INLINED_SCOPE1]] = !DILexicalBlock(scope: [[INLINED_SCOPE2:![0-9]*]], file: [[FILE]], line: 4, column: 4)
+; CHECK: [[INLINED_SCOPE2]] = !DILexicalBlock(scope: [[NEWSCOPE]], file: [[FILE]], line: 5, column: 5)
+
 define void @foo(i32 %arg1) !dbg !6 {
 entry:
   %var = add i32 0, 0, !dbg !11
   br i1 undef, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   ret void
 
 if.end:                                           ; preds = %entry
   call void @llvm.dbg.value(metadata i32 %arg1, metadata !9, metadata !DIExpression()), !dbg !11
   %add1 = add i32 %arg1, 1
   call void @sink(i32 %add1)
   call void @sink(i32 %add1), !dbg !11
   call void @llvm.dbg.value(metadata i32 %add1, metadata !9, metadata !DIExpression()), !dbg !11
   call void @llvm.dbg.value(metadata i32 %add1, metadata !9, metadata !DIExpression(DW_OP_constu, 1, DW_OP_plus, DW_OP_stack_value)), !dbg !11
+  call void @sink(i32 %add1), !dbg !13 ; inlined from @inline_me
+  call void @sink(i32 %add1), !dbg !14 ; not inlined, but inside some scope of foo
   ret void
 }
 
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 declare void @sink(i32) cold
 
+define void @inline_me() !dbg !12{
+  ret void
+}
+
 !llvm.dbg.cu = !{!0}
 !llvm.debugify = !{!3, !4}
 !llvm.module.flags = !{!5}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "debugify", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !2)
 !1 = !DIFile(filename: "<stdin>", directory: "/")
 !2 = !{}
 !3 = !{i32 7}
 !4 = !{i32 1}
 !5 = !{i32 2, !"Debug Info Version", i32 3}
 !6 = distinct !DISubprogram(name: "foo", linkageName: "foo", scope: null, file: !1, line: 1, type: !7, isLocal: false, isDefinition: true, scopeLine: 1, isOptimized: true, unit: !0, retainedNodes: !8)
 !7 = !DISubroutineType(types: !2)
 !8 = !{!9}
 !9 = !DILocalVariable(name: "1", scope: !6, file: !1, line: 1, type: !10)
 !10 = !DIBasicType(name: "ty32", size: 32, encoding: DW_ATE_unsigned)
 !11 = !DILocation(line: 1, column: 1, scope: !6)
+!12 = distinct !DISubprogram(name: "inline_me", linkageName: "inline_me", scope: null, file: !1, line: 1, type: !7, isLocal: false, isDefinition: true, scopeLine: 1, isOptimized: true, unit: !0, retainedNodes: !8)
+!13 = !DILocation(line: 2, column: 2, scope: !12, inlinedAt: !14)
+!14 = !DILocation(line: 3, column: 3, scope: !15)
+!15 = distinct !DILexicalBlock(scope: !16, file: !1, line: 4, column: 4)
+!16 = distinct !DILexicalBlock(scope: !6, file: !1, line: 5, column: 5)