[KeyInstr][DwarfDebug] Add is_stmt emission support

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -17,6 +17,7 @@
 #include "DwarfExpression.h"
 #include "DwarfUnit.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/Twine.h"
@@ -170,6 +171,9 @@ static cl::opt<DwarfDebug::MinimizeAddrInV5> MinimizeAddrInV5Option(
                           "Stuff")),
     cl::init(DwarfDebug::MinimizeAddrInV5::Default));
 
+static cl::opt<bool> KeyInstructionsAreStmts("dwarf-use-key-instructions",
+                                             cl::Hidden, cl::init(false));
+
 static constexpr unsigned ULEB128PadSize = 4;
 
 void DebugLocDwarfExpression::emitOp(uint8_t Op, const char *Comment) {
@@ -2070,6 +2074,10 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
   unsigned LastAsmLine =
       Asm->OutStreamer->getContext().getCurrentDwarfLoc().getLine();
 
+  bool IsKey = false;
+  if (KeyInstructionsAreStmts && DL && DL.getLine())
+    IsKey = KeyInstructions.contains(MI);
+
   if (!DL && MI == PrologEndLoc) {
     // In rare situations, we might want to place the end of the prologue
     // somewhere that doesn't have a source location already. It should be in
@@ -2084,17 +2092,22 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
       (!PrevInstBB ||
        PrevInstBB->getSectionID() == MI->getParent()->getSectionID());
   bool ForceIsStmt = ForceIsStmtInstrs.contains(MI);
-  if (DL == PrevInstLoc && PrevInstInSameSection && !ForceIsStmt) {
+  if (PrevInstInSameSection && !ForceIsStmt && DL.isSameSourceLocation(PrevInstLoc)) {
     // If we have an ongoing unspecified location, nothing to do here.
     if (!DL)
       return;
-    // We have an explicit location, same as the previous location.
-    // But we might be coming back to it after a line 0 record.
-    if ((LastAsmLine == 0 && DL.getLine() != 0) || Flags) {
-      // Reinstate the source location but not marked as a statement.
-      RecordSourceLine(DL, Flags);
+
+    // Skip this if the instruction is Key, else we might accidentally miss an
+    // is_stmt.
+    if (!IsKey) {
+      // We have an explicit location, same as the previous location.
+      // But we might be coming back to it after a line 0 record.
+      if ((LastAsmLine == 0 && DL.getLine() != 0) || Flags) {
+        // Reinstate the source location but not marked as a statement.
+        RecordSourceLine(DL, Flags);
+      }
+      return;
     }
-    return;
   }
 
   if (!DL) {
@@ -2141,11 +2154,17 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
     Flags |= DWARF2_FLAG_PROLOGUE_END | DWARF2_FLAG_IS_STMT;
     PrologEndLoc = nullptr;
   }
-  // If the line changed, we call that a new statement; unless we went to
-  // line 0 and came back, in which case it is not a new statement.
-  unsigned OldLine = PrevInstLoc ? PrevInstLoc.getLine() : LastAsmLine;
-  if (DL.getLine() && (DL.getLine() != OldLine || ForceIsStmt))
-    Flags |= DWARF2_FLAG_IS_STMT;
+
+  if (KeyInstructionsAreStmts) {
+    if (IsKey)
+      Flags |= DWARF2_FLAG_IS_STMT;
+  } else {
+    // If the line changed, we call that a new statement; unless we went to
+    // line 0 and came back, in which case it is not a new statement.
+    unsigned OldLine = PrevInstLoc ? PrevInstLoc.getLine() : LastAsmLine;
+    if (DL.getLine() && (DL.getLine() != OldLine || ForceIsStmt))
+      Flags |= DWARF2_FLAG_IS_STMT;
+  }
 
   RecordSourceLine(DL, Flags);
 
@@ -2338,6 +2357,137 @@ DwarfDebug::emitInitialLocDirective(const MachineFunction &MF, unsigned CUID) {
   return PrologEndLoc;
 }
 
+void DwarfDebug::computeKeyInstructions(const MachineFunction *MF) {
+  // New function - reset KeyInstructions.
+  KeyInstructions.clear();
+
+  // The current candidate is_stmt instructions for each source atom.
+  // Map {(InlinedAt, Group): (Rank, Instructions)}.
+  DenseMap<std::pair<DILocation *, uint32_t>,
+           std::pair<uint16_t, SmallVector<const MachineInstr *>>>
+      GroupCandidates;
+
+  // For each instruction:
+  //   * Skip insts without DebugLoc, AtomGroup or AtomRank, and line zeros.
+  //   * Check if insts in this group have been seen already in GroupCandidates.
+  //     * If this instr rank is equal, add this instruction to GroupCandidates.
+  //       Remove existing instructions from GroupCandidates if they have the
+  //       same parent.
+  //     * If this instr rank is higher (lower precedence), ignore it.
+  //     * If this instr rank is lower (higher precedence), erase existing
+  //       instructions from GroupCandidates and add this one.
+  //
+  // Then insert each GroupCandidates instruction into KeyInstructions.
+
+  for (auto &MBB : *MF) {
+    // Rather than apply is_stmt directly to Key Instructions, we "float"
+    // is_stmt up to the 1st instruction with the same line number in a
+    // contiguous block. That instruction is called the "buoy". The
+    // buoy gets reset if we encouner an instruction with an atom
+    // group.
+    const MachineInstr *Buoy = nullptr;
+    // The atom group number associated with Buoy which may be 0 if we haven't
+    // encountered an atom group yet in this blob of instructions with the same
+    // line number.
+    uint64_t BuoyAtom = 0;
+
+    for (auto &MI : MBB) {
+      if (MI.isMetaInstruction())
+        continue;
+
+      if (!MI.getDebugLoc() || !MI.getDebugLoc().getLine())
+        continue;
+
+      // Reset the Buoy to this instruction if it has a different line number.
+      if (!Buoy ||
+          Buoy->getDebugLoc().getLine() != MI.getDebugLoc().getLine()) {
+        Buoy = &MI;
+        BuoyAtom = 0; // Set later when we know which atom the buoy is used by.
+      }
+
+      // Call instructions are handled specially - we always mark them as key
+      // regardless of atom info.
+      const auto &TII =
+          *MI.getParent()->getParent()->getSubtarget().getInstrInfo();
+      bool IsCallLike = MI.isCall() || TII.isTailCall(MI);
+      if (IsCallLike) {
+        assert(MI.getDebugLoc() && "Unexpectedly missing DL");
+
+        // Calls are always key. Put the buoy (may not be the call) into
+        // KeyInstructions directly rather than the candidate map to avoid it
+        // being erased (and we may not have a group number for the call).
+        KeyInstructions.insert(Buoy);
+
+        // Avoid floating any future is_stmts up to the call.
+        Buoy = nullptr;
+        BuoyAtom = 0;
+
+        if (!MI.getDebugLoc()->getAtomGroup() ||
+            !MI.getDebugLoc()->getAtomRank())
+          continue;
+      }
+
+      auto *InlinedAt = MI.getDebugLoc()->getInlinedAt();
+      uint64_t Group = MI.getDebugLoc()->getAtomGroup();
+      uint8_t Rank = MI.getDebugLoc()->getAtomRank();
+      if (!Group || !Rank)
+        continue;
+
+      // Don't let is_stmts float past instructions from different source atoms.
+      if (BuoyAtom && BuoyAtom != Group) {
+        Buoy = &MI;
+        BuoyAtom = Group;
+      }
+
+      auto &[CandidateRank, CandidateInsts] =
+          GroupCandidates[{InlinedAt, Group}];
+
+      // If CandidateRank is zero then CandidateInsts should be empty: there
+      // are no other candidates for this group yet. If CandidateRank is nonzero
+      // then CandidateInsts shouldn't be empty: we've got existing candidate
+      // instructions.
+      assert((CandidateRank == 0 && CandidateInsts.empty()) ||
+             (CandidateRank != 0 && !CandidateInsts.empty()));
+
+      assert(Rank && "expected nonzero rank");
+      // If we've seen other instructions in this group with higher precedence
+      // (lower nonzero rank), don't add this one as a candidate.
+      if (CandidateRank && CandidateRank < Rank)
+        continue;
+
+      // If we've seen other instructions in this group of the same rank,
+      // discard any from this block (keeping the others). Else if we've
+      // seen other instructions in this group of lower precedence (higher
+      // rank), discard them all.
+      if (CandidateRank == Rank)
+        llvm::remove_if(CandidateInsts, [&MI](const MachineInstr *Candidate) {
+          return MI.getParent() == Candidate->getParent();
+        });
+      else if (CandidateRank > Rank)
+        CandidateInsts.clear();
+
+      if (Buoy) {
+        // Add this candidate.
+        CandidateInsts.push_back(Buoy);
+        CandidateRank = Rank;
+
+        assert(!BuoyAtom || BuoyAtom == MI.getDebugLoc()->getAtomGroup());
+        BuoyAtom = MI.getDebugLoc()->getAtomGroup();
+      } else {
+        // Don't add calls, because they've been dealt with already. This means
+        // CandidateInsts might now be empty - handle that.
+        assert(IsCallLike);
+        if (CandidateInsts.empty())
+          CandidateRank = 0;
+      }
+    }
+  }
+
+  for (const auto &[_, Insts] : GroupCandidates.values())
+    for (auto *I : Insts)
+      KeyInstructions.insert(I);
+}
+
 /// For the function \p MF, finds the set of instructions which may represent a
 /// change in line number from one or more of the preceding MBBs. Stores the
 /// resulting set of instructions, which should have is_stmt set, in
@@ -2496,7 +2646,10 @@ void DwarfDebug::beginFunctionImpl(const MachineFunction *MF) {
   PrologEndLoc = emitInitialLocDirective(
       *MF, Asm->OutStreamer->getContext().getDwarfCompileUnitID());
 
-  findForceIsStmtInstrs(MF);
+  if (KeyInstructionsAreStmts)
+    computeKeyInstructions(MF);
+  else
+    findForceIsStmtInstrs(MF);
 }
 
 unsigned

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.h

@@ -464,6 +464,10 @@ class DwarfDebug : public DebugHandlerBase {
   };
 
 private:
+  /// Instructions which should get is_stmt applied because they implement key
+  /// functionality for a source atom.
+  SmallDenseSet<const MachineInstr *> KeyInstructions;
+
   /// Force the use of DW_AT_ranges even for single-entry range lists.
   MinimizeAddrInV5 MinimizeAddr = MinimizeAddrInV5::Disabled;
 
@@ -701,6 +705,11 @@ class DwarfDebug : public DebugHandlerBase {
 
   void findForceIsStmtInstrs(const MachineFunction *MF);
 
+  /// Compute instructions which should get is_stmt applied because they
+  /// implement key functionality for a source location atom, store results in
+  /// DwarfDebug::KeyInstructions.
+  void computeKeyInstructions(const MachineFunction *MF);
+
 protected:
   /// Gather pre-function debug information.
   void beginFunctionImpl(const MachineFunction *MF) override;

llvm/test/DebugInfo/KeyInstructions/X86/dwarf-basic-ranks.ll

@@ -0,0 +1,69 @@
+; RUN: llc %s --filetype=obj -o - --dwarf-use-key-instructions \
+; RUN: | llvm-objdump -d - --no-show-raw-insn \
+; RUN: | FileCheck %s --check-prefix=OBJ
+
+; RUN: llc %s --filetype=obj -o - --dwarf-use-key-instructions \
+; RUN: | llvm-dwarfdump - --debug-line \
+; RUN: | FileCheck %s --check-prefix=DBG
+
+;; 1. [[gnu::nodebug]] void prologue_end();
+;; 2.
+;; 3. int f(int *a, int b, int c) {
+;; 4.   prologue_end();
+;; 5.   *a =
+;; 6.     b + c;
+;; 7.   return *a;
+;; 8. }
+;;
+;; The add and store are in the same group (1). The add (line 6) has lower
+;; precedence (rank 2) so should not get is_stmt applied.
+
+; OBJ: 0000000000000000 <_Z1fPiii>:
+; OBJ-NEXT:  0: pushq   %rbp
+; OBJ-NEXT:  1: pushq   %r14
+; OBJ-NEXT:  3: pushq   %rbx
+; OBJ-NEXT:  4: movl    %edx, %ebx
+; OBJ-NEXT:  6: movl    %esi, %ebp
+; OBJ-NEXT:  8: movq    %rdi, %r14
+; OBJ-NEXT:  b: callq   0x10 <_Z1fPiii+0x10>
+; OBJ-NEXT: 10: addl    %ebx, %ebp
+; OBJ-NEXT: 12: movl    %ebp, (%r14)
+; OBJ-NEXT: 15: movl    %ebp, %eax
+; OBJ-NEXT: 17: popq    %rbx
+; OBJ-NEXT: 18: popq    %r14
+; OBJ-NEXT: 1a: popq    %rbp
+
+; DBG:      Address            Line   Column File   ISA Discriminator OpIndex Flags
+; DBG-NEXT: ------------------ ------ ------ ------ --- ------------- ------- -------------
+; DBG-NEXT: 0x0000000000000000      3      0      0   0             0       0  is_stmt
+; DBG-NEXT: 0x000000000000000b      4      0      0   0             0       0  is_stmt prologue_end
+; DBG-NEXT: 0x0000000000000010      6      0      0   0             0       0
+; DBG-NEXT: 0x0000000000000012      5      0      0   0             0       0  is_stmt
+; DBG-NEXT: 0x0000000000000015      7      0      0   0             0       0  is_stmt
+; DBG-NEXT: 0x0000000000000017      7      0      0   0             0       0  epilogue_begin
+; DBG-NEXT: 0x000000000000001c      7      0      0   0             0       0  end_sequence
+
+target triple = "x86_64-unknown-linux-gnu"
+
+define hidden noundef i32 @_Z1fPiii(ptr %a, i32 %b, i32 %c) local_unnamed_addr !dbg !11 {
+entry:
+  tail call void @_Z12prologue_endv(), !dbg !DILocation(line: 4, scope: !11)
+  %add = add nsw i32 %c, %b,           !dbg !DILocation(line: 6, scope: !11, atomGroup: 1, atomRank: 2)
+  store i32 %add, ptr %a, align 4,     !dbg !DILocation(line: 5, scope: !11, atomGroup: 1, atomRank: 1)
+  ret i32 %add,                        !dbg !DILocation(line: 7, scope: !11, atomGroup: 2, atomRank: 1)
+}
+
+declare void @_Z12prologue_endv() local_unnamed_addr #1
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3}
+!llvm.ident = !{!10}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_17, file: !1, producer: "clang version 19.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "test.cpp", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!10 = !{!"clang version 19.0.0"}
+!11 = distinct !DISubprogram(name: "f", scope: !1, file: !1, line: 3, type: !12, scopeLine: 3, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+!12 = !DISubroutineType(types: !13)
+!13 = !{}

llvm/test/DebugInfo/KeyInstructions/X86/dwarf-basic.ll

@@ -0,0 +1,62 @@
+; RUN: llc %s --filetype=obj -o - --dwarf-use-key-instructions \
+; RUN: | llvm-objdump -d - --no-show-raw-insn \
+; RUN: | FileCheck %s --check-prefix=OBJ
+
+; RUN: llc %s --filetype=obj -o - --dwarf-use-key-instructions \
+; RUN: | llvm-dwarfdump - --debug-line \
+; RUN: | FileCheck %s --check-prefix=DBG
+
+;; 1. int f(int a) {
+;; 2.   int x = a + 1;
+;; 3.   return x;
+;; 4. }
+;; 5. int g(int b) {
+;; 6.   return f(b);
+;; 7. }
+;;
+;; Both functions contain 2 instructions in unique atom groups. In f we see
+;; groups 1 and 3, and in g we see {!18, 1} and 1. All of these instructions
+;; should get is_stmt.
+
+; OBJ: 0000000000000000 <_Z1fi>:
+; OBJ-NEXT: 0: leal    0x1(%rdi), %eax
+; OBJ-NEXT: 3: retq
+; OBJ: 0000000000000010 <_Z1gi>:
+; OBJ-NEXT: 10: leal    0x1(%rdi), %eax
+; OBJ-NEXT: 13: retq
+
+; DBG:      Address            Line   Column File   ISA Discriminator OpIndex Flags
+; DBG-NEXT: ------------------ ------ ------ ------ --- ------------- ------- -------------
+; DBG-NEXT: 0x0000000000000000      2      0      0   0             0       0  is_stmt prologue_end
+; DBG-NEXT: 0x0000000000000003      3      0      0   0             0       0  is_stmt
+; DBG-NEXT: 0x0000000000000010      2      0      0   0             0       0  is_stmt prologue_end
+; DBG-NEXT: 0x0000000000000013      6      0      0   0             0       0  is_stmt
+
+target triple = "x86_64-unknown-linux-gnu"
+
+define hidden noundef i32 @_Z1fi(i32 noundef %a) local_unnamed_addr !dbg !11 {
+entry:
+  %add = add nsw i32 %a, 1,   !dbg !DILocation(line: 2, scope: !11, atomGroup: 1, atomRank: 2)
+  ret i32 %add,               !dbg !DILocation(line: 3, scope: !11, atomGroup: 3, atomRank: 1)
+}
+
+define hidden noundef i32 @_Z1gi(i32 noundef %b) local_unnamed_addr !dbg !16 {
+entry:
+  %add.i = add nsw i32 %b, 1, !dbg !DILocation(line: 2, scope: !11, inlinedAt: !18, atomGroup: 1, atomRank: 2)
+  ret i32 %add.i,             !dbg !DILocation(line: 6, scope: !16, atomGroup: 1, atomRank: 1)
+}
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3}
+!llvm.ident = !{!10}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_17, file: !1, producer: "clang version 19.0.0", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "test.cpp", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!10 = !{!"clang version 19.0.0"}
+!11 = distinct !DISubprogram(name: "f", scope: !1, file: !1, line: 1, type: !12, scopeLine: 1, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+!12 = !DISubroutineType(types: !13)
+!13 = !{}
+!16 = distinct !DISubprogram(name: "g", scope: !1, file: !1, line: 5, type: !12, scopeLine: 5, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+!18 = distinct !DILocation(line: 6, scope: !16)