[NFC][KeyInstr] Add Atom Group (re)mapping

llvm/include/llvm/IR/ValueMap.h

@@ -87,6 +87,8 @@ class ValueMap {
   using ValueMapCVH = ValueMapCallbackVH<KeyT, ValueT, Config>;
   using MapT = DenseMap<ValueMapCVH, ValueT, DenseMapInfo<ValueMapCVH>>;
   using MDMapT = DenseMap<const Metadata *, TrackingMDRef>;
+  /// Map {(InlinedAt, old atom number) -> new atom number}.
+  using DMAtomT = SmallDenseMap<std::pair<Metadata *, uint64_t>, uint64_t>;
   using ExtraData = typename Config::ExtraData;
 
   MapT Map;
@@ -117,6 +119,8 @@ class ValueMap {
     return *MDMap;
   }
   std::optional<MDMapT> &getMDMap() { return MDMap; }
+  /// Map {(InlinedAt, old atom number) -> new atom number}.
+  DMAtomT AtomMap;
 
   /// Get the mapped metadata, if it's in the map.
   std::optional<Metadata *> getMappedMD(const Metadata *MD) const {
@@ -145,6 +149,7 @@ class ValueMap {
   void clear() {
     Map.clear();
     MDMap.reset();
+    AtomMap.clear();
   }
 
   /// Return 1 if the specified key is in the map, 0 otherwise.

llvm/include/llvm/Transforms/Utils/Cloning.h

@@ -22,6 +22,7 @@
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <functional>
@@ -117,9 +118,22 @@ struct ClonedCodeInfo {
 /// If you would like to collect additional information about the cloned
 /// function, you can specify a ClonedCodeInfo object with the optional fifth
 /// parameter.
+///
+/// \p MapAtoms indicates whether source location atoms should be mapped for
+/// later remapping. Must be true when you duplicate a code path and a source
+/// location is intended to appear twice in the generated instructions. Can be
+/// set to false if you are transplanting code from one place to another.
+/// Setting true (default) is always safe (won't produce incorrect debug info)
+/// but is sometimes unnecessary, causing extra work that could be avoided by
+/// setting the parameter to false.
 BasicBlock *CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
                             const Twine &NameSuffix = "", Function *F = nullptr,
-                            ClonedCodeInfo *CodeInfo = nullptr);
+                            ClonedCodeInfo *CodeInfo = nullptr,
+                            bool MapAtoms = true);
+
+/// Mark a cloned instruction as a new instance so that its source loc can
+/// be updated when remapped.
+void mapAtomInstance(const DebugLoc &DL, ValueToValueMapTy &VMap);
 
 /// Return a copy of the specified function and add it to that
 /// function's module.  Also, any references specified in the VMap are changed

llvm/include/llvm/Transforms/Utils/ValueMapper.h

@@ -105,6 +105,13 @@ enum RemapFlags {
   /// Any global values not in value map are mapped to null instead of mapping
   /// to self.  Illegal if RF_IgnoreMissingLocals is also set.
   RF_NullMapMissingGlobalValues = 8,
+
+  /// Do not remap source location atoms. Only safe if to do this if the cloned
+  /// instructions being remapped are inserted into a new function, or an
+  /// existing function where the inlined-at fields are updated. If in doubt,
+  /// don't use this flag. It's used when remapping is known to be un-necessary
+  /// to save some compile-time.
+  RF_DoNotRemapAtoms = 16,
 };
 
 inline RemapFlags operator|(RemapFlags LHS, RemapFlags RHS) {
@@ -284,6 +291,12 @@ inline void RemapInstruction(Instruction *I, ValueToValueMapTy &VM,
       .remapInstruction(*I);
 }
 
+/// Remap source location atom. Called by RemapInstruction. This updates the
+/// instruction's atom group number if it has been mapped (e.g. with
+/// llvm::mapAtomInstance), which is necessary to distinguish source code
+/// atoms on duplicated code paths.
+void RemapSourceAtom(Instruction *I, ValueToValueMapTy &VM);
+
 /// Remap the Values used in the DbgRecord \a DR using the value map \a
 /// VM.
 inline void RemapDbgRecord(Module *M, DbgRecord *DR, ValueToValueMapTy &VM,

llvm/lib/Transforms/Utils/CloneFunction.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/InstructionSimplify.h"
@@ -40,6 +41,27 @@ using namespace llvm;
 
 #define DEBUG_TYPE "clone-function"
 
+STATISTIC(RemappedAtomMax, "Highest global NextAtomGroup (after mapping)");
+
+void llvm::mapAtomInstance(const DebugLoc &DL, ValueToValueMapTy &VMap) {
+  auto CurGroup = DL.get()->getAtomGroup();
+  if (!CurGroup)
+    return;
+
+  // Try inserting a new entry. If there's already a mapping for this atom
+  // then there's nothing to do.
+  auto [It, Inserted] = VMap.AtomMap.insert({{DL.getInlinedAt(), CurGroup}, 0});
+  if (!Inserted)
+    return;
+
+  // Map entry to a new atom group.
+  uint64_t NewGroup = DL->getContext().incNextDILocationAtomGroup();
+  assert(NewGroup > CurGroup && "Next should always be greater than current");
+  It->second = NewGroup;
+
+  RemappedAtomMax = std::max<uint64_t>(NewGroup, RemappedAtomMax);
+}
+
 namespace {
 void collectDebugInfoFromInstructions(const Function &F,
                                       DebugInfoFinder &DIFinder) {
@@ -79,7 +101,7 @@ MetadataPredicate createIdentityMDPredicate(const Function &F,
 /// See comments in Cloning.h.
 BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
                                   const Twine &NameSuffix, Function *F,
-                                  ClonedCodeInfo *CodeInfo) {
+                                  ClonedCodeInfo *CodeInfo, bool MapAtoms) {
   BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "", F);
   NewBB->IsNewDbgInfoFormat = BB->IsNewDbgInfoFormat;
   if (BB->hasName())
@@ -98,6 +120,11 @@ BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
 
     VMap[&I] = NewInst; // Add instruction map to value.
 
+    if (MapAtoms) {
+      if (const DebugLoc &DL = NewInst->getDebugLoc())
+        mapAtomInstance(DL.get(), VMap);
+    }
+
     if (isa<CallInst>(I) && !I.isDebugOrPseudoInst()) {
       hasCalls = true;
       hasMemProfMetadata |= I.hasMetadata(LLVMContext::MD_memprof);

llvm/lib/Transforms/Utils/ValueMapper.cpp

@@ -37,6 +37,7 @@
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include <cassert>
 #include <limits>
@@ -1010,6 +1011,10 @@ void Mapper::remapInstruction(Instruction *I) {
       I->setMetadata(MI.first, New);
   }
 
+  // Remap source location atom instance.
+  if (!(Flags & RF_DoNotRemapAtoms))
+    RemapSourceAtom(I, getVM());
+
   if (!TypeMapper)
     return;
 
@@ -1292,3 +1297,24 @@ void ValueMapper::scheduleMapGlobalIFunc(GlobalIFunc &GI, Constant &Resolver,
 void ValueMapper::scheduleRemapFunction(Function &F, unsigned MCID) {
   getAsMapper(pImpl)->scheduleRemapFunction(F, MCID);
 }
+
+void llvm::RemapSourceAtom(Instruction *I, ValueToValueMapTy &VM) {
+  const DebugLoc &DL = I->getDebugLoc();
+  if (!DL)
+    return;
+
+  auto AtomGroup = DL->getAtomGroup();
+  if (!AtomGroup)
+    return;
+
+  auto R = VM.AtomMap.find({DL->getInlinedAt(), AtomGroup});
+  if (R == VM.AtomMap.end())
+    return;
+  AtomGroup = R->second;
+
+  // Remap the atom group and copy all other fields.
+  DILocation *New = DILocation::get(
+      I->getContext(), DL.getLine(), DL.getCol(), DL.getScope(),
+      DL.getInlinedAt(), DL.isImplicitCode(), AtomGroup, DL->getAtomRank());
+  I->setDebugLoc(New);
+}

llvm/unittests/Transforms/Utils/CloningTest.cpp

@@ -14,18 +14,21 @@
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Argument.h"
+#include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -1162,4 +1165,105 @@ declare i64 @foo(i32 noundef) local_unnamed_addr
   auto NewM = llvm::CloneModule(*M);
   EXPECT_FALSE(verifyModule(*NewM, &errs()));
 }
+
+TEST_F(CloneInstruction, cloneKeyInstructions) {
+  LLVMContext Context;
+
+  std::unique_ptr<Module> M = parseIR(Context, R"(
+    define void @test(ptr align 4 %dst) !dbg !3 {
+      store i64 1, ptr %dst, align 4, !dbg !6
+      store i64 2, ptr %dst, align 4, !dbg !7
+      store i64 3, ptr %dst, align 4, !dbg !8
+      ret void, !dbg !9
+    }
+
+    !llvm.dbg.cu = !{!0}
+    !llvm.module.flags = !{!2}
+    !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1)
+    !1 = !DIFile(filename: "test.cpp",  directory: "")
+    !2 = !{i32 1, !"Debug Info Version", i32 3}
+    !3 = distinct !DISubprogram(name: "test", scope: !0, unit: !0)
+    !4 = distinct !DISubprogram(name: "inlined", scope: !0, unit: !0, retainedNodes: !{!5})
+    !5 = !DILocalVariable(name: "awaitables", scope: !4)
+    !6 = !DILocation(line: 1, scope: !4, inlinedAt: !8, atomGroup: 1, atomRank: 1)
+    !7 = !DILocation(line: 2, scope: !3, atomGroup: 1, atomRank: 1)
+    !8 = !DILocation(line: 3, scope: !3, atomGroup: 1, atomRank: 1)
+    !9 = !DILocation(line: 4, scope: !3, atomGroup: 2, atomRank: 1)
+  )");
+
+  ASSERT_FALSE(verifyModule(*M, &errs()));
+
+#ifdef EXPERIMENTAL_KEY_INSTRUCTIONS
+#define EXPECT_ATOM(Inst, G)                                                   \
+  EXPECT_TRUE(Inst->getDebugLoc());                                            \
+  EXPECT_EQ(Inst->getDebugLoc()->getAtomGroup(), uint64_t(G));
+#else
+#define EXPECT_ATOM(Inst, G) (void)Inst;
+#endif
+
+  Function *F = M->getFunction("test");
+  BasicBlock *BB = &*F->begin();
+  ASSERT_TRUE(F);
+  Instruction *Store1 = &*BB->begin();
+  Instruction *Store2 = Store1->getNextNode();
+  Instruction *Store3 = Store2->getNextNode();
+  Instruction *Ret = Store3->getNextNode();
+
+  // Test the remapping works as expected outside of cloning.
+  ValueToValueMapTy VM;
+  // Store1 and Store2 have the same atomGroup number, but have different
+  // inlining scopes, so only Store1 should change group.
+  mapAtomInstance(Store1->getDebugLoc(), VM);
+  for (Instruction &I : *BB)
+    RemapSourceAtom(&I, VM);
+  EXPECT_ATOM(Store1, 3);
+  EXPECT_ATOM(Store2, 1);
+  EXPECT_ATOM(Store3, 1);
+  EXPECT_ATOM(Ret, 2);
+  VM.clear();
+
+  // Store2 and Store3 have the same group number; both should get remapped.
+  mapAtomInstance(Store2->getDebugLoc(), VM);
+  for (Instruction &I : *BB)
+    RemapSourceAtom(&I, VM);
+  EXPECT_ATOM(Store1, 3);
+  EXPECT_ATOM(Store2, 4);
+  EXPECT_ATOM(Store3, 4);
+  EXPECT_ATOM(Ret, 2);
+  VM.clear();
+
+  // Cloning BB with MapAtoms=false should clone the atom numbers.
+  BasicBlock *BB2 =
+      CloneBasicBlock(BB, VM, "", nullptr, nullptr, /*MapAtoms*/ false);
+  for (Instruction &I : *BB2)
+    RemapSourceAtom(&I, VM);
+  Store1 = &*BB2->begin();
+  Store2 = Store1->getNextNode();
+  Store3 = Store2->getNextNode();
+  Ret = Store3->getNextNode();
+  EXPECT_ATOM(Store1, 3);
+  EXPECT_ATOM(Store2, 4);
+  EXPECT_ATOM(Store3, 4);
+  EXPECT_ATOM(Ret, 2);
+  VM.clear();
+  delete BB2;
+
+  // Cloning BB with MapAtoms=true should map the atom numbers.
+  BasicBlock *BB3 =
+      CloneBasicBlock(BB, VM, "", nullptr, nullptr, /*MapAtoms*/ true);
+  for (Instruction &I : *BB3)
+    RemapSourceAtom(&I, VM);
+  Store1 = &*BB3->begin();
+  Store2 = Store1->getNextNode();
+  Store3 = Store2->getNextNode();
+  Ret = Store3->getNextNode();
+  EXPECT_ATOM(Store1, 5);
+  EXPECT_ATOM(Store2, 6);
+  EXPECT_ATOM(Store3, 6);
+  EXPECT_ATOM(Ret, 7);
+  VM.clear();
+  delete BB3;
+#undef EXPECT_ATOM
+}
+
 } // namespace