[llvm] [MachineBasicBlock] Fix use after free in SplitCriticalEdge (PR #68786)
Carl Ritson via llvm-commits
llvm-commits at lists.llvm.org
Thu Oct 12 22:05:02 PDT 2023
https://github.com/perlfu updated https://github.com/llvm/llvm-project/pull/68786
>From 2f753b4ae96604ba76202b5cf011be1ec0adf49f Mon Sep 17 00:00:00 2001
From: Carl Ritson <carl.ritson at amd.com>
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 15c4fcd8399c181..4f9df77141debc8 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<MachineInstr, ilist_sentinel_tracking<true>>;
@@ -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<Delegate *, 1> 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 7d3d8b6fba1b7df..38dffc031dcb66d 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -156,12 +156,14 @@ void ilist_traits<MachineInstr>::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<MachineInstr>::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<SparseBitVector<>> *LiveInSets) {
@@ -1170,51 +1197,23 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
ReplaceUsesOfBlockWith(Succ, NMBB);
- // If updateTerminator() removes instructions, we need to remove them from
- // SlotIndexes.
- SmallVector<MachineInstr*, 4> 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<MachineInstr*, 4> 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<MachineOperand, 4> 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 <carl.ritson at amd.com>
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 4f9df77141debc8..181156686c90819 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<MachineInstr, ilist_sentinel_tracking<true>>;
@@ -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<Delegate *, 1> 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 38dffc031dcb66d..11b85251d814fde 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -156,14 +156,12 @@ void ilist_traits<MachineInstr>::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<MachineInstr>::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 <carl.ritson at amd.com>
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 181156686c90819..15c4fcd8399c181 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 11b85251d814fde..14d9bb292ddf2e8 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<MachineOperand, 4> Cond;
const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo();
TII->insertBranch(*NMBB, Succ, nullptr, Cond, DL);
More information about the llvm-commits
mailing list