From 2f753b4ae96604ba76202b5cf011be1ec0adf49f Mon Sep 17 00:00:00 2001 From: Carl Ritson Date: Wed, 11 Oct 2023 19:10:50 +0900 Subject: [PATCH 1/3] [MachineBasicBlock] Fix use after free in SplitCriticalEdge Remove use after free when attempting to update SlotIndexes in MachineBasicBlock::SplitCriticalEdge. Add delegate mechanism for MachineBasicBlock to observe instructions added or removed when calling target specific functions for updating branches and update SlotIndexes appropriately. --- llvm/include/llvm/CodeGen/MachineBasicBlock.h | 37 +++++++++++ llvm/lib/CodeGen/MachineBasicBlock.cpp | 63 ++++++++++--------- 2 files changed, 69 insertions(+), 31 deletions(-) diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h index 15c4fcd8399c1..4f9df77141deb 100644 --- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h +++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h @@ -14,6 +14,7 @@ #define LLVM_CODEGEN_MACHINEBASICBLOCK_H #include "llvm/ADT/GraphTraits.h" +#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SparseBitVector.h" #include "llvm/ADT/ilist.h" #include "llvm/ADT/iterator_range.h" @@ -106,6 +107,16 @@ class MachineBasicBlock : PhysReg(PhysReg), LaneMask(LaneMask) {} }; + class Delegate { + virtual void anchor(); + + public: + virtual ~Delegate() = default; + + virtual void MBB_NoteInsertion(MachineInstr &MI) = 0; + virtual void MBB_NoteRemoval(MachineInstr &MI) = 0; + }; + private: using Instructions = ilist>; @@ -233,6 +244,9 @@ class MachineBasicBlock /// Return a formatted string to identify this block and its parent function. std::string getFullName() const; + /// Delegates to inform of instruction addition and removal. + SmallPtrSet TheDelegates; + /// Test whether this block is used as something other than the target /// of a terminator, exception-handling target, or jump table. This is /// either the result of an IR-level "blockaddress", or some form @@ -1184,6 +1198,29 @@ class MachineBasicBlock /// MachineBranchProbabilityInfo class. BranchProbability getSuccProbability(const_succ_iterator Succ) const; + void resetDelegate(Delegate *delegate) { + assert(TheDelegates.count(delegate) && + "Only an existing delegate can perform reset!"); + TheDelegates.erase(delegate); + } + + void addDelegate(Delegate *delegate) { + assert(delegate && !TheDelegates.count(delegate) && + "Attempted to add null delegate, or to change it without " + "first resetting it!"); + TheDelegates.insert(delegate); + } + + void noteInsertion(MachineInstr &MI) { + for (auto *TheDelegate : TheDelegates) + TheDelegate->MBB_NoteInsertion(MI); + } + + void noteRemoval(MachineInstr &MI) { + for (auto *TheDelegate : TheDelegates) + TheDelegate->MBB_NoteRemoval(MI); + } + private: /// Return probability iterator corresponding to the I successor iterator. probability_iterator getProbabilityIterator(succ_iterator I); diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp index 7d3d8b6fba1b7..38dffc031dcb6 100644 --- a/llvm/lib/CodeGen/MachineBasicBlock.cpp +++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp @@ -156,12 +156,14 @@ void ilist_traits::addNodeToList(MachineInstr *N) { MachineFunction *MF = Parent->getParent(); N->addRegOperandsToUseLists(MF->getRegInfo()); MF->handleInsertion(*N); + Parent->noteInsertion(*N); } /// When we remove an instruction from a basic block list, we update its parent /// pointer and remove its operands from reg use/def lists if appropriate. void ilist_traits::removeNodeFromList(MachineInstr *N) { assert(N->getParent() && "machine instruction not in a basic block"); + N->getParent()->noteRemoval(*N); // Remove from the use/def lists. if (MachineFunction *MF = N->getMF()) { @@ -1097,6 +1099,31 @@ static bool jumpTableHasOtherUses(const MachineFunction &MF, return false; } +class MBBSplitCriticalEdgeDelegate : public MachineBasicBlock::Delegate { +private: + MachineBasicBlock *MBB; + SlotIndexes *Indexes; + +public: + MBBSplitCriticalEdgeDelegate(MachineBasicBlock *Block, + SlotIndexes *IndexesToUpdate) + : MBB(Block), Indexes(IndexesToUpdate) { + MBB->addDelegate(this); + } + + ~MBBSplitCriticalEdgeDelegate() { MBB->resetDelegate(this); } + + void MBB_NoteInsertion(MachineInstr &MI) override { + if (Indexes) + Indexes->insertMachineInstrInMaps(MI); + } + + void MBB_NoteRemoval(MachineInstr &MI) override { + if (Indexes) + Indexes->removeMachineInstrFromMaps(MI); + } +}; + MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge( MachineBasicBlock *Succ, Pass &P, std::vector> *LiveInSets) { @@ -1170,51 +1197,23 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge( ReplaceUsesOfBlockWith(Succ, NMBB); - // If updateTerminator() removes instructions, we need to remove them from - // SlotIndexes. - SmallVector Terminators; - if (Indexes) { - for (MachineInstr &MI : - llvm::make_range(getFirstInstrTerminator(), instr_end())) - Terminators.push_back(&MI); - } - // Since we replaced all uses of Succ with NMBB, that should also be treated // as the fallthrough successor if (Succ == PrevFallthrough) PrevFallthrough = NMBB; - if (!ChangedIndirectJump) + if (!ChangedIndirectJump) { + MBBSplitCriticalEdgeDelegate SlotUpdater(this, Indexes); updateTerminator(PrevFallthrough); - - if (Indexes) { - SmallVector NewTerminators; - for (MachineInstr &MI : - llvm::make_range(getFirstInstrTerminator(), instr_end())) - NewTerminators.push_back(&MI); - - for (MachineInstr *Terminator : Terminators) { - if (!is_contained(NewTerminators, Terminator)) - Indexes->removeMachineInstrFromMaps(*Terminator); - } } // Insert unconditional "jump Succ" instruction in NMBB if necessary. NMBB->addSuccessor(Succ); if (!NMBB->isLayoutSuccessor(Succ)) { + MBBSplitCriticalEdgeDelegate SlotUpdater(NMBB, Indexes); SmallVector Cond; const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo(); TII->insertBranch(*NMBB, Succ, nullptr, Cond, DL); - - if (Indexes) { - for (MachineInstr &MI : NMBB->instrs()) { - // Some instructions may have been moved to NMBB by updateTerminator(), - // so we first remove any instruction that already has an index. - if (Indexes->hasIndex(MI)) - Indexes->removeMachineInstrFromMaps(MI); - Indexes->insertMachineInstrInMaps(MI); - } - } } // Fix PHI nodes in Succ so they refer to NMBB instead of this. @@ -1743,3 +1742,5 @@ bool MachineBasicBlock::sizeWithoutDebugLargerThan(unsigned Limit) const { const MBBSectionID MBBSectionID::ColdSectionID(MBBSectionID::SectionType::Cold); const MBBSectionID MBBSectionID::ExceptionSectionID(MBBSectionID::SectionType::Exception); + +void MachineBasicBlock::Delegate::anchor() {} From bd96a52e007df640f7681ae8f4a8c2ea1eb94a66 Mon Sep 17 00:00:00 2001 From: Carl Ritson Date: Thu, 12 Oct 2023 15:46:24 +0900 Subject: [PATCH 2/3] Address reviewer feedback: - Use MachineFunction delegate instead --- llvm/include/llvm/CodeGen/MachineBasicBlock.h | 36 ------------------- llvm/lib/CodeGen/MachineBasicBlock.cpp | 20 +++++------ 2 files changed, 8 insertions(+), 48 deletions(-) diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h index 4f9df77141deb..181156686c908 100644 --- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h +++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h @@ -107,16 +107,6 @@ class MachineBasicBlock : PhysReg(PhysReg), LaneMask(LaneMask) {} }; - class Delegate { - virtual void anchor(); - - public: - virtual ~Delegate() = default; - - virtual void MBB_NoteInsertion(MachineInstr &MI) = 0; - virtual void MBB_NoteRemoval(MachineInstr &MI) = 0; - }; - private: using Instructions = ilist>; @@ -244,9 +234,6 @@ class MachineBasicBlock /// Return a formatted string to identify this block and its parent function. std::string getFullName() const; - /// Delegates to inform of instruction addition and removal. - SmallPtrSet TheDelegates; - /// Test whether this block is used as something other than the target /// of a terminator, exception-handling target, or jump table. This is /// either the result of an IR-level "blockaddress", or some form @@ -1198,29 +1185,6 @@ class MachineBasicBlock /// MachineBranchProbabilityInfo class. BranchProbability getSuccProbability(const_succ_iterator Succ) const; - void resetDelegate(Delegate *delegate) { - assert(TheDelegates.count(delegate) && - "Only an existing delegate can perform reset!"); - TheDelegates.erase(delegate); - } - - void addDelegate(Delegate *delegate) { - assert(delegate && !TheDelegates.count(delegate) && - "Attempted to add null delegate, or to change it without " - "first resetting it!"); - TheDelegates.insert(delegate); - } - - void noteInsertion(MachineInstr &MI) { - for (auto *TheDelegate : TheDelegates) - TheDelegate->MBB_NoteInsertion(MI); - } - - void noteRemoval(MachineInstr &MI) { - for (auto *TheDelegate : TheDelegates) - TheDelegate->MBB_NoteRemoval(MI); - } - private: /// Return probability iterator corresponding to the I successor iterator. probability_iterator getProbabilityIterator(succ_iterator I); diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp index 38dffc031dcb6..11b85251d814f 100644 --- a/llvm/lib/CodeGen/MachineBasicBlock.cpp +++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp @@ -156,14 +156,12 @@ void ilist_traits::addNodeToList(MachineInstr *N) { MachineFunction *MF = Parent->getParent(); N->addRegOperandsToUseLists(MF->getRegInfo()); MF->handleInsertion(*N); - Parent->noteInsertion(*N); } /// When we remove an instruction from a basic block list, we update its parent /// pointer and remove its operands from reg use/def lists if appropriate. void ilist_traits::removeNodeFromList(MachineInstr *N) { assert(N->getParent() && "machine instruction not in a basic block"); - N->getParent()->noteRemoval(*N); // Remove from the use/def lists. if (MachineFunction *MF = N->getMF()) { @@ -1099,26 +1097,26 @@ static bool jumpTableHasOtherUses(const MachineFunction &MF, return false; } -class MBBSplitCriticalEdgeDelegate : public MachineBasicBlock::Delegate { +class MBBSplitCriticalEdgeDelegate : public MachineFunction::Delegate { private: - MachineBasicBlock *MBB; + MachineFunction *MF; SlotIndexes *Indexes; public: - MBBSplitCriticalEdgeDelegate(MachineBasicBlock *Block, + MBBSplitCriticalEdgeDelegate(MachineBasicBlock *MBB, SlotIndexes *IndexesToUpdate) - : MBB(Block), Indexes(IndexesToUpdate) { - MBB->addDelegate(this); + : MF(MBB->getParent()), Indexes(IndexesToUpdate) { + MF->setDelegate(this); } - ~MBBSplitCriticalEdgeDelegate() { MBB->resetDelegate(this); } + ~MBBSplitCriticalEdgeDelegate() { MF->resetDelegate(this); } - void MBB_NoteInsertion(MachineInstr &MI) override { + void MF_HandleInsertion(MachineInstr &MI) override { if (Indexes) Indexes->insertMachineInstrInMaps(MI); } - void MBB_NoteRemoval(MachineInstr &MI) override { + void MF_HandleRemoval(MachineInstr &MI) override { if (Indexes) Indexes->removeMachineInstrFromMaps(MI); } @@ -1742,5 +1740,3 @@ bool MachineBasicBlock::sizeWithoutDebugLargerThan(unsigned Limit) const { const MBBSectionID MBBSectionID::ColdSectionID(MBBSectionID::SectionType::Cold); const MBBSectionID MBBSectionID::ExceptionSectionID(MBBSectionID::SectionType::Exception); - -void MachineBasicBlock::Delegate::anchor() {} From 6e2ba5aedde4669bfb646aa02896d927a89d876d Mon Sep 17 00:00:00 2001 From: Carl Ritson Date: Fri, 13 Oct 2023 14:02:43 +0900 Subject: [PATCH 3/3] Address reviewer nits: - Remove unnecessary include - Change MBB pointer to MF reference - Rename Delegate - Tidy constructor parameter names --- llvm/include/llvm/CodeGen/MachineBasicBlock.h | 1 - llvm/lib/CodeGen/MachineBasicBlock.cpp | 17 ++++++++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h index 181156686c908..15c4fcd8399c1 100644 --- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h +++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h @@ -14,7 +14,6 @@ #define LLVM_CODEGEN_MACHINEBASICBLOCK_H #include "llvm/ADT/GraphTraits.h" -#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SparseBitVector.h" #include "llvm/ADT/ilist.h" #include "llvm/ADT/iterator_range.h" diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp index 11b85251d814f..14d9bb292ddf2 100644 --- a/llvm/lib/CodeGen/MachineBasicBlock.cpp +++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp @@ -1097,19 +1097,18 @@ static bool jumpTableHasOtherUses(const MachineFunction &MF, return false; } -class MBBSplitCriticalEdgeDelegate : public MachineFunction::Delegate { +class SlotIndexUpdateDelegate : public MachineFunction::Delegate { private: - MachineFunction *MF; + MachineFunction &MF; SlotIndexes *Indexes; public: - MBBSplitCriticalEdgeDelegate(MachineBasicBlock *MBB, - SlotIndexes *IndexesToUpdate) - : MF(MBB->getParent()), Indexes(IndexesToUpdate) { - MF->setDelegate(this); + SlotIndexUpdateDelegate(MachineFunction &MF, SlotIndexes *Indexes) + : MF(MF), Indexes(Indexes) { + MF.setDelegate(this); } - ~MBBSplitCriticalEdgeDelegate() { MF->resetDelegate(this); } + ~SlotIndexUpdateDelegate() { MF.resetDelegate(this); } void MF_HandleInsertion(MachineInstr &MI) override { if (Indexes) @@ -1201,14 +1200,14 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge( PrevFallthrough = NMBB; if (!ChangedIndirectJump) { - MBBSplitCriticalEdgeDelegate SlotUpdater(this, Indexes); + SlotIndexUpdateDelegate SlotUpdater(*MF, Indexes); updateTerminator(PrevFallthrough); } // Insert unconditional "jump Succ" instruction in NMBB if necessary. NMBB->addSuccessor(Succ); if (!NMBB->isLayoutSuccessor(Succ)) { - MBBSplitCriticalEdgeDelegate SlotUpdater(NMBB, Indexes); + SlotIndexUpdateDelegate SlotUpdater(*MF, Indexes); SmallVector Cond; const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo(); TII->insertBranch(*NMBB, Succ, nullptr, Cond, DL);