[llvm-branch-commits] [llvm] [KeyInstr][DwarfDebug] Add is_stmt emission support (PR #133495)

[Jeremy Morse via llvm-branch-commits]

================
@@ -2333,6 +2352,170 @@ DwarfDebug::emitInitialLocDirective(const MachineFunction &MF, unsigned CUID) {
   return PrologEndLoc;
 }
 
+void DwarfDebug::findKeyInstructions(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 KeyInstructions.
+  //       Remove existing instructions from KeyInstructions 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 KeyInstructions. Add this instr to 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 instruciton if it has a different line number.
+      if (!Buoy ||
+          Buoy->getDebugLoc().getLine() != MI.getDebugLoc().getLine()) {
+        Buoy = &MI;
+        BuoyAtom = 0;
+      }
+
+      // Call instructions are handled specially - we always mark them as key
+      // regardless of atom info.
+      const auto &TII =
+          *MI.getParent()->getParent()->getSubtarget().getInstrInfo();
+      if (MI.isCall() || TII.isTailCall(MI)) {
+        assert(MI.getDebugLoc() && "Unexpectedly missing DL");
+
+        // Calls are always key.
+        KeyInstructions.insert(Buoy);
+
+        uint64_t Group = MI.getDebugLoc()->getAtomGroup();
+        uint8_t Rank = MI.getDebugLoc()->getAtomRank();
+        if (Group && Rank) {
+          auto *InlinedAt = MI.getDebugLoc()->getInlinedAt();
+          auto &[CandidateRank, CandidateInsts] = GroupCandidates[{InlinedAt, Group}];
+
+          // This looks similar to the non-call handling code, except that
+          // we don't put the call into CandidateInsts so that they can't be
+          // made un-key. As a result, we also have to take special care not
+          // to erase the is_stmt from the buoy, and prevent that happening
+          // in the future.
+
+          if (CandidateRank == Rank) {
+            // We've seen other instructions in this group of this rank. Discard
+            // ones we've seen in this block, keep the others.
+            assert(!CandidateInsts.empty());
+            SmallVector<const MachineInstr *> Insts;
+            Insts.reserve(CandidateInsts.size());
+            for (auto &PrevInst : CandidateInsts) {
+              if (PrevInst->getParent() != MI.getParent())
+                Insts.push_back(PrevInst);
+              else if (PrevInst != Buoy)
+                KeyInstructions.erase(PrevInst);
+            }
+
+            if (Insts.empty()) {
+              CandidateInsts.clear();
+              CandidateRank = 0;
+            } else {
+              CandidateInsts = std::move(Insts);
+            }
+
+          } else if (CandidateRank > Rank) {
+            // We've seen other instructions in this group of lower precedence
+            // (higher rank). Discard them.
+            for (auto *Supplanted : CandidateInsts) {
+              // Don't erase the is_stmt we're using for this call.
+              if (Supplanted != Buoy)
+                KeyInstructions.erase(Supplanted);
+            }
+            CandidateInsts.clear();
+            CandidateRank = 0;
+          }
+        }
+
+        // Avoid floating any future is_stmts up to the call.
+        Buoy = nullptr;
+        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 = MI.getDebugLoc()->getAtomGroup();
----------------
jmorse wrote:

```suggestion
        BuoyAtom = Group;
```
Avoids the appearance that this could be different from the one earlier.

https://github.com/llvm/llvm-project/pull/133495