[KeyInstr] Add Atom Group waterline to LLVMContext (#133478)

Source location atoms are identified by a function-local number and the
DILocation's InlinedAt field. The front end is responsible for assigning source
atom numbers, but certain optimisations need to assign new atom numbers to some
instructions. Most often code duplication optimisations like loop
unroll. Tracking a global maximum value (waterline) means we can easily
(cheaply) get new numbers that don't clash in any function.

The waterline is managed through DILocation creation,
LLVMContext::incNextAtomGroup, and LLVMContext::updateAtomGroupWaterline.

Add unittest.

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

Author: OCHyams

llvm/include/llvm/IR/LLVMContext.h

@@ -335,6 +335,14 @@ class LLVMContext {
   StringRef getDefaultTargetFeatures();
   void setDefaultTargetFeatures(StringRef Features);
 
+  /// Key Instructions: update the highest number atom group emitted for any
+  /// function.
+  void updateDILocationAtomGroupWaterline(uint64_t G);
+
+  /// Key Instructions: get the next free atom group number and increment
+  /// the global tracker.
+  uint64_t incNextDILocationAtomGroup();
+
 private:
   // Module needs access to the add/removeModule methods.
   friend class Module;

llvm/lib/IR/DebugInfoMetadata.cpp

@@ -74,6 +74,9 @@ DILocation::DILocation(LLVMContext &C, StorageType Storage, unsigned Line,
 #ifdef EXPERIMENTAL_KEY_INSTRUCTIONS
   assert(AtomRank <= 7 && "AtomRank number should fit in 3 bits");
 #endif
+  if (AtomGroup)
+    C.updateDILocationAtomGroupWaterline(AtomGroup + 1);
+
   assert((MDs.size() == 1 || MDs.size() == 2) &&
          "Expected a scope and optional inlined-at");
   // Set line and column.

llvm/lib/IR/LLVMContext.cpp

@@ -377,3 +377,11 @@ StringRef LLVMContext::getDefaultTargetFeatures() {
 void LLVMContext::setDefaultTargetFeatures(StringRef Features) {
   pImpl->DefaultTargetFeatures = Features;
 }
+
+void LLVMContext::updateDILocationAtomGroupWaterline(uint64_t V) {
+  pImpl->NextAtomGroup = std::max(pImpl->NextAtomGroup, V);
+}
+
+uint64_t LLVMContext::incNextDILocationAtomGroup() {
+  return pImpl->NextAtomGroup++;
+}

llvm/lib/IR/LLVMContextImpl.h

@@ -1880,6 +1880,16 @@ class LLVMContextImpl {
 
   std::string DefaultTargetCPU;
   std::string DefaultTargetFeatures;
+
+  /// The next available source atom group number. The front end is responsible
+  /// for assigning source atom numbers, but certain optimisations need to
+  /// assign new group numbers to a set of instructions. Most often code
+  /// duplication optimisations like loop unroll. Tracking a global maximum
+  /// value means we can know (cheaply) we're never using a group number that's
+  /// already used within this function.
+  ///
+  /// Start a 1 because 0 means the source location isn't part of an atom group.
+  uint64_t NextAtomGroup = 1;
 };
 
 } // end namespace llvm

llvm/unittests/IR/MetadataTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Metadata.h"
+#include "../lib/IR/LLVMContextImpl.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Constants.h"
@@ -1591,6 +1592,44 @@ TEST_F(DILocationTest, discriminatorSpecialCases) {
   EXPECT_EQ(std::nullopt, L4->cloneByMultiplyingDuplicationFactor(0x1000));
 }
 
+TEST_F(DILocationTest, KeyInstructions) {
+  Context.pImpl->NextAtomGroup = 1;
+
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 1u);
+  DILocation *A1 =
+      DILocation::get(Context, 1, 0, getSubprogram(), nullptr, false, 1, 2);
+  // The group is only applied to the DILocation if we've built LLVM with
+  // EXPERIMENTAL_KEY_INSTRUCTIONS.
+#ifdef EXPERIMENTAL_KEY_INSTRUCTIONS
+  EXPECT_EQ(A1->getAtomGroup(), 1u);
+  EXPECT_EQ(A1->getAtomRank(), 2u);
+#else
+  EXPECT_EQ(A1->getAtomGroup(), 0u);
+  EXPECT_EQ(A1->getAtomRank(), 0u);
+#endif
+
+  // Group number 1 has been "used" so next available is 2.
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 2u);
+
+  // Set a group number higher than current + 1, then check the waterline.
+  DILocation::get(Context, 2, 0, getSubprogram(), nullptr, false, 5, 1);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 6u);
+
+  // The waterline should be unchanged (group <= next).
+  DILocation::get(Context, 3, 0, getSubprogram(), nullptr, false, 4, 1);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 6u);
+  DILocation::get(Context, 3, 0, getSubprogram(), nullptr, false, 5, 1);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 6u);
+
+  // Check the waterline gets incremented by 1.
+  EXPECT_EQ(Context.incNextDILocationAtomGroup(), 6u);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 7u);
+
+  Context.updateDILocationAtomGroupWaterline(8);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 8u);
+  Context.updateDILocationAtomGroupWaterline(7);
+  EXPECT_EQ(Context.pImpl->NextAtomGroup, 8u);
+}
 
 typedef MetadataTest GenericDINodeTest;