[CodeExtractor] Improve debug info for input values.

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

@@ -290,7 +290,8 @@ class CodeExtractorAnalysisCache {
     void emitFunctionBody(const ValueSet &inputs, const ValueSet &outputs,
                           const ValueSet &StructValues, Function *newFunction,
                           StructType *StructArgTy, BasicBlock *header,
-                          const ValueSet &SinkingCands);
+                          const ValueSet &SinkingCands,
+                          SmallVectorImpl<Value *> &NewValues);
 
     /// Generates a Basic Block that calls the extracted function.
     CallInst *emitReplacerCall(const ValueSet &inputs, const ValueSet &outputs,

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -1239,11 +1239,19 @@ static void eraseDebugIntrinsicsWithNonLocalRefs(Function &F) {
   }
 }
 
-/// Fix up the debug info in the old and new functions by pointing line
-/// locations and debug intrinsics to the new subprogram scope, and by deleting
-/// intrinsics which point to values outside of the new function.
+/// Fix up the debug info in the old and new functions. Following changes are
+/// done.
+/// 1. If a debug record points to a value that has been replaced, update the
+/// record to use the new value.
+/// 2. If an Input value that has been replaced was used as a location of a
+/// debug record in the Parent function, then materealize a similar record in
+/// the new function.
+/// 3. Point line locations and debug intrinsics to the new subprogram scope
+/// 4. Remove intrinsics which point to values outside of the new function.
 static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
-                                         CallInst &TheCall) {
+                                         CallInst &TheCall,
+                                         const SetVector<Value *> &Inputs,
+                                         ArrayRef<Value *> NewValues) {
   DISubprogram *OldSP = OldFunc.getSubprogram();
   LLVMContext &Ctx = OldFunc.getContext();
 
@@ -1270,14 +1278,48 @@ static void fixupDebugInfoPostExtraction(Function &OldFunc, Function &NewFunc,
       /*LineNo=*/0, SPType, /*ScopeLine=*/0, DINode::FlagZero, SPFlags);
   NewFunc.setSubprogram(NewSP);
 
+  auto UpdateOrInsertDebugRecord = [&](auto *DR, Value *OldLoc, Value *NewLoc,
+                                       DIExpression *Expr, bool Declare) {
+    if (DR->getParent()->getParent() == &NewFunc)
+      DR->replaceVariableLocationOp(OldLoc, NewLoc);
+    else {
+      if (Declare)
+        DIB.insertDeclare(NewLoc, DR->getVariable(), Expr, DR->getDebugLoc(),
+                          &NewFunc.getEntryBlock());
+      else
+        DIB.insertDbgValueIntrinsic(
+            NewLoc, DR->getVariable(), Expr, DR->getDebugLoc(),
+            NewFunc.getEntryBlock().getTerminator()->getIterator());
+    }
+  };
+  for (auto [Input, NewVal] : zip_equal(Inputs, NewValues)) {
+    SmallVector<DbgVariableIntrinsic *, 1> DbgUsers;
+    SmallVector<DbgVariableRecord *, 1> DPUsers;
+    findDbgUsers(DbgUsers, Input, &DPUsers);
+    DIExpression *Expr = DIB.createExpression();
+
+    // Iterate the debud users of the Input values. If they are in the extracted
+    // function then update their location with the new value. If they are in
+    // the parent function then create a similar debug record.
+    for (auto *DVI : DbgUsers)
+      UpdateOrInsertDebugRecord(DVI, Input, NewVal, Expr,
+                                isa<DbgDeclareInst>(DVI));
+    for (auto *DVR : DPUsers)
+      UpdateOrInsertDebugRecord(DVR, Input, NewVal, Expr, DVR->isDbgDeclare());
+  }
+
   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)))
+    // Location is invalid if it isn't a constant, an instruction or an
+    // argument, or is an instruction/argument but isn't in the new function.
+    if (!Location || (!isa<Constant>(Location) && !isa<Argument>(Location) &&
+                      !isa<Instruction>(Location)))
       return true;
-    Instruction *LocationInst = dyn_cast<Instruction>(Location);
-    return LocationInst && LocationInst->getFunction() != &NewFunc;
+
+    if (Argument *Arg = dyn_cast<Argument>(Location))
+      return Arg->getParent() != &NewFunc;
+    else if (Instruction *LocationInst = dyn_cast<Instruction>(Location))
+      return LocationInst->getFunction() != &NewFunc;
+    return false;
   };
 
   // Debug intrinsics in the new function need to be updated in one of two
@@ -1506,9 +1548,10 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
       inputs, outputs, EntryFreq, oldFunction->getName() + "." + SuffixToUse,
       StructValues, StructTy);
   newFunction->IsNewDbgInfoFormat = oldFunction->IsNewDbgInfoFormat;
+  SmallVector<Value *> NewValues;
 
   emitFunctionBody(inputs, outputs, StructValues, newFunction, StructTy, header,
-                   SinkingCands);
+                   SinkingCands, NewValues);
 
   std::vector<Value *> Reloads;
   CallInst *TheCall = emitReplacerCall(
@@ -1518,7 +1561,8 @@ CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC,
   insertReplacerCall(oldFunction, header, TheCall->getParent(), outputs,
                      Reloads, ExitWeights);
 
-  fixupDebugInfoPostExtraction(*oldFunction, *newFunction, *TheCall);
+  fixupDebugInfoPostExtraction(*oldFunction, *newFunction, *TheCall, inputs,
+                               NewValues);
 
   LLVM_DEBUG(if (verifyFunction(*newFunction, &errs())) {
     newFunction->dump();
@@ -1583,7 +1627,8 @@ Type *CodeExtractor::getSwitchType() {
 void CodeExtractor::emitFunctionBody(
     const ValueSet &inputs, const ValueSet &outputs,
     const ValueSet &StructValues, Function *newFunction,
-    StructType *StructArgTy, BasicBlock *header, const ValueSet &SinkingCands) {
+    StructType *StructArgTy, BasicBlock *header, const ValueSet &SinkingCands,
+    SmallVectorImpl<Value *> &NewValues) {
   Function *oldFunction = header->getParent();
   LLVMContext &Context = oldFunction->getContext();
 
@@ -1615,7 +1660,6 @@ void CodeExtractor::emitFunctionBody(
 
   // Rewrite all users of the inputs in the extracted region to use the
   // arguments (or appropriate addressing into struct) instead.
-  SmallVector<Value *> NewValues;
   for (unsigned i = 0, e = inputs.size(), aggIdx = 0; i != e; ++i) {
     Value *RewriteVal;
     if (StructValues.contains(inputs[i])) {

llvm/test/Transforms/CodeExtractor/LoopExtractor_alloca.ll

@@ -19,6 +19,7 @@
 
 ; CHECK-LABEL: define internal void @test.loop1(ptr %v1)
 ; CHECK-NEXT: newFuncRoot:
+; CHECK-NEXT: #dbg_value
 ; CHECK-NEXT:   br
 
 define void @test() {

llvm/test/Transforms/CodeExtractor/input-value-debug.ll

@@ -0,0 +1,54 @@
+; RUN: opt -passes=hotcoldsplit -hotcoldsplit-threshold=0 -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @foo(i32 %a, i32 %b) !dbg !2 {
+entry:
+   %1 = alloca i32, i64 1, align 4
+   %2 = alloca i32, i64 1, align 4
+   store i32 %a, ptr %1, align 4
+   #dbg_declare(ptr %1, !8, !DIExpression(), !1)
+   #dbg_value(i32 %b, !9, !DIExpression(), !1)
+   %tobool = icmp eq i32 %a, 0
+   br i1 %tobool, label %if.then, label %if.end
+if.then:                                          ; preds = %entry
+  ret void
+
+if.end:                                           ; preds = %entry
+   store i32 10, ptr %1, align 4
+   %3 = add i32 %b, 1
+   store i32 1, ptr %2, align 4
+   call void @sink(i32 %3)
+   #dbg_declare(ptr %2, !10, !DIExpression(), !1)
+   ret void
+}
+
+declare void @sink(i32) cold
+
+!llvm.dbg.cu = !{!6}
+!llvm.module.flags = !{!0}
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = !DILocation(line: 11, column: 7, scope: !2)
+!2 = distinct !DISubprogram(name: "foo", scope: !3, file: !3, type: !4, spFlags: DISPFlagDefinition, unit: !6)
+!3 = !DIFile(filename: "test.c", directory: "")
+!4 = !DISubroutineType(cc: DW_CC_program, types: !5)
+!5 = !{}
+!6 = distinct !DICompileUnit(language: DW_LANG_C, file: !3)
+!7 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+!8 = !DILocalVariable(name: "a", scope: !2, file: !3, type: !7)
+!9 = !DILocalVariable(name: "b", scope: !2, file: !3, type: !7)
+!10 = !DILocalVariable(name: "c", scope: !2, file: !3, type: !7)
+
+; CHECK: define {{.*}}@foo.cold.1(ptr %[[ARG0:.*]], i32 %[[ARG1:.*]], ptr %[[ARG2:.*]]){{.*}} !dbg ![[FN:.*]] {
+; CHECK-NEXT: newFuncRoot:
+; CHECK-NEXT: #dbg_declare(ptr %[[ARG0]], ![[V1:[0-9]+]], {{.*}})
+; CHECK-NEXT: #dbg_value(i32 %[[ARG1]], ![[V2:[0-9]+]], {{.*}})
+; CHECK-NEXT: br
+; CHECK: if.end:
+; CHECK:     #dbg_declare(ptr %[[ARG2]], ![[V3:[0-9]+]], {{.*}})
+; CHECK: }
+
+; CHECK: ![[V1]] = !DILocalVariable(name: "a", scope: ![[FN]]{{.*}})
+; CHECK: ![[V2]] = !DILocalVariable(name: "b", scope: ![[FN]]{{.*}})
+; CHECK: ![[V3]] = !DILocalVariable(name: "c", scope: ![[FN]]{{.*}})

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

@@ -8,10 +8,6 @@ target triple = "x86_64-apple-macosx10.14.0"
 
 ; CHECK-LABEL: define {{.*}}@foo.cold.1
 
-; - The llvm.dbg.value intrinsic pointing to an argument in @foo (%arg1) is
-;   dropped
-; CHECK-NOT: #dbg_value
-
 ; - Instructions without locations in the original function have no
 ;   location in the new function
 ; CHECK:      [[ADD1:%.*]] = add i32 %{{.*}}, 1{{$}}

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

@@ -11,6 +11,7 @@
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
@@ -731,4 +732,87 @@ TEST(CodeExtractor, OpenMPAggregateArgs) {
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
+
+TEST(CodeExtractor, ArgsDebugInfo) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"ir(
+
+  target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+  target triple = "x86_64-unknown-linux-gnu"
+
+  define void @foo(i32 %a, i32 %b) !dbg !2 {
+    %1 = alloca i32, i64 1, align 4, !dbg !1
+    store i32 %a, ptr %1, align 4, !dbg !1
+    #dbg_declare(ptr %1, !8, !DIExpression(), !1)
+    #dbg_value(i32 %b, !9, !DIExpression(), !1)
+    br label %entry
+
+  entry:
+    br label %extract
+
+  extract:
+    store i32 10, ptr %1, align 4, !dbg !1
+    %2 = add i32 %b, 1, !dbg !1
+    br label %exit
+
+  exit:
+    ret void
+  }
+  !llvm.dbg.cu = !{!6}
+  !llvm.module.flags = !{!0}
+  !0 = !{i32 2, !"Debug Info Version", i32 3}
+  !1 = !DILocation(line: 11, column: 7, scope: !2)
+  !2 = distinct !DISubprogram(name: "foo", scope: !3, file: !3, type: !4, spFlags: DISPFlagDefinition, unit: !6)
+  !3 = !DIFile(filename: "test.f90", directory: "")
+  !4 = !DISubroutineType(cc: DW_CC_program, types: !5)
+  !5 = !{null}
+  !6 = distinct !DICompileUnit(language: DW_LANG_Fortran95, file: !3)
+  !7 = !DIBasicType(name: "integer", size: 32, encoding: DW_ATE_signed)
+  !8 = !DILocalVariable(name: "a", scope: !2, file: !3, type: !7)
+  !9 = !DILocalVariable(name: "b", scope: !2, file: !3, type: !7)
+
+  )ir",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 1> Blocks{getBlockByName(Func, "extract")};
+
+  auto TestExtracted = [&](bool AggregateArgs) {
+    CodeExtractor CE(Blocks, /* DominatorTree */ nullptr, AggregateArgs);
+    EXPECT_TRUE(CE.isEligible());
+    CodeExtractorAnalysisCache CEAC(*Func);
+    SetVector<Value *> Inputs, Outputs, SinkingCands, HoistingCands;
+    BasicBlock *CommonExit = nullptr;
+    CE.findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
+    CE.findInputsOutputs(Inputs, Outputs, SinkingCands);
+    Function *Outlined = CE.extractCodeRegion(CEAC, Inputs, Outputs);
+    Outlined->dump();
+    EXPECT_TRUE(Outlined);
+    BasicBlock &EB = Outlined->getEntryBlock();
+    Instruction *Term = EB.getTerminator();
+    EXPECT_TRUE(Term);
+    EXPECT_TRUE(Term->hasDbgRecords());
+    for (DbgVariableRecord &DVR : filterDbgVars(Term->getDbgRecordRange())) {
+      DILocalVariable *Var = DVR.getVariable();
+      EXPECT_TRUE(Var);
+      if (DVR.isDbgDeclare())
+        EXPECT_TRUE(Var->getName() == "a");
+      else
+        EXPECT_TRUE(Var->getName() == "b");
+      for (Value *Loc : DVR.location_ops()) {
+        if (Instruction *I = dyn_cast<Instruction>(Loc))
+          EXPECT_TRUE(I->getParent() == &EB);
+        else if (Argument *A = dyn_cast<Argument>(Loc))
+          EXPECT_TRUE(A->getParent() == Outlined);
+      }
+    }
+    EXPECT_FALSE(verifyFunction(*Outlined));
+  };
+  // Test with both true and false for AggregateArgs.
+  TestExtracted(true);
+  TestExtracted(false);
+  EXPECT_FALSE(verifyFunction(*Func));
+}
+
 } // end anonymous namespace