[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

3 files changed, 21 insertions, 0 deletions

diff --git a/llvm/lib/IR/DebugInfoMetadata.cpp b/llvm/lib/IR/DebugInfoMetadata.cpp
index e8a58f1..0973fcf 100644
--- a/llvm/lib/IR/DebugInfoMetadata.cpp
+++ b/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.
diff --git a/llvm/lib/IR/LLVMContext.cpp b/llvm/lib/IR/LLVMContext.cpp
index 447e5d9..57532cd 100644
--- a/llvm/lib/IR/LLVMContext.cpp
+++ b/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++;
+}
diff --git a/llvm/lib/IR/LLVMContextImpl.h b/llvm/lib/IR/LLVMContextImpl.h
index 5c2b5cd..21f5c06 100644
--- a/llvm/lib/IR/LLVMContextImpl.h
+++ b/llvm/lib/IR/LLVMContextImpl.h
@@ -1880,6 +1880,16 @@ public:
 
   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