Merge pull request #5777 from felipepiovezan/felipe/cherrypick_hotcold_fixes_0421

[CherryPick] Fixes for HotColdSplit debug information propagation

Author: felipepiovezan

llvm/include/llvm/IR/DebugInfoMetadata.h

@@ -1590,6 +1590,13 @@ class DILocalScope : public DIScope {
   /// chain.
   DISubprogram *getSubprogram() const;
 
+  /// Traverses the scope chain rooted at RootScope until it hits a Subprogram,
+  /// recreating the chain with "NewSP" instead.
+  static DILocalScope *
+  cloneScopeForSubprogram(DILocalScope &RootScope, DISubprogram &NewSP,
+                          LLVMContext &Ctx,
+                          DenseMap<const MDNode *, MDNode *> &Cache);
+
   /// Get the first non DILexicalBlockFile scope of this scope.
   ///
   /// Return this if it's not a \a DILexicalBlockFIle; otherwise, look up the
@@ -2126,6 +2133,11 @@ class DILexicalBlockBase : public DILocalScope {
 
   Metadata *getRawScope() const { return getOperand(1); }
 
+  void replaceScope(DIScope *Scope) {
+    assert(!isUniqued());
+    setOperand(1, Scope);
+  }
+
   static bool classof(const Metadata *MD) {
     return MD->getMetadataID() == DILexicalBlockKind ||
            MD->getMetadataID() == DILexicalBlockFileKind;

llvm/include/llvm/IR/DebugLoc.h

@@ -86,6 +86,13 @@ namespace llvm {
     /// 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

llvm/lib/IR/DebugInfoMetadata.cpp

@@ -937,6 +937,35 @@ DILocalScope *DILocalScope::getNonLexicalBlockFileScope() const {
   return const_cast<DILocalScope *>(this);
 }
 
+DILocalScope *DILocalScope::cloneScopeForSubprogram(
+    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);
+  }
+
+  // Recreate the scope chain, bottom-up, starting at the new subprogram (or a
+  // cached result).
+  DIScope *UpdatedScope = CachedResult ? CachedResult : &NewSP;
+  for (DIScope *ScopeToUpdate : reverse(ScopeChain)) {
+    TempMDNode ClonedScope = ScopeToUpdate->clone();
+    cast<DILexicalBlockBase>(*ClonedScope).replaceScope(UpdatedScope);
+    UpdatedScope =
+        cast<DIScope>(MDNode::replaceWithUniqued(std::move(ClonedScope)));
+    Cache[ScopeToUpdate] = UpdatedScope;
+  }
+
+  return cast<DILocalScope>(UpdatedScope);
+}
+
 DISubprogram::DISPFlags DISubprogram::getFlag(StringRef Flag) {
   return StringSwitch<DISPFlags>(Flag)
 #define HANDLE_DISP_FLAG(ID, NAME) .Case("DISPFlag" #NAME, SPFlag##NAME)

llvm/lib/IR/DebugLoc.cpp

@@ -67,6 +67,46 @@ void DebugLoc::setImplicitCode(bool ImplicitCode) {
   }
 }
 
+DebugLoc DebugLoc::replaceInlinedAtSubprogram(
+    const DebugLoc &RootLoc, DISubprogram &NewSP, LLVMContext &Ctx,
+    DenseMap<const MDNode *, MDNode *> &Cache) {
+  SmallVector<DILocation *> LocChain;
+  DILocation *CachedResult = nullptr;
+
+  // Collect the inline chain, stopping if we find a location that has already
+  // been processed.
+  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, that is,
+    // the DILocation whose scope ends in the Subprogram to be replaced.
+    DILocation *LocToUpdate = LocChain.pop_back_val();
+    DIScope *NewScope = DILocalScope::cloneScopeForSubprogram(
+        *LocToUpdate->getScope(), NewSP, Ctx, Cache);
+    UpdatedLoc = DILocation::get(Ctx, LocToUpdate->getLine(),
+                                 LocToUpdate->getColumn(), NewScope);
+    Cache[LocToUpdate] = UpdatedLoc;
+  }
+
+  // Recreate the location chain, bottom-up, starting at the new scope (or a
+  // cached result).
+  for (const DILocation *LocToUpdate : reverse(LocChain)) {
+    UpdatedLoc =
+        DILocation::get(Ctx, LocToUpdate->getLine(), LocToUpdate->getColumn(),
+                        LocToUpdate->getScope(), UpdatedLoc);
+    Cache[LocToUpdate] = UpdatedLoc;
+  }
+
+  return UpdatedLoc;
+}
+
 DebugLoc DebugLoc::appendInlinedAt(const DebugLoc &DL, DILocation *InlinedAt,
                                    LLVMContext &Ctx,
                                    DenseMap<const MDNode *, MDNode *> &Cache) {

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -1551,18 +1551,25 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
   //     point to a variable in the wrong scope.
   SmallDenseMap<DINode *, DINode *> RemappedMetadata;
   SmallVector<Instruction *, 4> DebugIntrinsicsToDelete;
+  DenseMap<const MDNode *, MDNode *> Cache;
   for (Instruction &I : instructions(NewFunc)) {
     auto *DII = dyn_cast<DbgInfoIntrinsic>(&I);
     if (!DII)
       continue;
 
-    // Point the intrinsic to a fresh label within the new function.
+    // Point the intrinsic to a fresh label within the new function if the
+    // intrinsic was not inlined from some other function.
     if (auto *DLI = dyn_cast<DbgLabelInst>(&I)) {
+      if (DLI->getDebugLoc().getInlinedAt())
+        continue;
       DILabel *OldLabel = DLI->getLabel();
       DINode *&NewLabel = RemappedMetadata[OldLabel];
-      if (!NewLabel)
-        NewLabel = DILabel::get(Ctx, NewSP, OldLabel->getName(),
+      if (!NewLabel) {
+        DILocalScope *NewScope = DILocalScope::cloneScopeForSubprogram(
+            *OldLabel->getScope(), *NewSP, Ctx, Cache);
+        NewLabel = DILabel::get(Ctx, NewScope, OldLabel->getName(),
                                 OldLabel->getFile(), OldLabel->getLine());
+      }
       DLI->setArgOperand(0, MetadataAsValue::get(Ctx, NewLabel));
       continue;
     }
@@ -1583,17 +1590,23 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
       DebugIntrinsicsToDelete.push_back(DVI);
       continue;
     }
-
-    // Point the intrinsic to a fresh variable within the new function.
-    DILocalVariable *OldVar = DVI->getVariable();
-    DINode *&NewVar = RemappedMetadata[OldVar];
-    if (!NewVar)
-      NewVar = DIB.createAutoVariable(
-          NewSP, OldVar->getName(), OldVar->getFile(), OldVar->getLine(),
-          OldVar->getType(), /*AlwaysPreserve=*/false, DINode::FlagZero,
-          OldVar->getAlignInBits());
-    DVI->setVariable(cast<DILocalVariable>(NewVar));
+    // 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));
+    }
   }
+
   for (auto *DII : DebugIntrinsicsToDelete)
     DII->eraseFromParent();
   DIB.finalizeSubprogram(NewSP);
@@ -1602,13 +1615,13 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
   // function.
   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 * {
+    auto updateLoopInfoLoc = [&Ctx, &Cache, NewSP](Metadata *MD) -> Metadata * {
       if (auto *Loc = dyn_cast_or_null<DILocation>(MD))
-        return DILocation::get(Ctx, Loc->getLine(), Loc->getColumn(), NewSP,
-                               nullptr);
+        return DebugLoc::replaceInlinedAtSubprogram(Loc, *NewSP, Ctx, Cache);
       return MD;
     };
     updateLoopMetadataDebugLocations(I, updateLoopInfoLoc);

llvm/test/Transforms/HotColdSplit/split-out-dbg-label.ll

@@ -12,11 +12,18 @@ target triple = "x86_64-apple-macosx10.14.0"
 
 ; CHECK-LABEL: define {{.*}}@foo.cold.1
 ; CHECK: llvm.dbg.label(metadata [[LABEL:![0-9]+]]), !dbg [[LINE:![0-9]+]]
+; CHECK: llvm.dbg.label(metadata [[LABEL_IN_INLINE_ME:![0-9]+]]), !dbg [[LINE2:![0-9]+]]
+; CHECK: llvm.dbg.label(metadata [[SCOPED_LABEL:![0-9]+]]), !dbg [[LINE]]
 
 ; CHECK: [[FILE:![0-9]+]] = !DIFile
+; CHECK: [[INLINE_ME_SCOPE:![0-9]+]] = distinct !DISubprogram(name: "inline_me"
 ; CHECK: [[SCOPE:![0-9]+]] = distinct !DISubprogram(name: "foo.cold.1"
 ; CHECK: [[LINE]] = !DILocation(line: 1, column: 1, scope: [[SCOPE]]
 ; CHECK: [[LABEL]] = !DILabel(scope: [[SCOPE]], name: "bye", file: [[FILE]], line: 28
+; CHECK: [[LABEL_IN_INLINE_ME]] = !DILabel(scope: [[INLINE_ME_SCOPE]], name: "label_in_@inline_me", file: [[FILE]], line: 29
+; CHECK: [[LINE2]] = !DILocation(line: 2, column: 2, scope: [[INLINE_ME_SCOPE]], inlinedAt: [[LINE]]
+; CHECK: [[SCOPED_LABEL]] = !DILabel(scope: [[SCOPE_IN_FOO:![0-9]+]], name: "scoped_label_in_foo", file: [[FILE]], line: 30
+; CHECK: [[SCOPE_IN_FOO]] = !DILexicalBlock(scope: [[SCOPE]], file: [[FILE]], line: 31, column: 31)
 
 define void @foo(i32 %arg1) !dbg !6 {
 entry:
@@ -28,6 +35,8 @@ if.then:                                          ; preds = %entry
 
 if.end:                                           ; preds = %entry
   call void @llvm.dbg.label(metadata !12), !dbg !11
+  call void @llvm.dbg.label(metadata !14), !dbg !15
+  call void @llvm.dbg.label(metadata !16), !dbg !11
   call void @sink()
   ret void
 }
@@ -36,6 +45,10 @@ declare void @llvm.dbg.label(metadata)
 
 declare void @sink() cold
 
+define void @inline_me() !dbg !13 {
+  ret void
+}
+
 !llvm.dbg.cu = !{!0}
 !llvm.debugify = !{!3, !4}
 !llvm.module.flags = !{!5}
@@ -53,3 +66,8 @@ declare void @sink() cold
 !10 = !DIBasicType(name: "ty32", size: 32, encoding: DW_ATE_unsigned)
 !11 = !DILocation(line: 1, column: 1, scope: !6)
 !12 = !DILabel(scope: !6, name: "bye", file: !1, line: 28)
+!13 = 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)
+!14 = !DILabel(scope: !13, name: "label_in_@inline_me", file: !1, line: 29)
+!15 = !DILocation(line: 2, column: 2, scope: !13, inlinedAt: !11)
+!16 = !DILabel(scope: !17, name: "scoped_label_in_foo", file: !1, line: 30)
+!17 = distinct !DILexicalBlock(scope: !6, file: !1, line: 31, column: 31)

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

@@ -28,15 +28,42 @@ target triple = "x86_64-apple-macosx10.14.0"
 ; - 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]+]]
+; CHECK-NEXT: call void @sink(i32 [[ADD1]]), !dbg [[LINE3:![0-9]+]]
+
+; CHECK-NEXT: call void @llvm.dbg.value(metadata i32 [[ADD1]]
+; CHECK-SAME:      metadata [[VAR_FROM_INLINE_ME:![0-9]+]]
+; CHECK-SAME:      !dbg [[LINE2]]
+
+; CHECK-NEXT: call void @llvm.dbg.value(metadata i32 [[ADD1]]
+; CHECK-SAME:      metadata [[VAR2:![0-9]+]]
+; CHECK-SAME:     !dbg [[LINE4:![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: [[INLINE_ME_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: [[INLINE_ME_SCOPE]]
+; CHECK-SAME:                            inlinedAt: [[LINE3]]
+; CHECK: [[LINE3]] =          !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)
+
+; CHECK: [[VAR_FROM_INLINE_ME]] = !DILocalVariable(name: "var_from_inline_me",
+; CHECK-SAME:                                      scope: [[INLINE_ME_SCOPE]]
+
+; CHECK: [[VAR2]] = !DILocalVariable(name: "var_from_scope_in_foo",
+; CHECK-SAME:                        scope: [[NEWSCOPE2:![0-9]+]]
+; CHECK: [[NEWSCOPE2]] = !DILexicalBlock(scope: [[NEWSCOPE]], file: [[FILE]], line: 7, column: 7)
+; CHECK: [[LINE4]] =     !DILocation(line: 6, column: 6, scope: [[NEWSCOPE2]]
+
 define void @foo(i32 %arg1) !dbg !6 {
 entry:
   %var = add i32 0, 0, !dbg !11
@@ -52,13 +79,21 @@ if.end:                                           ; preds = %entry
   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
+  call void @llvm.dbg.value(metadata i32 %add1, metadata !17, metadata !DIExpression()), !dbg !13 ; variable from @inline_me, should preserve scope in !17.
+  call void @llvm.dbg.value(metadata i32 %add1, metadata !18, metadata !DIExpression()), !dbg !19 ; variable 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}
@@ -75,3 +110,12 @@ declare void @sink(i32) cold
 !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)
+!17 = !DILocalVariable(name: "var_from_inline_me", scope: !12, file: !1, line: 1, type: !10)
+!18 = !DILocalVariable(name: "var_from_scope_in_foo", scope: !20, file: !1, line: 1, type: !10)
+!19 = !DILocation(line: 6, column: 6, scope: !20)
+!20 = distinct !DILexicalBlock(scope: !6, file: !1, line: 7, column: 7)

llvm/test/Transforms/HotColdSplit/update-split-loop-metadata.ll

@@ -11,8 +11,10 @@ target triple = "x86_64-apple-macosx10.14.0"
 
 ; The scope for these debug locations should be @basic.cold.1, not @basic.
 ; CHECK: [[SCOPE:![0-9]+]] = distinct !DISubprogram(name: "basic.cold.1"
-; CHECK: [[LOOP_MD]] = distinct !{[[LOOP_MD]], [[LINE:![0-9]+]], [[LINE]]}
+; CHECK: [[LOOP_MD]] = distinct !{[[LOOP_MD]], [[LINE:![0-9]+]], [[LINE2:![0-9]+]]}
 ; CHECK: [[LINE]] = !DILocation(line: 1, column: 1, scope: [[SCOPE]])
+; CHECK: [[LINE2]] = !DILocation(line: 2, column: 2, scope: [[LEX_SCOPE:![0-9]+]])
+; CHECK: [[LEX_SCOPE]] = !DILexicalBlock(scope: [[SCOPE]], file: !{{[0-9]+}}, line: 3, column: 3)
 
 define void @basic(i32* %p, i32 %k) !dbg !6 {
 entry:
@@ -55,4 +57,6 @@ declare void @sink() cold
 !7 = !DISubroutineType(types: !2)
 !8 = !{}
 !9 = !DILocation(line: 1, column: 1, scope: !6)
-!10 = distinct !{!10, !9, !9}
+!10 = distinct !{!10, !9, !11}
+!11 = !DILocation(line: 2, column: 2, scope: !12)
+!12 = !DILexicalBlock(scope: !6, file: !1, line: 3, column: 3)