diff --git a/llvm/include/llvm/IR/ValueMap.h b/llvm/include/llvm/IR/ValueMap.h index d12d639aaa888..1a11718bfcdae 100644 --- a/llvm/include/llvm/IR/ValueMap.h +++ b/llvm/include/llvm/IR/ValueMap.h @@ -87,6 +87,8 @@ class ValueMap { using ValueMapCVH = ValueMapCallbackVH; using MapT = DenseMap>; using MDMapT = DenseMap; + /// Map {(InlinedAt, old atom number) -> new atom number}. + using DMAtomT = SmallDenseMap, uint64_t>; using ExtraData = typename Config::ExtraData; MapT Map; @@ -117,6 +119,8 @@ class ValueMap { return *MDMap; } std::optional &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 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..24f7ae7f07c24 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 @@ -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 diff --git a/llvm/include/llvm/Transforms/Utils/ValueMapper.h b/llvm/include/llvm/Transforms/Utils/ValueMapper.h index 560df1d3f7f29..ac71996872277 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 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, diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp index 35550642a4365..7b3e0729f5a74 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().incNextDILocationAtomGroup(); + assert(NewGroup > CurGroup && "Next should always be greater than current"); + It->second = NewGroup; + + RemappedAtomMax = std::max(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(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 #include @@ -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 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