From 73f5f83b192b0a27f7edae5365c247961d9f1bd9 Mon Sep 17 00:00:00 2001 From: Haohai Wen Date: Thu, 4 Jul 2024 09:52:38 +0800 Subject: [PATCH] [BasicBlockSections] Using MBBSectionID as DenseMap key (#97295) getSectionIDNum may return same value for two different MBBSectionID. e.g. A Cold type MBBSectionID with number 0 and a Default type MBBSectionID with number 2 get same value 2 from getSectionIDNum. This may lead to overwrite of MBBSectionRanges. Using MBBSectionID itself as DenseMap key is better choice. --- llvm/include/llvm/CodeGen/AsmPrinter.h | 4 +-- llvm/include/llvm/CodeGen/MachineBasicBlock.h | 26 ++++++++++++++----- llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp | 17 +++++++----- .../CodeGen/AsmPrinter/DwarfCompileUnit.cpp | 2 +- llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp | 4 +-- llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp | 4 +-- 6 files changed, 37 insertions(+), 20 deletions(-) diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h index a60dce30c4a6c9..dc00bd57d655d1 100644 --- a/llvm/include/llvm/CodeGen/AsmPrinter.h +++ b/llvm/include/llvm/CodeGen/AsmPrinter.h @@ -137,7 +137,7 @@ class AsmPrinter : public MachineFunctionPass { MCSymbol *BeginLabel, *EndLabel; }; - MapVector MBBSectionRanges; + MapVector MBBSectionRanges; /// Map global GOT equivalent MCSymbols to GlobalVariables and keep track of /// its number of uses by other globals. @@ -157,7 +157,7 @@ class AsmPrinter : public MachineFunctionPass { /// Map a basic block section ID to the exception symbol associated with that /// section. Map entries are assigned and looked up via /// AsmPrinter::getMBBExceptionSym. - DenseMap MBBSectionExceptionSyms; + DenseMap MBBSectionExceptionSyms; // The symbol used to represent the start of the current BB section of the // function. This is used to calculate the size of the BB section. diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h index e4919ecabd7053..562d37ef32f543 100644 --- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h +++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h @@ -13,6 +13,7 @@ #ifndef LLVM_CODEGEN_MACHINEBASICBLOCK_H #define LLVM_CODEGEN_MACHINEBASICBLOCK_H +#include "llvm/ADT/DenseMapInfo.h" #include "llvm/ADT/GraphTraits.h" #include "llvm/ADT/SparseBitVector.h" #include "llvm/ADT/ilist.h" @@ -74,6 +75,25 @@ struct MBBSectionID { MBBSectionID(SectionType T) : Type(T), Number(0) {} }; +template <> struct DenseMapInfo { + using TypeInfo = DenseMapInfo; + using NumberInfo = DenseMapInfo; + + static inline MBBSectionID getEmptyKey() { + return MBBSectionID(NumberInfo::getEmptyKey()); + } + static inline MBBSectionID getTombstoneKey() { + return MBBSectionID(NumberInfo::getTombstoneKey()); + } + static unsigned getHashValue(const MBBSectionID &SecID) { + return detail::combineHashValue(TypeInfo::getHashValue(SecID.Type), + NumberInfo::getHashValue(SecID.Number)); + } + static bool isEqual(const MBBSectionID &LHS, const MBBSectionID &RHS) { + return LHS == RHS; + } +}; + // This structure represents the information for a basic block pertaining to // the basic block sections profile. struct UniqueBBID { @@ -658,12 +678,6 @@ class MachineBasicBlock /// Returns the section ID of this basic block. MBBSectionID getSectionID() const { return SectionID; } - /// Returns the unique section ID number of this basic block. - unsigned getSectionIDNum() const { - return ((unsigned)MBBSectionID::SectionType::Cold) - - ((unsigned)SectionID.Type) + SectionID.Number; - } - /// Sets the fixed BBID of this basic block. void setBBID(const UniqueBBID &V) { assert(!BBID.has_value() && "Cannot change BBID."); diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp index c52cbff689dc5d..1391893e55a52a 100644 --- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp +++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp @@ -1374,7 +1374,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) { OutStreamer->emitULEB128IntValue(MBBSectionRanges.size()); } // Number of blocks in each MBB section. - MapVector MBBSectionNumBlocks; + MapVector MBBSectionNumBlocks; const MCSymbol *PrevMBBEndSymbol = nullptr; if (!Features.MultiBBRange) { OutStreamer->AddComment("function address"); @@ -1388,7 +1388,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) { BBCount++; if (MBB.isEndSection()) { // Store each section's basic block count when it ends. - MBBSectionNumBlocks[MBB.getSectionIDNum()] = BBCount; + MBBSectionNumBlocks[MBB.getSectionID()] = BBCount; // Reset the count for the next section. BBCount = 0; } @@ -1404,8 +1404,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) { OutStreamer->AddComment("base address"); OutStreamer->emitSymbolValue(MBBSymbol, getPointerSize()); OutStreamer->AddComment("number of basic blocks"); - OutStreamer->emitULEB128IntValue( - MBBSectionNumBlocks[MBB.getSectionIDNum()]); + OutStreamer->emitULEB128IntValue(MBBSectionNumBlocks[MBB.getSectionID()]); PrevMBBEndSymbol = MBBSymbol; } // TODO: Remove this check when version 1 is deprecated. @@ -1855,7 +1854,9 @@ void AsmPrinter::emitFunctionBody() { OutContext); OutStreamer->emitELFSize(CurrentSectionBeginSym, SizeExp); } - MBBSectionRanges[MBB.getSectionIDNum()] = + assert(!MBBSectionRanges.contains(MBB.getSectionID()) && + "Overwrite section range"); + MBBSectionRanges[MBB.getSectionID()] = MBBSectionRange{CurrentSectionBeginSym, MBB.getEndSymbol()}; } } @@ -1972,7 +1973,9 @@ void AsmPrinter::emitFunctionBody() { for (auto &Handler : Handlers) Handler->markFunctionEnd(); - MBBSectionRanges[MF->front().getSectionIDNum()] = + assert(!MBBSectionRanges.contains(MF->front().getSectionID()) && + "Overwrite section range"); + MBBSectionRanges[MF->front().getSectionID()] = MBBSectionRange{CurrentFnBegin, CurrentFnEnd}; // Print out jump tables referenced by the function. @@ -2536,7 +2539,7 @@ bool AsmPrinter::doFinalization(Module &M) { } MCSymbol *AsmPrinter::getMBBExceptionSym(const MachineBasicBlock &MBB) { - auto Res = MBBSectionExceptionSyms.try_emplace(MBB.getSectionIDNum()); + auto Res = MBBSectionExceptionSyms.try_emplace(MBB.getSectionID()); if (Res.second) Res.first->second = createTempSymbol("exception"); return Res.first->second; diff --git a/llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp b/llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp index c1e7f01f0eba59..c1e8355353cfdc 100644 --- a/llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp +++ b/llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp @@ -684,7 +684,7 @@ void DwarfCompileUnit::attachRangesOrLowHighPC( // the order of blocks will be frozen beyond this point. do { if (MBB->sameSection(EndMBB) || MBB->isEndSection()) { - auto MBBSectionRange = Asm->MBBSectionRanges[MBB->getSectionIDNum()]; + auto MBBSectionRange = Asm->MBBSectionRanges[MBB->getSectionID()]; List.push_back( {MBB->sameSection(BeginMBB) ? BeginLabel : MBBSectionRange.BeginLabel, diff --git a/llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp b/llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp index 2addf938c8b638..80cd5ec501f25d 100644 --- a/llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp +++ b/llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp @@ -1713,7 +1713,7 @@ bool DwarfDebug::buildLocationList(SmallVectorImpl &DebugLoc, const MCSymbol *EndLabel; if (std::next(EI) == Entries.end()) { const MachineBasicBlock &EndMBB = Asm->MF->back(); - EndLabel = Asm->MBBSectionRanges[EndMBB.getSectionIDNum()].EndLabel; + EndLabel = Asm->MBBSectionRanges[EndMBB.getSectionID()].EndLabel; if (EI->isClobber()) EndMI = EI->getInstr(); } @@ -2064,7 +2064,7 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) { bool PrevInstInSameSection = (!PrevInstBB || - PrevInstBB->getSectionIDNum() == MI->getParent()->getSectionIDNum()); + PrevInstBB->getSectionID() == MI->getParent()->getSectionID()); if (DL == PrevInstLoc && PrevInstInSameSection) { // If we have an ongoing unspecified location, nothing to do here. if (!DL) diff --git a/llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp b/llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp index 32239535e4d022..1c603f5988ad1e 100644 --- a/llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp +++ b/llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp @@ -253,8 +253,8 @@ void EHStreamer::computeCallSiteTable( // We start a call-site range upon function entry and at the beginning of // every basic block section. CallSiteRanges.push_back( - {Asm->MBBSectionRanges[MBB.getSectionIDNum()].BeginLabel, - Asm->MBBSectionRanges[MBB.getSectionIDNum()].EndLabel, + {Asm->MBBSectionRanges[MBB.getSectionID()].BeginLabel, + Asm->MBBSectionRanges[MBB.getSectionID()].EndLabel, Asm->getMBBExceptionSym(MBB), CallSites.size()}); PreviousIsInvoke = false; SawPotentiallyThrowing = false;