[llvm] [NFC][KeyInstr] Add Atom Group (re)mapping (PR #133479)
Orlando Cazalet-Hyams via llvm-commits
llvm-commits at lists.llvm.org
Tue May 6 03:52:24 PDT 2025
https://github.com/OCHyams updated https://github.com/llvm/llvm-project/pull/133479
>From c5829702f686b69f727f1718fcb2524675f4bdbd Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 20 Mar 2025 14:19:34 +0000
Subject: [PATCH 1/4] [NFC][KeyInstr] Add Atom Group (re)mapping
Add:
mapAtomInstance - map the atom group number to a new group.
RemapSourceAtom - apply the mapped atom group number to this instruction.
Modify:
CloneBasicBlock - Call mapAtomInstance on cloned instruction's DebugLocs
if MapAtoms is true (default). Setting to false could
lead to a degraded debugging experience. See code comment.
Optimisations like loop unroll that duplicate instructions need to remap source
atom groups so that each duplicated source construct instance is considered
distinct when determining is_stmt locations.
This commit adds the remapping functionality and a unittest.
---
llvm/include/llvm/IR/ValueMap.h | 5 +
llvm/include/llvm/Transforms/Utils/Cloning.h | 15 ++-
.../llvm/Transforms/Utils/ValueMapper.h | 10 ++
llvm/lib/Transforms/Utils/CloneFunction.cpp | 29 ++++-
llvm/lib/Transforms/Utils/ValueMapper.cpp | 26 +++++
.../Transforms/Utils/CloningTest.cpp | 104 ++++++++++++++++++
6 files changed, 187 insertions(+), 2 deletions(-)
diff --git a/llvm/include/llvm/IR/ValueMap.h b/llvm/include/llvm/IR/ValueMap.h
index d12d639aaa888..398be19331aba 100644
--- a/llvm/include/llvm/IR/ValueMap.h
+++ b/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 = DenseMap<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.
diff --git a/llvm/include/llvm/Transforms/Utils/Cloning.h b/llvm/include/llvm/Transforms/Utils/Cloning.h
index ec1a1d5faa7e9..5e931c6583510 100644
--- a/llvm/include/llvm/Transforms/Utils/Cloning.h
+++ b/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,21 @@ 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.
+///
+/// Set \p MapAtoms to false to skip mapping source atoms for later remapping.
+/// Incorrectly setting false may harm the debugging experience. It's safe to
+/// set false if the cloned basic block is destined for a different function or
+/// if the original block is deleted. Setting true (default) is always safe
+/// (correct) but sometimes unecessary; this option reduces the compile-time
+/// impact of Key Instruction in such cases.
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
diff --git a/llvm/include/llvm/Transforms/Utils/ValueMapper.h b/llvm/include/llvm/Transforms/Utils/ValueMapper.h
index 560df1d3f7f29..63680038172e3 100644
--- a/llvm/include/llvm/Transforms/Utils/ValueMapper.h
+++ b/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 atom instances. 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 for compiler performance reasons rather
+ /// than correctness.
+ RF_DoNotRemapAtoms = 16,
};
inline RemapFlags operator|(RemapFlags LHS, RemapFlags RHS) {
@@ -284,6 +291,9 @@ inline void RemapInstruction(Instruction *I, ValueToValueMapTy &VM,
.remapInstruction(*I);
}
+/// Remap source atom. Called by RemapInstruction.
+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,
diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 35550642a4365..139c689126d94 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/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().incNextAtomGroup();
+ 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);
diff --git a/llvm/lib/Transforms/Utils/ValueMapper.cpp b/llvm/lib/Transforms/Utils/ValueMapper.cpp
index 5e50536a99206..827d708f79e97 100644
--- a/llvm/lib/Transforms/Utils/ValueMapper.cpp
+++ b/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);
+}
diff --git a/llvm/unittests/Transforms/Utils/CloningTest.cpp b/llvm/unittests/Transforms/Utils/CloningTest.cpp
index 03769ff59e372..8f6ec12ec77aa 100644
--- a/llvm/unittests/Transforms/Utils/CloningTest.cpp
+++ b/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
>From 580ad45f0b525c9069a24f0bfd4b9ae15f0bd285 Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 1 May 2025 14:11:38 +0100
Subject: [PATCH 2/4] SmallDenseMap and comment
---
llvm/include/llvm/IR/ValueMap.h | 2 +-
llvm/include/llvm/Transforms/Utils/Cloning.h | 13 +++++++------
2 files changed, 8 insertions(+), 7 deletions(-)
diff --git a/llvm/include/llvm/IR/ValueMap.h b/llvm/include/llvm/IR/ValueMap.h
index 398be19331aba..1a11718bfcdae 100644
--- a/llvm/include/llvm/IR/ValueMap.h
+++ b/llvm/include/llvm/IR/ValueMap.h
@@ -88,7 +88,7 @@ class ValueMap {
using MapT = DenseMap<ValueMapCVH, ValueT, DenseMapInfo<ValueMapCVH>>;
using MDMapT = DenseMap<const Metadata *, TrackingMDRef>;
/// Map {(InlinedAt, old atom number) -> new atom number}.
- using DMAtomT = DenseMap<std::pair<Metadata *, uint64_t>, uint64_t>;
+ using DMAtomT = SmallDenseMap<std::pair<Metadata *, uint64_t>, uint64_t>;
using ExtraData = typename Config::ExtraData;
MapT Map;
diff --git a/llvm/include/llvm/Transforms/Utils/Cloning.h b/llvm/include/llvm/Transforms/Utils/Cloning.h
index 5e931c6583510..24f7ae7f07c24 100644
--- a/llvm/include/llvm/Transforms/Utils/Cloning.h
+++ b/llvm/include/llvm/Transforms/Utils/Cloning.h
@@ -119,12 +119,13 @@ struct ClonedCodeInfo {
/// function, you can specify a ClonedCodeInfo object with the optional fifth
/// parameter.
///
-/// Set \p MapAtoms to false to skip mapping source atoms for later remapping.
-/// Incorrectly setting false may harm the debugging experience. It's safe to
-/// set false if the cloned basic block is destined for a different function or
-/// if the original block is deleted. Setting true (default) is always safe
-/// (correct) but sometimes unecessary; this option reduces the compile-time
-/// impact of Key Instruction in such cases.
+/// \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,
>From 88ee5deba58372edc218e626e9773001bbf844fd Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 1 May 2025 14:14:10 +0100
Subject: [PATCH 3/4] another comment
---
llvm/include/llvm/Transforms/Utils/ValueMapper.h | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/Transforms/Utils/ValueMapper.h b/llvm/include/llvm/Transforms/Utils/ValueMapper.h
index 63680038172e3..b25f6b964f011 100644
--- a/llvm/include/llvm/Transforms/Utils/ValueMapper.h
+++ b/llvm/include/llvm/Transforms/Utils/ValueMapper.h
@@ -106,11 +106,11 @@ enum RemapFlags {
/// to self. Illegal if RF_IgnoreMissingLocals is also set.
RF_NullMapMissingGlobalValues = 8,
- /// Do not remap atom instances. Only safe if to do this if the cloned
+ /// 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 for compiler performance reasons rather
- /// than correctness.
+ /// don't use this flag. It's used when remapping is known to be un-necessary
+ /// to save some compile-time.
RF_DoNotRemapAtoms = 16,
};
>From 0847aa54164c5b7b37d1dc50c9ef028912b6b64f Mon Sep 17 00:00:00 2001
From: Orlando Cazalet-Hyams <orlando.hyams at sony.com>
Date: Thu, 1 May 2025 14:24:41 +0100
Subject: [PATCH 4/4] less terse comment
---
llvm/include/llvm/Transforms/Utils/ValueMapper.h | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/llvm/include/llvm/Transforms/Utils/ValueMapper.h b/llvm/include/llvm/Transforms/Utils/ValueMapper.h
index b25f6b964f011..ac71996872277 100644
--- a/llvm/include/llvm/Transforms/Utils/ValueMapper.h
+++ b/llvm/include/llvm/Transforms/Utils/ValueMapper.h
@@ -291,7 +291,10 @@ inline void RemapInstruction(Instruction *I, ValueToValueMapTy &VM,
.remapInstruction(*I);
}
-/// Remap source atom. Called by RemapInstruction.
+/// 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
More information about the llvm-commits
mailing list