[KeyInstr][Clang] Ret atom (#134652)

This patch is part of a stack that teaches Clang to generate Key Instructions
metadata for C and C++.

When returning a value, stores to the `retval` allocas and branches to `return`
block are put in the same atom group. They are both rank 1, which could in
theory introduce an extra step in some optimized code. This low risk currently
feels an acceptable for keeping the code a bit simpler (as opposed to adding
scaffolding to make the store rank 2).

In the case of a single return (no control flow) the return instruction inherits
the atom group of the branch to the return block when the blocks get folded
togather.

RFC:
https://discourse.llvm.org/t/rfc-improving-is-stmt-placement-for-better-interactive-debugging/82668

The feature is only functional in LLVM if LLVM is built with CMake flag
LLVM_EXPERIMENTAL_KEY_INSTRUCTIONs. Eventually that flag will be removed.

Author: OCHyams

clang/lib/CodeGen/CGCall.cpp

@@ -3929,9 +3929,9 @@ llvm::Value *CodeGenFunction::EmitCMSEClearRecord(llvm::Value *Src,
   return R;
 }
 
-void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
-                                         bool EmitRetDbgLoc,
-                                         SourceLocation EndLoc) {
+void CodeGenFunction::EmitFunctionEpilog(
+    const CGFunctionInfo &FI, bool EmitRetDbgLoc, SourceLocation EndLoc,
+    uint64_t RetKeyInstructionsSourceAtom) {
   if (FI.isNoReturn()) {
     // Noreturn functions don't return.
     EmitUnreachable(EndLoc);
@@ -3946,7 +3946,11 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   // Functions with no result always return void.
   if (!ReturnValue.isValid()) {
-    Builder.CreateRetVoid();
+    auto *I = Builder.CreateRetVoid();
+    if (RetKeyInstructionsSourceAtom)
+      addInstToSpecificSourceAtom(I, nullptr, RetKeyInstructionsSourceAtom);
+    else
+      addInstToNewSourceAtom(I, nullptr);
     return;
   }
 
@@ -4126,6 +4130,12 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   if (RetDbgLoc)
     Ret->setDebugLoc(std::move(RetDbgLoc));
+
+  llvm::Value *Backup = RV ? Ret->getOperand(0) : nullptr;
+  if (RetKeyInstructionsSourceAtom)
+    addInstToSpecificSourceAtom(Ret, Backup, RetKeyInstructionsSourceAtom);
+  else
+    addInstToNewSourceAtom(Ret, Backup);
 }
 
 void CodeGenFunction::EmitReturnValueCheck(llvm::Value *RV) {

clang/lib/CodeGen/CGStmt.cpp

@@ -1605,6 +1605,7 @@ static bool isSwiftAsyncCallee(const CallExpr *CE) {
 /// if the function returns void, or may be missing one if the function returns
 /// non-void.  Fun stuff :).
 void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
+  ApplyAtomGroup Grp(getDebugInfo());
   if (requiresReturnValueCheck()) {
     llvm::Constant *SLoc = EmitCheckSourceLocation(S.getBeginLoc());
     auto *SLocPtr =
@@ -1680,16 +1681,19 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
     // If this function returns a reference, take the address of the expression
     // rather than the value.
     RValue Result = EmitReferenceBindingToExpr(RV);
-    Builder.CreateStore(Result.getScalarVal(), ReturnValue);
+    auto *I = Builder.CreateStore(Result.getScalarVal(), ReturnValue);
+    addInstToCurrentSourceAtom(I, I->getValueOperand());
   } else {
     switch (getEvaluationKind(RV->getType())) {
     case TEK_Scalar: {
       llvm::Value *Ret = EmitScalarExpr(RV);
-      if (CurFnInfo->getReturnInfo().getKind() == ABIArgInfo::Indirect)
+      if (CurFnInfo->getReturnInfo().getKind() == ABIArgInfo::Indirect) {
         EmitStoreOfScalar(Ret, MakeAddrLValue(ReturnValue, RV->getType()),
                           /*isInit*/ true);
-      else
-        Builder.CreateStore(Ret, ReturnValue);
+      } else {
+        auto *I = Builder.CreateStore(Ret, ReturnValue);
+        addInstToCurrentSourceAtom(I, I->getValueOperand());
+      }
       break;
     }
     case TEK_Complex:

clang/lib/CodeGen/CodeGenFunction.cpp

@@ -444,8 +444,10 @@ void CodeGenFunction::FinishFunction(SourceLocation EndLoc) {
 
   // Reset the debug location to that of the simple 'return' expression, if any
   // rather than that of the end of the function's scope '}'.
+  uint64_t RetKeyInstructionsAtomGroup = Loc ? Loc->getAtomGroup() : 0;
   ApplyDebugLocation AL(*this, Loc);
-  EmitFunctionEpilog(*CurFnInfo, EmitRetDbgLoc, EndLoc);
+  EmitFunctionEpilog(*CurFnInfo, EmitRetDbgLoc, EndLoc,
+                     RetKeyInstructionsAtomGroup);
   EmitEndEHSpec(CurCodeDecl);
 
   assert(EHStack.empty() &&

clang/lib/CodeGen/CodeGenFunction.h

@@ -2551,9 +2551,12 @@ class CodeGenFunction : public CodeGenTypeCache {
                           const FunctionArgList &Args);
 
   /// EmitFunctionEpilog - Emit the target specific LLVM code to return the
-  /// given temporary.
+  /// given temporary. Specify the source location atom group (Key Instructions
+  /// debug info feature) for the `ret` using \p RetKeyInstructionsSourceAtom.
+  /// If it's 0, the `ret` will get added to a new source atom group.
   void EmitFunctionEpilog(const CGFunctionInfo &FI, bool EmitRetDbgLoc,
-                          SourceLocation EndLoc);
+                          SourceLocation EndLoc,
+                          uint64_t RetKeyInstructionsSourceAtom);
 
   /// Emit a test that checks if the return value \p RV is nonnull.
   void EmitReturnValueCheck(llvm::Value *RV);

clang/test/DebugInfo/KeyInstructions/agg.c

@@ -24,6 +24,8 @@ void fun(Struct a) {
 // CHECK: %matins = insertelement <25 x float> %3, float 0.000000e+00, i64 0, !dbg [[G4R2:!.*]]
 // CHECK: store <25 x float> %matins, ptr @m{{.*}}, !dbg [[G4R1:!.*]]
   m[0][0] = 0;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -32,3 +34,4 @@ void fun(Struct a) {
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
 // CHECK: [[G4R2]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 2)
 // CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)

clang/test/DebugInfo/KeyInstructions/assign-scalar.c

@@ -58,6 +58,7 @@ void fun() {
 // CHECK-NEXT: %inc4 = add i64 %7, 1, !dbg [[G11R2:!.*]]
 // CHECK-NEXT: store i64 %inc4, ptr @i{{.*}}, !dbg [[G11R1:!.*]]
     g = ++h, ++i;
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -83,3 +84,4 @@ void fun() {
 // CHECK: [[load_i_loc]] = !DILocation(line: [[#]], column: [[#]], scope: ![[#]])
 // CHECK: [[G11R2]] = !DILocation({{.*}}, atomGroup: 11, atomRank: 2)
 // CHECK: [[G11R1]] = !DILocation({{.*}}, atomGroup: 11, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 12, atomRank: 1)

clang/test/DebugInfo/KeyInstructions/bitfield.cpp

@@ -10,7 +10,10 @@ void foo(int x, S s) {
 // CHECK: %bf.set = or i8 %bf.clear, %bf.value, !dbg [[G1R2:!.*]]
 // CHECK: store i8 %bf.set, ptr %s, align 4, !dbg [[G1R1:!.*]]
   s.a = x;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)

clang/test/DebugInfo/KeyInstructions/builtin.c

@@ -64,6 +64,7 @@ void fun() {
 // CHECK-NEXT: %4 = trunc i32 %3 to i8, !dbg [[G15R2:!.*]]
 // CHECK-NEXT: call void @llvm.memset{{.*}}, !dbg [[G15R1:!.*]]
     __builtin_memset(f4, v, sizeof(float) * 4);
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -86,3 +87,4 @@ void fun() {
 // CHECK: [[G15R3]] = !DILocation({{.*}}, atomGroup: 15, atomRank: 3)
 // CHECK: [[G15R2]] = !DILocation({{.*}}, atomGroup: 15, atomRank: 2)
 // CHECK: [[G15R1]] = !DILocation({{.*}}, atomGroup: 15, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 16, atomRank: 1)

clang/test/DebugInfo/KeyInstructions/cast.c

@@ -13,8 +13,10 @@ void a() {
 // CHECK: %conv = fptosi float %0 to i32{{.*}}, !dbg [[G1R2:!.*]]
 // CHECK: store i32 %conv, ptr %a{{.*}}, !dbg [[G1R1:!.*]]
     int a = g;
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R3]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 3)
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)

clang/test/DebugInfo/KeyInstructions/coerced-packed.c

@@ -14,7 +14,9 @@ void f() {
 // CHECK: [[call:%.*]] = call i40{{.*}}getS{{.*}}, !dbg [[G1R2:!.*]]
 // CHECK: store i40 [[call]], ptr %s, align 1, !dbg [[G1R1:!.*]]
     S s = getS();
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[RET]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)