[KeyInstr][Clang] Ret atom

[OCHyams]

[KeyInstr][Clang] Ret atom

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

The Key Instructions project is introduced, including a "quick summary" section
at the top which adds context for this PR, here:
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.

The Clang-side work is demoed here:
#130943

[KeyInstr][Clang] Update tests with ret atoms

[OCHyams]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• [NFC][KeyInstr][Clang] Add some additional tests #134654
• [KeyInstr][Clang] Coerced store atoms #134653
• [KeyInstr][Clang] Ret atom #134652 👈 (View in Graphite)
• [KeyIntsr][Clang] Builtins alloca auto-init atom #134651
• [KeyInstr][Clang] Assign matrix element atom #134650
• [KeyInstr][Clang] Assign vector element atom #134649
• [KeyInstr][Clang] Bitfield atom #134648
• [KeyInstr][Clang] For range stmt atoms #134647
• [KeyInstr][Clang] For stmt atom #134646
• [KeyInstr][Clang] While stmt atom #134645
• [KeyInstr][Clang] Do stmt atom #134644
• [KeyInstr][Clang] Switch stmt atom #134643
• [KeyInstr][Clang] If stmt atom #134642
• [KeyInstr][Clang] Catch variable init atom #134641
• [KeyInstr][Clang] Member initalization atom #134640
• [KeyInstr][Clang] Aggregate init + copy #134639
• [KeyInstr] Complex assignment atoms #134638
• [KeyInstr][Clang] Assignment atom group #134637
• [KeyInstr][Clang] Static variable init atom #134636
• [KeyInstr][Clang] Agg init atom #134635
• [Clang][NFC] Move some static functions into CodeGenFunction #134634
• [KeyInstr][Clang] Scalar init atom #134633
• [KeyInstr][Clang] Add ApplyAtomGroup #134632
• [KeyInstr][Clang] Add Clang option -g[no-]key-instructions #134627
• [KeyInstr][DwarfDebug] Add is_stmt emission support #133495
• [KeyInstr] Add MIR parser support #133494
• [KeyInstr] Remap cloned PHIs in BreakCriticalEdges #133493
• Propagate DebugLocs on phis in BreakCriticalEdges #133492
• [KeyInstr][LoopUnswitch] Remap cloned instructions' atoms #133491
• [KeyInstr][LoopRotate] Remap atoms of duplicated instructions #133490
• [KeyInstr][LoopUnroll] Remap atoms while unrolling #133489
• [KeyInstr][JumpThreading] Remap atoms duping bb with cond br on phi into pred #133488
• [KeyInstr][JumpThreading] Remap atoms after threading #133487
• [KeyInstr][JumpThreading] Remap atoms in blocks duplicated for threading #133486
• [KeyInstr][Inline] Don't propagate atoms to inlined nodebug instructions #133485
• [KeyInstr][SimplifyCFG] Remap atoms after duplication for threading #133484
• [KeyInstr][debugify] Add --debugify-atoms to add key instructions metadata #133483
• [KeyInstr][SimplifyCFG] Remap atoms when folding br to common succ into pred #133482
• [KeyInstr] Inline atom info #133481
• [KeyInstr] Merge atoms in DILocation::getMergedLocation #133480
• [NFC][KeyInstr] Add Atom Group (re)mapping #133479
• [KeyInstr] Add Atom Group waterline to LLVMContext #133478
• [KeyInstr] Add fields to DILocation behind compile time flag #133477
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-clang-codegen

Author: Orlando Cazalet-Hyams (OCHyams)

Changes

[KeyInstr][Clang] Ret atom

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

The Key Instructions project is introduced, including a "quick summary" section
at the top which adds context for this PR, here:
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.

The Clang-side work is demoed here:
#130943

[KeyInstr][Clang] Update tests with ret atoms

---

Patch is 21.44 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/134652.diff

21 Files Affected:

• (modified) clang/lib/CodeGen/CGCall.cpp (+5-1)
• (modified) clang/lib/CodeGen/CGCleanup.cpp (+2)
• (modified) clang/lib/CodeGen/CGStmt.cpp (+9-4)
• (modified) clang/lib/CodeGen/CodeGenFunction.cpp (+16)
• (modified) clang/test/KeyInstructions/agg.c (+3)
• (modified) clang/test/KeyInstructions/assign-scalar.c (+3)
• (modified) clang/test/KeyInstructions/bitfield.cpp (+3)
• (modified) clang/test/KeyInstructions/builtin.c (+3)
• (modified) clang/test/KeyInstructions/complex.c (+3)
• (modified) clang/test/KeyInstructions/do.c (+3)
• (modified) clang/test/KeyInstructions/for.c (+3)
• (modified) clang/test/KeyInstructions/if.c (+4-1)
• (modified) clang/test/KeyInstructions/init-agg.c (+3)
• (modified) clang/test/KeyInstructions/init-member.cpp (+2)
• (modified) clang/test/KeyInstructions/init-scalar.c (+2-2)
• (modified) clang/test/KeyInstructions/init-static.cpp (+2-1)
• (added) clang/test/KeyInstructions/return-va-arg.c (+25)
• (added) clang/test/KeyInstructions/return.c (+90)
• (modified) clang/test/KeyInstructions/switch.c (+3)
• (modified) clang/test/KeyInstructions/try-catch.cpp (+3)
• (modified) clang/test/KeyInstructions/while.c (+3)

diff --git a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp
index 7aa77e55dbfcc..dba3fadba4f60 100644
--- a/clang/lib/CodeGen/CGCall.cpp
+++ b/clang/lib/CodeGen/CGCall.cpp
@@ -3883,7 +3883,8 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   // Functions with no result always return void.
   if (!ReturnValue.isValid()) {
-    Builder.CreateRetVoid();
+    auto *I = Builder.CreateRetVoid();
+    addRetToOverrideOrNewSourceAtom(I, nullptr);
     return;
   }
 
@@ -4065,6 +4066,9 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   if (RetDbgLoc)
     Ret->setDebugLoc(std::move(RetDbgLoc));
+
+  llvm::Value *Backup = RV ? Ret->getOperand(0) : nullptr;
+  addRetToOverrideOrNewSourceAtom(cast<llvm::ReturnInst>(Ret), Backup);
 }
 
 void CodeGenFunction::EmitReturnValueCheck(llvm::Value *RV) {
diff --git a/clang/lib/CodeGen/CGCleanup.cpp b/clang/lib/CodeGen/CGCleanup.cpp
index 7e1c5b7da9552..7292dcd47172c 100644
--- a/clang/lib/CodeGen/CGCleanup.cpp
+++ b/clang/lib/CodeGen/CGCleanup.cpp
@@ -1118,6 +1118,8 @@ void CodeGenFunction::EmitBranchThroughCleanup(JumpDest Dest) {
 
   // Create the branch.
   llvm::BranchInst *BI = Builder.CreateBr(Dest.getBlock());
+  // This is the primary instruction for this atom, acting in place of a ret.
+  addInstToCurrentSourceAtom(BI, nullptr);
 
   // Calculate the innermost active normal cleanup.
   EHScopeStack::stable_iterator
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
index 9292be24fc12e..21c2dd14799dd 100644
--- a/clang/lib/CodeGen/CGStmt.cpp
+++ b/clang/lib/CodeGen/CGStmt.cpp
@@ -1594,6 +1594,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 =
@@ -1603,6 +1604,7 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
     CGM.getSanitizerMetadata()->disableSanitizerForGlobal(SLocPtr);
     assert(ReturnLocation.isValid() && "No valid return location");
     Builder.CreateStore(SLocPtr, ReturnLocation);
+    //*OCH?*//
   }
 
   // Returning from an outlined SEH helper is UB, and we already warn on it.
@@ -1669,16 +1671,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:
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index b176227657f24..a4d2a48d77a17 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -36,6 +36,7 @@
 #include "clang/CodeGen/CGFunctionInfo.h"
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
@@ -339,6 +340,15 @@ llvm::DebugLoc CodeGenFunction::EmitReturnBlock() {
       // later by the actual 'ret' instruction.
       llvm::DebugLoc Loc = BI->getDebugLoc();
       Builder.SetInsertPoint(BI->getParent());
+
+      // Key Instructions: If there's only one `ret` then we want to put the
+      // instruction in the same source atom group as the store to the ret-value
+      // alloca and unconditional `br` to the return block that we're about to
+      // delete. It all comes from the same source (`return (value)`).
+      if (auto *DI = getDebugInfo(); DI && BI->getDebugLoc())
+        DI->setRetInstSourceAtomOverride(
+            BI->getDebugLoc().get()->getAtomGroup());
+
       BI->eraseFromParent();
       delete ReturnBlock.getBlock();
       ReturnBlock = JumpDest();
@@ -1543,6 +1553,12 @@ void CodeGenFunction::GenerateCode(GlobalDecl GD, llvm::Function *Fn,
       Bypasses.Init(CGM, Body);
   }
 
+  // Finalize function debug info on exit.
+  auto Cleanup = llvm::make_scope_exit([this] {
+    if (CGDebugInfo *DI = getDebugInfo())
+      DI->completeFunction();
+  });
+
   // Emit the standard function prologue.
   StartFunction(GD, ResTy, Fn, FnInfo, Args, Loc, BodyRange.getBegin());
 
diff --git a/clang/test/KeyInstructions/agg.c b/clang/test/KeyInstructions/agg.c
index 6caf84e89537f..5772922787fd5 100644
--- a/clang/test/KeyInstructions/agg.c
+++ b/clang/test/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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/assign-scalar.c b/clang/test/KeyInstructions/assign-scalar.c
index 1f1fe8fda39e6..801f3c1391e02 100644
--- a/clang/test/KeyInstructions/assign-scalar.c
+++ b/clang/test/KeyInstructions/assign-scalar.c
@@ -34,6 +34,8 @@ void fun() {
 // CHECK: %dec = add i64 %3, -1, !dbg [[G6R2:!.*]]
 // CHECK: store i64 %dec, ptr @g{{.*}}, !dbg [[G6R1:!.*]]
     g--;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -46,3 +48,4 @@ void fun() {
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
 // CHECK: [[G6R2]] = !DILocation({{.*}}, atomGroup: 6, atomRank: 2)
 // CHECK: [[G6R1]] = !DILocation({{.*}}, atomGroup: 6, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/bitfield.cpp b/clang/test/KeyInstructions/bitfield.cpp
index 0586050ba8397..be470dd7ca029 100644
--- a/clang/test/KeyInstructions/bitfield.cpp
+++ b/clang/test/KeyInstructions/bitfield.cpp
@@ -7,7 +7,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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/builtin.c b/clang/test/KeyInstructions/builtin.c
index 5129a4ac2c482..dbf4e287d58f2 100644
--- a/clang/test/KeyInstructions/builtin.c
+++ b/clang/test/KeyInstructions/builtin.c
@@ -57,6 +57,8 @@ void fun() {
 
 // CHECK: call void @llvm.memset{{.*}}, !dbg [[G14R1:!.*]]
     __builtin___memset_chk(f4, 0, sizeof(float), -1);
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -75,3 +77,4 @@ void fun() {
 // CHECK: [[G12R1]] = !DILocation({{.*}}, atomGroup: 12, atomRank: 1)
 // CHECK: [[G13R1]] = !DILocation({{.*}}, atomGroup: 13, atomRank: 1)
 // CHECK: [[G14R1]] = !DILocation({{.*}}, atomGroup: 14, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/complex.c b/clang/test/KeyInstructions/complex.c
index b97314e815bdc..86dd098ed93fc 100644
--- a/clang/test/KeyInstructions/complex.c
+++ b/clang/test/KeyInstructions/complex.c
@@ -28,6 +28,8 @@ void test() {
 // CHECK: %add = fadd float %0, %1, !dbg [[G4R2:!.*]]
 // CHECK: store float %add, ptr getelementptr inbounds nuw ({ float, float }, ptr @ci, i32 0, i32 1){{.*}}, !dbg [[G4R1:!.*]]
   __imag ci = __imag ci + __imag ci;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -38,3 +40,4 @@ void test() {
 // 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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/do.c b/clang/test/KeyInstructions/do.c
index 4b7a38cb35985..9778394c5d005 100644
--- a/clang/test/KeyInstructions/do.c
+++ b/clang/test/KeyInstructions/do.c
@@ -25,9 +25,12 @@ void a(int A) {
 // CHECK: %tobool = icmp ne i32 %dec, 0, !dbg [[G2R1:!.*]]
 // CHECK: br i1 %tobool, label %do.body, label %do.end, !dbg [[G3R1:!.*]], !llvm.loop
     do { } while (--A);
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/for.c b/clang/test/KeyInstructions/for.c
index 3221ece69a717..1336a461eae2b 100644
--- a/clang/test/KeyInstructions/for.c
+++ b/clang/test/KeyInstructions/for.c
@@ -27,6 +27,8 @@ void a(int A) {
 // CHECK: %inc = add{{.*}}, !dbg [[G4R2:!.*]]
 // CHECK: store i32 %inc, ptr %i{{.*}}, !dbg [[G4R1:!.*]]
     for (int i = 0; i < A; ++i) { }
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -35,3 +37,4 @@ void a(int A) {
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
 // CHECK: [[G4R2]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 2)
 // CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/if.c b/clang/test/KeyInstructions/if.c
index ccc7eb9253198..6fff140db9d01 100644
--- a/clang/test/KeyInstructions/if.c
+++ b/clang/test/KeyInstructions/if.c
@@ -31,7 +31,9 @@ void a(int A) {
 // CHECK-CXX: br i1 %tobool4, label %if.then5, label %if.end6{{.*}}, !dbg [[G5R1:!.*]]
     if (int B = A; B)
         ;
-#endif 
+#endif
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -44,3 +46,4 @@ void a(int A) {
 // CHECK-CXX: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
 // CHECK-CXX: [[G5R2]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 2)
 // CHECK-CXX: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-agg.c b/clang/test/KeyInstructions/init-agg.c
index dc3ccaedc57b5..1fd0b4909f6ad 100644
--- a/clang/test/KeyInstructions/init-agg.c
+++ b/clang/test/KeyInstructions/init-agg.c
@@ -33,6 +33,8 @@ void a() {
 
 // CHECK: store i8 -86, ptr %uninit{{.*}}, !dbg [[G5R1:!.*]], !annotation
     char uninit; // -ftrivial-auto-var-init=pattern
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -41,3 +43,4 @@ void a() {
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
 // CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-member.cpp b/clang/test/KeyInstructions/init-member.cpp
index 60949bd8a604a..a830fd06acb35 100644
--- a/clang/test/KeyInstructions/init-member.cpp
+++ b/clang/test/KeyInstructions/init-member.cpp
@@ -17,6 +17,8 @@ void fun() {
 
 // CHECK: store i32 1, ptr %x{{.*}}, !dbg [[G1R1:!.*]]
 // CHECK: store float 5.000000e+00, ptr %y{{.*}}, !dbg [[G2R1:!.*]]
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-scalar.c b/clang/test/KeyInstructions/init-scalar.c
index c212c2a4fd623..439283b80e3da 100644
--- a/clang/test/KeyInstructions/init-scalar.c
+++ b/clang/test/KeyInstructions/init-scalar.c
@@ -10,10 +10,10 @@ void a() {
 // CHECK: %add = add {{.*}}, !dbg [[G2R2:!.*]]
 // CHECK: store i32 %add, ptr %B, align 4, !dbg [[G2R1:!.*]]
     int B = 2 * A + 1;
-// CHECK-TODO: ret{{.*}}, !dbg [[G3R1:!.*]]
+// CHECK: ret{{.*}}, !dbg [[G3R1:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
-// CHECK-TODO: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
diff --git a/clang/test/KeyInstructions/init-static.cpp b/clang/test/KeyInstructions/init-static.cpp
index 82e14b59df5e5..5f5ea75662038 100644
--- a/clang/test/KeyInstructions/init-static.cpp
+++ b/clang/test/KeyInstructions/init-static.cpp
@@ -5,8 +5,9 @@ void g(int *a) {
     // CHECK: %2 = load ptr, ptr %a.addr{{.*}}, !dbg [[G1R2:!.*]]
     // CHECK: store ptr %2, ptr @_ZZ1gPiE1b{{.*}}, !dbg [[G1R1:!.*]]
     static int &b = *a;
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
-
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/return-va-arg.c b/clang/test/KeyInstructions/return-va-arg.c
new file mode 100644
index 0000000000000..2864489afd738
--- /dev/null
+++ b/clang/test/KeyInstructions/return-va-arg.c
@@ -0,0 +1,25 @@
+// RUN: %clang -gkey-instructions -gno-column-info -x c++ %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --implicit-check-not atomGroup --implicit-check-not atomRank
+
+// RUN: %clang -gkey-instructions -gno-column-info -x c %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --implicit-check-not atomGroup --implicit-check-not atomRank
+
+typedef struct {
+  struct{} a;
+  double b;
+} s1;
+
+s1 f(int z, ...) {
+  __builtin_va_list list;
+  __builtin_va_start(list, z);
+// CHECK: vaarg.end:
+// CHECK-NEXT: %vaarg.addr = phi ptr
+// CHECK-NEXT: call void @llvm.memcpy{{.*}}, !dbg [[G1R1:!.*]]
+// CHECK-NEXT: {{.*}} = getelementptr{{.*}}
+// CHECK-NEXT: [[LOAD:%.*]] = load double{{.*}}, !dbg [[G1R2:!.*]]
+// CHECK-NEXT: ret double [[LOAD]], !dbg [[G1R1]]
+  return __builtin_va_arg(list, s1);
+}
+
+// CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
diff --git a/clang/test/KeyInstructions/return.c b/clang/test/KeyInstructions/return.c
new file mode 100644
index 0000000000000..24c58450d37ed
--- /dev/null
+++ b/clang/test/KeyInstructions/return.c
@@ -0,0 +1,90 @@
+// RUN: %clang -gkey-instructions -gno-column-info -x c++ %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --check-prefixes=CHECK,CHECK-CXX
+
+// RUN: %clang -gkey-instructions -gno-column-info -x c %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s
+
+// Check the stores to `retval` allocas and branches to `return` block are 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).
+
+// Also check that 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.
+
+#ifdef __cplusplus
+#define nomangle extern "C"
+#else
+#define nomangle
+#endif
+
+int g;
+nomangle float a() {
+// CHECK: float @a()
+  if (g)
+// CHECK: if.then:
+// CHECK-NEXT: %1 = load i32, ptr @g{{.*}}, !dbg [[G2R3:!.*]]
+// CHECK-NEXT: %conv = sitofp i32 %1 to float{{.*}}, !dbg [[G2R2:!.*]]
+// CHECK-NEXT: store float %conv, ptr %retval{{.*}}, !dbg [[G2R1:!.*]]
+// CHECK-NEXT: br label %return{{.*}}, !dbg [[G2R1]]
+    return g;
+// CHECK: if.end:
+// CHECK-NEXT: store float 1.000000e+00, ptr %retval{{.*}}, !dbg [[G3R1:!.*]]
+// CHECK-NEXT: br label %return, !dbg [[G3R1]]
+
+// CHECK: return:
+// CHECK-NEXT:  %2 = load float, ptr %retval{{.*}}, !dbg [[G4R2:!.*]]
+// CHECK-NEXT:  ret float %2{{.*}}, !dbg [[G4R1:!.*]]
+  return 1;
+}
+
+// CHECK: void @b()
+// CHECK: ret void{{.*}}, !dbg [[B_G1R1:!.*]]
+nomangle void b() { return; }
+
+// CHECK: i32 @c()
+// CHECK: %add = add{{.*}}, !dbg [[C_G1R2:!.*]]
+// CHECK: ret i32 %add{{.*}}, !dbg [[C_G1R1:!.*]]
+nomangle int c() { return g + 1; }
+
+// NOTE: (return) (g = 1) are two separate atoms.
+// CHECK: i32 @d()
+// CHECK: store{{.*}}, !dbg [[D_G2R1:!.*]]
+// CHECK: ret i32 1{{.*}}, !dbg [[D_G1R1:!.*]]
+nomangle int d() { return g = 1; }
+
+// The implicit return here get the line number of the closing brace; make it
+// key to match existing behaviour.
+// CHECK: ret void, !dbg [[E_G1R1:!.*]]
+nomangle void e() {}
+
+#ifdef __cplusplus
+// CHECK-CXX: ptr @_Z1fRi
+int &f(int &r) {
+// Include ctrl-flow to stop ret value store being elided.
+    if (r)
+// CHECK-CXX: if.then:
+// CHECK-CXX-NEXT: %2 = load ptr, ptr %r.addr{{.*}}, !dbg [[F_G2R2:!.*]]
+// CHECK-CXX-NEXT: store ptr %2, ptr %retval{{.*}}, !dbg [[F_G2R1:!.*]]
+// CHECK-CXX-NEXT: br label %return, !dbg [[F_G2R1:!.*]]
+    return r;
+  return g;
+}
+#endif
+
+// CHECK: [[G2R3]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 3)
+// CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
+// CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[G4R2]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 2)
+// CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[B_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[C_G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
+// CHECK: [[C_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[D_G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[D_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[E_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK-CXX: [[F_G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
+// CHECK-CXX: [[F_G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
diff --git a/clang/test/KeyInstructions/switch.c b/clang/test/KeyInstructions/switch.c
index cff6b834106e9..32485156a8be1 100644
--- a/clang/test/KeyInstructions/switch.c
+++ b/clang/test/KeyInstructions/switch.c
@@ -41,6 +41,8 @@ void a(int A, int B) {
     } break;
     default: break;
     }
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
@@ -49,3 +51,4 @@ void a(int A, int B) {
 // CHECK: [[G3R2]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 2)
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
 // CHECK-CXX: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]]...
[truncated]

[llvmbot]

@llvm/pr-subscribers-clang

Author: Orlando Cazalet-Hyams (OCHyams)

Changes

[KeyInstr][Clang] Ret atom

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

The Key Instructions project is introduced, including a "quick summary" section
at the top which adds context for this PR, here:
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.

The Clang-side work is demoed here:
#130943

[KeyInstr][Clang] Update tests with ret atoms

---

Patch is 21.44 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/134652.diff

21 Files Affected:

• (modified) clang/lib/CodeGen/CGCall.cpp (+5-1)
• (modified) clang/lib/CodeGen/CGCleanup.cpp (+2)
• (modified) clang/lib/CodeGen/CGStmt.cpp (+9-4)
• (modified) clang/lib/CodeGen/CodeGenFunction.cpp (+16)
• (modified) clang/test/KeyInstructions/agg.c (+3)
• (modified) clang/test/KeyInstructions/assign-scalar.c (+3)
• (modified) clang/test/KeyInstructions/bitfield.cpp (+3)
• (modified) clang/test/KeyInstructions/builtin.c (+3)
• (modified) clang/test/KeyInstructions/complex.c (+3)
• (modified) clang/test/KeyInstructions/do.c (+3)
• (modified) clang/test/KeyInstructions/for.c (+3)
• (modified) clang/test/KeyInstructions/if.c (+4-1)
• (modified) clang/test/KeyInstructions/init-agg.c (+3)
• (modified) clang/test/KeyInstructions/init-member.cpp (+2)
• (modified) clang/test/KeyInstructions/init-scalar.c (+2-2)
• (modified) clang/test/KeyInstructions/init-static.cpp (+2-1)
• (added) clang/test/KeyInstructions/return-va-arg.c (+25)
• (added) clang/test/KeyInstructions/return.c (+90)
• (modified) clang/test/KeyInstructions/switch.c (+3)
• (modified) clang/test/KeyInstructions/try-catch.cpp (+3)
• (modified) clang/test/KeyInstructions/while.c (+3)

diff --git a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp
index 7aa77e55dbfcc..dba3fadba4f60 100644
--- a/clang/lib/CodeGen/CGCall.cpp
+++ b/clang/lib/CodeGen/CGCall.cpp
@@ -3883,7 +3883,8 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   // Functions with no result always return void.
   if (!ReturnValue.isValid()) {
-    Builder.CreateRetVoid();
+    auto *I = Builder.CreateRetVoid();
+    addRetToOverrideOrNewSourceAtom(I, nullptr);
     return;
   }
 
@@ -4065,6 +4066,9 @@ void CodeGenFunction::EmitFunctionEpilog(const CGFunctionInfo &FI,
 
   if (RetDbgLoc)
     Ret->setDebugLoc(std::move(RetDbgLoc));
+
+  llvm::Value *Backup = RV ? Ret->getOperand(0) : nullptr;
+  addRetToOverrideOrNewSourceAtom(cast<llvm::ReturnInst>(Ret), Backup);
 }
 
 void CodeGenFunction::EmitReturnValueCheck(llvm::Value *RV) {
diff --git a/clang/lib/CodeGen/CGCleanup.cpp b/clang/lib/CodeGen/CGCleanup.cpp
index 7e1c5b7da9552..7292dcd47172c 100644
--- a/clang/lib/CodeGen/CGCleanup.cpp
+++ b/clang/lib/CodeGen/CGCleanup.cpp
@@ -1118,6 +1118,8 @@ void CodeGenFunction::EmitBranchThroughCleanup(JumpDest Dest) {
 
   // Create the branch.
   llvm::BranchInst *BI = Builder.CreateBr(Dest.getBlock());
+  // This is the primary instruction for this atom, acting in place of a ret.
+  addInstToCurrentSourceAtom(BI, nullptr);
 
   // Calculate the innermost active normal cleanup.
   EHScopeStack::stable_iterator
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
index 9292be24fc12e..21c2dd14799dd 100644
--- a/clang/lib/CodeGen/CGStmt.cpp
+++ b/clang/lib/CodeGen/CGStmt.cpp
@@ -1594,6 +1594,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 =
@@ -1603,6 +1604,7 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
     CGM.getSanitizerMetadata()->disableSanitizerForGlobal(SLocPtr);
     assert(ReturnLocation.isValid() && "No valid return location");
     Builder.CreateStore(SLocPtr, ReturnLocation);
+    //*OCH?*//
   }
 
   // Returning from an outlined SEH helper is UB, and we already warn on it.
@@ -1669,16 +1671,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:
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
index b176227657f24..a4d2a48d77a17 100644
--- a/clang/lib/CodeGen/CodeGenFunction.cpp
+++ b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -36,6 +36,7 @@
 #include "clang/CodeGen/CGFunctionInfo.h"
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
@@ -339,6 +340,15 @@ llvm::DebugLoc CodeGenFunction::EmitReturnBlock() {
       // later by the actual 'ret' instruction.
       llvm::DebugLoc Loc = BI->getDebugLoc();
       Builder.SetInsertPoint(BI->getParent());
+
+      // Key Instructions: If there's only one `ret` then we want to put the
+      // instruction in the same source atom group as the store to the ret-value
+      // alloca and unconditional `br` to the return block that we're about to
+      // delete. It all comes from the same source (`return (value)`).
+      if (auto *DI = getDebugInfo(); DI && BI->getDebugLoc())
+        DI->setRetInstSourceAtomOverride(
+            BI->getDebugLoc().get()->getAtomGroup());
+
       BI->eraseFromParent();
       delete ReturnBlock.getBlock();
       ReturnBlock = JumpDest();
@@ -1543,6 +1553,12 @@ void CodeGenFunction::GenerateCode(GlobalDecl GD, llvm::Function *Fn,
       Bypasses.Init(CGM, Body);
   }
 
+  // Finalize function debug info on exit.
+  auto Cleanup = llvm::make_scope_exit([this] {
+    if (CGDebugInfo *DI = getDebugInfo())
+      DI->completeFunction();
+  });
+
   // Emit the standard function prologue.
   StartFunction(GD, ResTy, Fn, FnInfo, Args, Loc, BodyRange.getBegin());
 
diff --git a/clang/test/KeyInstructions/agg.c b/clang/test/KeyInstructions/agg.c
index 6caf84e89537f..5772922787fd5 100644
--- a/clang/test/KeyInstructions/agg.c
+++ b/clang/test/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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/assign-scalar.c b/clang/test/KeyInstructions/assign-scalar.c
index 1f1fe8fda39e6..801f3c1391e02 100644
--- a/clang/test/KeyInstructions/assign-scalar.c
+++ b/clang/test/KeyInstructions/assign-scalar.c
@@ -34,6 +34,8 @@ void fun() {
 // CHECK: %dec = add i64 %3, -1, !dbg [[G6R2:!.*]]
 // CHECK: store i64 %dec, ptr @g{{.*}}, !dbg [[G6R1:!.*]]
     g--;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -46,3 +48,4 @@ void fun() {
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
 // CHECK: [[G6R2]] = !DILocation({{.*}}, atomGroup: 6, atomRank: 2)
 // CHECK: [[G6R1]] = !DILocation({{.*}}, atomGroup: 6, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/bitfield.cpp b/clang/test/KeyInstructions/bitfield.cpp
index 0586050ba8397..be470dd7ca029 100644
--- a/clang/test/KeyInstructions/bitfield.cpp
+++ b/clang/test/KeyInstructions/bitfield.cpp
@@ -7,7 +7,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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/builtin.c b/clang/test/KeyInstructions/builtin.c
index 5129a4ac2c482..dbf4e287d58f2 100644
--- a/clang/test/KeyInstructions/builtin.c
+++ b/clang/test/KeyInstructions/builtin.c
@@ -57,6 +57,8 @@ void fun() {
 
 // CHECK: call void @llvm.memset{{.*}}, !dbg [[G14R1:!.*]]
     __builtin___memset_chk(f4, 0, sizeof(float), -1);
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -75,3 +77,4 @@ void fun() {
 // CHECK: [[G12R1]] = !DILocation({{.*}}, atomGroup: 12, atomRank: 1)
 // CHECK: [[G13R1]] = !DILocation({{.*}}, atomGroup: 13, atomRank: 1)
 // CHECK: [[G14R1]] = !DILocation({{.*}}, atomGroup: 14, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/complex.c b/clang/test/KeyInstructions/complex.c
index b97314e815bdc..86dd098ed93fc 100644
--- a/clang/test/KeyInstructions/complex.c
+++ b/clang/test/KeyInstructions/complex.c
@@ -28,6 +28,8 @@ void test() {
 // CHECK: %add = fadd float %0, %1, !dbg [[G4R2:!.*]]
 // CHECK: store float %add, ptr getelementptr inbounds nuw ({ float, float }, ptr @ci, i32 0, i32 1){{.*}}, !dbg [[G4R1:!.*]]
   __imag ci = __imag ci + __imag ci;
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -38,3 +40,4 @@ void test() {
 // 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: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/do.c b/clang/test/KeyInstructions/do.c
index 4b7a38cb35985..9778394c5d005 100644
--- a/clang/test/KeyInstructions/do.c
+++ b/clang/test/KeyInstructions/do.c
@@ -25,9 +25,12 @@ void a(int A) {
 // CHECK: %tobool = icmp ne i32 %dec, 0, !dbg [[G2R1:!.*]]
 // CHECK: br i1 %tobool, label %do.body, label %do.end, !dbg [[G3R1:!.*]], !llvm.loop
     do { } while (--A);
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/for.c b/clang/test/KeyInstructions/for.c
index 3221ece69a717..1336a461eae2b 100644
--- a/clang/test/KeyInstructions/for.c
+++ b/clang/test/KeyInstructions/for.c
@@ -27,6 +27,8 @@ void a(int A) {
 // CHECK: %inc = add{{.*}}, !dbg [[G4R2:!.*]]
 // CHECK: store i32 %inc, ptr %i{{.*}}, !dbg [[G4R1:!.*]]
     for (int i = 0; i < A; ++i) { }
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -35,3 +37,4 @@ void a(int A) {
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
 // CHECK: [[G4R2]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 2)
 // CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/if.c b/clang/test/KeyInstructions/if.c
index ccc7eb9253198..6fff140db9d01 100644
--- a/clang/test/KeyInstructions/if.c
+++ b/clang/test/KeyInstructions/if.c
@@ -31,7 +31,9 @@ void a(int A) {
 // CHECK-CXX: br i1 %tobool4, label %if.then5, label %if.end6{{.*}}, !dbg [[G5R1:!.*]]
     if (int B = A; B)
         ;
-#endif 
+#endif
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
@@ -44,3 +46,4 @@ void a(int A) {
 // CHECK-CXX: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
 // CHECK-CXX: [[G5R2]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 2)
 // CHECK-CXX: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-agg.c b/clang/test/KeyInstructions/init-agg.c
index dc3ccaedc57b5..1fd0b4909f6ad 100644
--- a/clang/test/KeyInstructions/init-agg.c
+++ b/clang/test/KeyInstructions/init-agg.c
@@ -33,6 +33,8 @@ void a() {
 
 // CHECK: store i8 -86, ptr %uninit{{.*}}, !dbg [[G5R1:!.*]], !annotation
     char uninit; // -ftrivial-auto-var-init=pattern
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
@@ -41,3 +43,4 @@ void a() {
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
 // CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
 // CHECK: [[G5R1]] = !DILocation({{.*}}, atomGroup: 5, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-member.cpp b/clang/test/KeyInstructions/init-member.cpp
index 60949bd8a604a..a830fd06acb35 100644
--- a/clang/test/KeyInstructions/init-member.cpp
+++ b/clang/test/KeyInstructions/init-member.cpp
@@ -17,6 +17,8 @@ void fun() {
 
 // CHECK: store i32 1, ptr %x{{.*}}, !dbg [[G1R1:!.*]]
 // CHECK: store float 5.000000e+00, ptr %y{{.*}}, !dbg [[G2R1:!.*]]
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/init-scalar.c b/clang/test/KeyInstructions/init-scalar.c
index c212c2a4fd623..439283b80e3da 100644
--- a/clang/test/KeyInstructions/init-scalar.c
+++ b/clang/test/KeyInstructions/init-scalar.c
@@ -10,10 +10,10 @@ void a() {
 // CHECK: %add = add {{.*}}, !dbg [[G2R2:!.*]]
 // CHECK: store i32 %add, ptr %B, align 4, !dbg [[G2R1:!.*]]
     int B = 2 * A + 1;
-// CHECK-TODO: ret{{.*}}, !dbg [[G3R1:!.*]]
+// CHECK: ret{{.*}}, !dbg [[G3R1:!.*]]
 }
 
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
 // CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
 // CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
-// CHECK-TODO: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
diff --git a/clang/test/KeyInstructions/init-static.cpp b/clang/test/KeyInstructions/init-static.cpp
index 82e14b59df5e5..5f5ea75662038 100644
--- a/clang/test/KeyInstructions/init-static.cpp
+++ b/clang/test/KeyInstructions/init-static.cpp
@@ -5,8 +5,9 @@ void g(int *a) {
     // CHECK: %2 = load ptr, ptr %a.addr{{.*}}, !dbg [[G1R2:!.*]]
     // CHECK: store ptr %2, ptr @_ZZ1gPiE1b{{.*}}, !dbg [[G1R1:!.*]]
     static int &b = *a;
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
 // CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
-
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]], atomRank: [[#]])
diff --git a/clang/test/KeyInstructions/return-va-arg.c b/clang/test/KeyInstructions/return-va-arg.c
new file mode 100644
index 0000000000000..2864489afd738
--- /dev/null
+++ b/clang/test/KeyInstructions/return-va-arg.c
@@ -0,0 +1,25 @@
+// RUN: %clang -gkey-instructions -gno-column-info -x c++ %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --implicit-check-not atomGroup --implicit-check-not atomRank
+
+// RUN: %clang -gkey-instructions -gno-column-info -x c %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --implicit-check-not atomGroup --implicit-check-not atomRank
+
+typedef struct {
+  struct{} a;
+  double b;
+} s1;
+
+s1 f(int z, ...) {
+  __builtin_va_list list;
+  __builtin_va_start(list, z);
+// CHECK: vaarg.end:
+// CHECK-NEXT: %vaarg.addr = phi ptr
+// CHECK-NEXT: call void @llvm.memcpy{{.*}}, !dbg [[G1R1:!.*]]
+// CHECK-NEXT: {{.*}} = getelementptr{{.*}}
+// CHECK-NEXT: [[LOAD:%.*]] = load double{{.*}}, !dbg [[G1R2:!.*]]
+// CHECK-NEXT: ret double [[LOAD]], !dbg [[G1R1]]
+  return __builtin_va_arg(list, s1);
+}
+
+// CHECK: [[G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
diff --git a/clang/test/KeyInstructions/return.c b/clang/test/KeyInstructions/return.c
new file mode 100644
index 0000000000000..24c58450d37ed
--- /dev/null
+++ b/clang/test/KeyInstructions/return.c
@@ -0,0 +1,90 @@
+// RUN: %clang -gkey-instructions -gno-column-info -x c++ %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s --check-prefixes=CHECK,CHECK-CXX
+
+// RUN: %clang -gkey-instructions -gno-column-info -x c %s -gmlt -S -emit-llvm -o - \
+// RUN: | FileCheck %s
+
+// Check the stores to `retval` allocas and branches to `return` block are 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).
+
+// Also check that 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.
+
+#ifdef __cplusplus
+#define nomangle extern "C"
+#else
+#define nomangle
+#endif
+
+int g;
+nomangle float a() {
+// CHECK: float @a()
+  if (g)
+// CHECK: if.then:
+// CHECK-NEXT: %1 = load i32, ptr @g{{.*}}, !dbg [[G2R3:!.*]]
+// CHECK-NEXT: %conv = sitofp i32 %1 to float{{.*}}, !dbg [[G2R2:!.*]]
+// CHECK-NEXT: store float %conv, ptr %retval{{.*}}, !dbg [[G2R1:!.*]]
+// CHECK-NEXT: br label %return{{.*}}, !dbg [[G2R1]]
+    return g;
+// CHECK: if.end:
+// CHECK-NEXT: store float 1.000000e+00, ptr %retval{{.*}}, !dbg [[G3R1:!.*]]
+// CHECK-NEXT: br label %return, !dbg [[G3R1]]
+
+// CHECK: return:
+// CHECK-NEXT:  %2 = load float, ptr %retval{{.*}}, !dbg [[G4R2:!.*]]
+// CHECK-NEXT:  ret float %2{{.*}}, !dbg [[G4R1:!.*]]
+  return 1;
+}
+
+// CHECK: void @b()
+// CHECK: ret void{{.*}}, !dbg [[B_G1R1:!.*]]
+nomangle void b() { return; }
+
+// CHECK: i32 @c()
+// CHECK: %add = add{{.*}}, !dbg [[C_G1R2:!.*]]
+// CHECK: ret i32 %add{{.*}}, !dbg [[C_G1R1:!.*]]
+nomangle int c() { return g + 1; }
+
+// NOTE: (return) (g = 1) are two separate atoms.
+// CHECK: i32 @d()
+// CHECK: store{{.*}}, !dbg [[D_G2R1:!.*]]
+// CHECK: ret i32 1{{.*}}, !dbg [[D_G1R1:!.*]]
+nomangle int d() { return g = 1; }
+
+// The implicit return here get the line number of the closing brace; make it
+// key to match existing behaviour.
+// CHECK: ret void, !dbg [[E_G1R1:!.*]]
+nomangle void e() {}
+
+#ifdef __cplusplus
+// CHECK-CXX: ptr @_Z1fRi
+int &f(int &r) {
+// Include ctrl-flow to stop ret value store being elided.
+    if (r)
+// CHECK-CXX: if.then:
+// CHECK-CXX-NEXT: %2 = load ptr, ptr %r.addr{{.*}}, !dbg [[F_G2R2:!.*]]
+// CHECK-CXX-NEXT: store ptr %2, ptr %retval{{.*}}, !dbg [[F_G2R1:!.*]]
+// CHECK-CXX-NEXT: br label %return, !dbg [[F_G2R1:!.*]]
+    return r;
+  return g;
+}
+#endif
+
+// CHECK: [[G2R3]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 3)
+// CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
+// CHECK: [[G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
+// CHECK: [[G4R2]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 2)
+// CHECK: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[B_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[C_G1R2]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 2)
+// CHECK: [[C_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[D_G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
+// CHECK: [[D_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK: [[E_G1R1]] = !DILocation({{.*}}, atomGroup: 1, atomRank: 1)
+// CHECK-CXX: [[F_G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
+// CHECK-CXX: [[F_G2R1]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 1)
diff --git a/clang/test/KeyInstructions/switch.c b/clang/test/KeyInstructions/switch.c
index cff6b834106e9..32485156a8be1 100644
--- a/clang/test/KeyInstructions/switch.c
+++ b/clang/test/KeyInstructions/switch.c
@@ -41,6 +41,8 @@ void a(int A, int B) {
     } break;
     default: break;
     }
+
+// CHECK: ret{{.*}}, !dbg [[RET:!.*]]
 }
 
 // CHECK: [[G2R2]] = !DILocation({{.*}}, atomGroup: 2, atomRank: 2)
@@ -49,3 +51,4 @@ void a(int A, int B) {
 // CHECK: [[G3R2]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 2)
 // CHECK: [[G3R1]] = !DILocation({{.*}}, atomGroup: 3, atomRank: 1)
 // CHECK-CXX: [[G4R1]] = !DILocation({{.*}}, atomGroup: 4, atomRank: 1)
+// CHECK: [[RET:!.*]] = !DILocation({{.*}}, atomGroup: [[#]]...
[truncated]