[coro] [async] Transfer the debug information from the coro.end.async and coro.async.supend intrinsic to the inlined functions

[aschwaighofer]

Summary

We want to transfer the debug location from the intrinsic instruction
rather than the ones from the function we inline.

rdar://90859530

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static bool replaceCoroEndAsync(AnyCoroEndInst *End) {
 
   auto *EndAsync = dyn_cast<CoroAsyncEndInst>(End);
   if (!EndAsync) {
     Builder.CreateRetVoid();
     return true /*needs cleanup of coro.end block*/;
   }
 
   auto *MustTailCallFunc = EndAsync->getMustTailCallFunction();
+  auto Loc = EndAsync->getDebugLoc();
   if (!MustTailCallFunc) {
     Builder.CreateRetVoid();
     return true /*needs cleanup of coro.end block*/;
   }
 
   // Move the must tail call from the predecessor block into the end block.
   auto *CoroEndBlock = End->getParent();
   auto *MustTailCallFuncBlock = CoroEndBlock->getSinglePredecessor();
   assert(MustTailCallFuncBlock && "Must have a single predecessor block");
   auto It = MustTailCallFuncBlock->getTerminator()->getIterator();
   auto *MustTailCall = cast<CallInst>(&*std::prev(It));
   CoroEndBlock->getInstList().splice(
       End->getIterator(), MustTailCallFuncBlock->getInstList(), MustTailCall);
 
   // Insert the return instruction.
   Builder.SetInsertPoint(End);
   Builder.CreateRetVoid();
   InlineFunctionInfo FnInfo;
 
   // Remove the rest of the block, by splitting it into an unreachable block.
   auto *BB = End->getParent();
   BB->splitBasicBlock(End);
   BB->getTerminator()->eraseFromParent();
 
+  auto InlineRegionBlock = MustTailCall->getParent();
+  auto InlineRegionStart =
+      MustTailCall->getIterator() == InlineRegionBlock->begin()
+          ? None
+          : Optional<BasicBlock::iterator>(
+                std::prev(MustTailCall->getIterator()));
+
   auto InlineRes = InlineFunction(*MustTailCall, FnInfo);
   assert(InlineRes.isSuccess() && "Expected inlining to succeed");
   (void)InlineRes;
 
+  // Apply the debug location of the coro.end.async instruction to the inlined
+  // function.
+
+  auto InlineRegionBegin =
+      InlineRegionStart ? ++(*InlineRegionStart) : InlineRegionBlock->begin();
+  while (Loc) {
+    InlineRegionBegin->setDebugLoc(Loc);
+    if (InlineRegionBegin->isTerminator())
+      break;
+    ++InlineRegionBegin;
+  }
+
   // We have cleaned up the coro.end block above.
   return false;
 }
 
 /// Replace a non-unwind call to llvm.coro.end.
 static void replaceFallthroughCoroEnd(AnyCoroEndInst *End,
                                       const coro::Shape &Shape, Value *FramePtr,
                                       bool InResume, CallGraph *CG) {
for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
     auto *ReturnBB =
         BasicBlock::Create(F.getContext(), "coro.return", &F, NewSuspendBB);
     Branch->setSuccessor(0, ReturnBB);
 
     IRBuilder<> Builder(ReturnBB);
 
     // Insert the call to the tail call function and inline it.
     auto *Fn = Suspend->getMustTailCallFunction();
+    auto Loc = Suspend->getDebugLoc();
     SmallVector<Value *, 8> Args(Suspend->args());
     auto FnArgs = ArrayRef<Value *>(Args).drop_front(
         CoroSuspendAsyncInst::MustTailCallFuncArg + 1);
     auto *TailCall =
         coro::createMustTailCall(Suspend->getDebugLoc(), Fn, FnArgs, Builder);
     Builder.CreateRetVoid();
+
+    auto InlineRegionBlock = TailCall->getParent();
+    auto InlineRegionStart =
+        TailCall->getIterator() == InlineRegionBlock->begin()
+            ? None
+            : Optional<BasicBlock::iterator>(
+                  std::prev(TailCall->getIterator()));
+
     InlineFunctionInfo FnInfo;
     auto InlineRes = InlineFunction(*TailCall, FnInfo);
     assert(InlineRes.isSuccess() && "Expected inlining to succeed");
     (void)InlineRes;
 
+    // Apply the debug location of the async.suspend instruction to the inlined
+    // function.
+    auto InlineRegionBegin =
+        InlineRegionStart ? ++(*InlineRegionStart) : InlineRegionBlock->begin();
+    while (Loc) {
+      InlineRegionBegin->setDebugLoc(Loc);
+      if (InlineRegionBegin->isTerminator())
+        break;
+      ++InlineRegionBegin;
+    }
+
     // Replace the lvm.coro.async.resume intrisic call.
     replaceAsyncResumeFunction(Suspend, Continuation);
   }
 
   assert(Clones.size() == Shape.CoroSuspends.size());
   for (size_t Idx = 0, End = Shape.CoroSuspends.size(); Idx != End; ++Idx) {
     auto *Suspend = Shape.CoroSuspends[Idx];
     auto *Clone = Clones[Idx];

llvm/test/Transforms/Coroutines/coro-async.ll

<{ i32 trunc (
          i64 ptrtoint (void (i8*, %async.task*, %async.actor*)* @my_async_function_pa to i64),
          i64 ptrtoint (i32* getelementptr inbounds (<{ i32, i32 }>, <{ i32, i32 }>* @my_async_function_pa_fp, i32 0, i32 1) to i64)
        )
      to i32),
      i32 8
 }>
 
 ; Function that implements the dispatch to the callee function.
-define swiftcc void @my_async_function.my_other_async_function_fp.apply(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
+define swiftcc void @my_async_function.my_other_async_function_fp.apply(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) !dbg !10 {
+  %callee = bitcast i8* %fnPtr to void(i8*, %async.task*, %async.actor*)* , !dbg !13
+  tail call swiftcc void %callee(i8* %async.ctxt, %async.task* %task, %async.actor* %actor) , !dbg !13
+  ret void
+}
+
+; Another time without debug info.
+define swiftcc void @my_async_function.my_other_async_function_fp.apply2(i8* %fnPtr, i8* %async.ctxt, %async.task* %task, %async.actor* %actor) {
   %callee = bitcast i8* %fnPtr to void(i8*, %async.task*, %async.actor*)*
   tail call swiftcc void %callee(i8* %async.ctxt, %async.task* %task, %async.actor* %actor)
   ret void
 }
 
 declare void @some_user(i64)
 declare void @some_may_write(i64*)
 

 ; CHECK:   [[RETURN_TO_CALLER_ADDR:%.*]] = bitcast i8* [[TYPED_RETURN_TO_CALLER_ADDR]] to i8**
 ; CHECK:   store i8* bitcast (void (i8*, i8*, i8*)* @my_async_functionTQ0_ to i8*), i8** [[RETURN_TO_CALLER_ADDR]]
 ; CHECK:   [[CALLER_CONTEXT_ADDR:%.*]] = bitcast i8* [[CALLEE_CTXT]] to i8**
 ; CHECK:   store i8* %async.ctxt, i8** [[CALLER_CONTEXT_ADDR]]
 ; Make sure the spill is underaligned to the max context alignment (16).
 ; CHECK-O0:   [[VECTOR_SPILL:%.*]] = load <4 x double>, <4 x double>* {{.*}}
 ; CHECK-O0:   [[VECTOR_SPILL_ADDR:%.*]] = getelementptr inbounds %my_async_function.Frame, %my_async_function.Frame* {{.*}}, i32 0, i32 1
 ; CHECK-O0:   store <4 x double> [[VECTOR_SPILL]], <4 x double>* [[VECTOR_SPILL_ADDR]], align 16
-; CHECK:   tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor)
-; CHECK:   ret void
+; CHECK:   tail call swiftcc void @asyncSuspend(i8* [[CALLEE_CTXT]], %async.task* %task, %async.actor* %actor), !dbg ![[DBG1:[0-9]+]]
+; CHECK:   ret void, !dbg ![[DBG1]]
 ; CHECK: }
 
 ; CHECK-LABEL: define internal swiftcc void @my_async_functionTQ0_(i8* nocapture readonly swiftasync %0, i8* %1, i8* nocapture readnone %2)
 ; CHECK-O0-LABEL: define internal swiftcc void @my_async_functionTQ0_(i8* swiftasync %0, i8* %1, i8* %2)
 ; CHECK-SAME: !dbg ![[SP2:[0-9]+]] {
 ; CHECK: entryresume.0:
 ; CHECK:   [[CALLER_CONTEXT_ADDR:%.*]] = bitcast i8* %0 to i8**
 ; CHECK:   [[CALLER_CONTEXT:%.*]] = load i8*, i8** [[CALLER_CONTEXT_ADDR]]
entry:
   %callee_context.caller_context.addr.1 = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0.1, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr.1
   %resume_proj_fun.2 = bitcast i8*(i8*)* @resume_context_projection to i8*
   %callee.2 = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
   %res.2 = call {i8*, i8*, i8*} (i32, i8*, i8*, ...) @llvm.coro.suspend.async(i32 0,
                                                   i8* %resume.func_ptr.1,
                                                   i8* %resume_proj_fun.2,
                                                   void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
-                                                  i8* %callee.2, i8* %callee_context, %async.task* %task.casted, %async.actor *%actor)
+                                                  i8* %callee.2, i8* %callee_context, %async.task* %task.casted, %async.actor *%actor), !dbg !9
 
   call void @llvm.coro.async.context.dealloc(i8* %callee_context)
   %continuation_actor_arg = extractvalue {i8*, i8*, i8*} %res.2, 1
   %actor.2 =  bitcast i8* %continuation_actor_arg to %async.actor*
 
   tail call swiftcc void @asyncReturn(i8* %async.ctxt, %async.task* %task.2, %async.actor* %actor.2)
   call i1 @llvm.coro.end(i8* %hdl, i1 0)
   unreachable
entry:
   %callee_context.caller_context.addr = getelementptr inbounds %async.ctxt, %async.ctxt* %callee_context.0, i32 0, i32 0
   store i8* %async.ctxt, i8** %callee_context.caller_context.addr
   %resume_proj_fun = bitcast i8*(i8*)* @resume_context_projection to i8*
   %callee = bitcast void(i8*, %async.task*, %async.actor*)* @asyncSuspend to i8*
   %res = call {i8*, i8*, i8*, i8*} (i32, i8*, i8*, ...)
          @llvm.coro.suspend.async.sl_p0i8p0i8p0i8p0i8s(i32 256, ;; swiftasync at 0 and swiftself at 1 in resume function
                                                  i8* %resume.func_ptr,
                                                  i8* %resume_proj_fun,
-                                                 void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply,
+                                                 void (i8*, i8*, %async.task*, %async.actor*)* @my_async_function.my_other_async_function_fp.apply2,
                                                  i8* %callee, i8* %callee_context, %async.task* %task, %async.actor *%actor)
 
   call void @llvm.coro.async.context.dealloc(i8* %callee_context)
   %continuation_task_arg = extractvalue {i8*, i8*, i8*, i8*} %res, 3
   %task.2 =  bitcast i8* %continuation_task_arg to %async.task*
   %val = load i64, i64* %proj.1
   call void @some_user(i64 %val)
   %val.2 = load i64, i64* %proj.2

 
 !llvm.dbg.cu = !{!2}
 !llvm.module.flags = !{!0}
 
 !0 = !{i32 2, !"Debug Info Version", i32 3}
 ; CHECK: ![[SP1]] = distinct !DISubprogram(name: "my_async_function",
 ; CHECK-SAME:                              linkageName: "my_async_function",
 ; CHECK-SAME:                              scopeLine: 1
+; CHECK ![[DBG1]] = !DILocation(line: 2, scope: !9)
 !1 = distinct !DISubprogram(name: "my_async_function",
                             linkageName: "my_async_function",
                             scope: !2, file: !3, line: 1, type: !4,
                             scopeLine: 1, spFlags: DISPFlagDefinition, unit: !2)
 ; CHECK: ![[SP2]] = distinct !DISubprogram(name: "my_async_function",
 ; CHECK-SAME:                              linkageName: "my_async_functionTQ0_",
 ; CHECK-SAME:                              scopeLine: 2
 !2 = distinct !DICompileUnit(language: DW_LANG_Swift, file: !3, emissionKind: FullDebug)
!6 = distinct !DISubprogram(name: "my_async_function2",
                             scope: !2, file: !3, line: 1, type: !4,
                             scopeLine: 1, spFlags: DISPFlagDefinition, unit: !2)
 ; CHECK: ![[SP4]] = distinct !DISubprogram(name: "my_async_function2",
 ; CHECK-SAME:                              linkageName: "my_async_function2.resume.0",
 ; CHECK-SAME:                              scopeLine: 1
 !7 = !DILexicalBlockFile(scope: !6, file: !8, discriminator: 0)
 !8 = !DIFile(filename: "/tmp/fake.cpp", directory: "/")
 !9 = !DILocation(line: 2, column: 0, scope: !7)
+
+
+!10 = distinct !DISubprogram(name: "my_async_function.my_other_async_function_fp.apply",
+                            linkageName: "my_async_function.my_other_async_function_fp.apply",
+                            scope: !2, file: !3, line: 1, type: !4,
+                            scopeLine: 1, spFlags: DISPFlagDefinition, unit: !2)
+!11 = !DILexicalBlockFile(scope: !10, file: !8, discriminator: 0)
+!12 = !DIFile(filename: "/tmp/fake2.cpp", directory: "/")
+!13 = !DILocation(line: 2, column: 0, scope: !11)