Skip to content

[MachineBasicBlock] Fix use after free in SplitCriticalEdge #66188

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

[MachineBasicBlock] Fix use after free in SplitCriticalEdge #66188

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c7b575f
[MachineBasicBlock] Fix use after free in SplitCriticalEdge
perlfu Sep 13, 2023
bec43c3
Address reviewer feedback.
perlfu Sep 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions llvm/include/llvm/CodeGen/MachineBasicBlock.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -688,12 +688,15 @@ class MachineBasicBlock
}

/// Update the terminator instructions in block to account for changes to
/// block layout which may have been made. PreviousLayoutSuccessor should be
/// set to the block which may have been used as fallthrough before the block
/// layout was modified. If the block previously fell through to that block,
/// it may now need a branch. If it previously branched to another block, it
/// may now be able to fallthrough to the current layout successor.
void updateTerminator(MachineBasicBlock *PreviousLayoutSuccessor);
/// block layout which may have been made. \p PreviousLayoutSuccessor should
/// be set to the block which may have been used as fallthrough before the
/// block layout was modified. If the block previously fell through to that
/// block, it may now need a branch. If it previously branched to another
/// block, it may now be able to fallthrough to the current layout successor.
/// SlotIndexes provided by \p Indexes will be updated to reflect any
/// instructions inserted or removed.
void updateTerminator(MachineBasicBlock *PreviousLayoutSuccessor,
SlotIndexes *Indexes = nullptr);

// Machine-CFG mutators

Expand Down
15 changes: 12 additions & 3 deletions llvm/include/llvm/CodeGen/TargetInstrInfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ class ScheduleDAGMI;
class ScheduleHazardRecognizer;
class SDNode;
class SelectionDAG;
class SlotIndexes;
class SMSchedule;
class SwingSchedulerDAG;
class RegScavenger;
Expand Down Expand Up @@ -679,18 +680,24 @@ class TargetInstrInfo : public MCInstrInfo {
/// Remove the branching code at the end of the specific MBB.
/// This is only invoked in cases where analyzeBranch returns success. It
/// returns the number of instructions that were removed.
/// If \p Indexes is non-null, then instructions will also be removed from
/// SlotIndexes.
/// If \p BytesRemoved is non-null, report the change in code size from the
/// removed instructions.
virtual unsigned removeBranch(MachineBasicBlock &MBB,
SlotIndexes *Indexes = nullptr,
int *BytesRemoved = nullptr) const {
llvm_unreachable("Target didn't implement TargetInstrInfo::removeBranch!");
}

/// Insert branch code into the end of the specified MachineBasicBlock. The
/// operands to this method are the same as those returned by analyzeBranch.
/// This is only invoked in cases where analyzeBranch returns success. It
/// returns the number of instructions inserted. If \p BytesAdded is non-null,
/// report the change in code size from the added instructions.
/// returns the number of instructions inserted.
/// If \p Indexes is non-null, then added instructions will be inserted into
/// SlotIndexes.
/// If \p BytesAdded is non-null, report the change in code size from the
/// added instructions.
///
/// It is also invoked by tail merging to add unconditional branches in
/// cases where analyzeBranch doesn't apply because there was no original
Expand All @@ -703,16 +710,18 @@ class TargetInstrInfo : public MCInstrInfo {
MachineBasicBlock *FBB,
ArrayRef<MachineOperand> Cond,
const DebugLoc &DL,
SlotIndexes *Indexes = nullptr,
int *BytesAdded = nullptr) const {
llvm_unreachable("Target didn't implement TargetInstrInfo::insertBranch!");
}

unsigned insertUnconditionalBranch(MachineBasicBlock &MBB,
MachineBasicBlock *DestBB,
const DebugLoc &DL,
SlotIndexes *Indexes = nullptr,
int *BytesAdded = nullptr) const {
return insertBranch(MBB, DestBB, nullptr, ArrayRef<MachineOperand>(), DL,
BytesAdded);
Indexes, BytesAdded);
}

/// Object returned by analyzeLoopForPipelining. Allows software pipelining
Expand Down
8 changes: 5 additions & 3 deletions llvm/lib/CodeGen/BranchRelaxation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,21 +364,23 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
MachineBasicBlock *DestBB) {
unsigned &BBSize = BlockInfo[MBB->getNumber()].Size;
int NewBrSize = 0;
TII->insertUnconditionalBranch(*MBB, DestBB, DL, &NewBrSize);
TII->insertUnconditionalBranch(*MBB, DestBB, DL, /*Indexes=*/nullptr,
&NewBrSize);
BBSize += NewBrSize;
};
auto insertBranch = [&](MachineBasicBlock *MBB, MachineBasicBlock *TBB,
MachineBasicBlock *FBB,
SmallVectorImpl<MachineOperand>& Cond) {
unsigned &BBSize = BlockInfo[MBB->getNumber()].Size;
int NewBrSize = 0;
TII->insertBranch(*MBB, TBB, FBB, Cond, DL, &NewBrSize);
TII->insertBranch(*MBB, TBB, FBB, Cond, DL, /*Indexes=*/nullptr,
&NewBrSize);
BBSize += NewBrSize;
};
auto removeBranch = [&](MachineBasicBlock *MBB) {
unsigned &BBSize = BlockInfo[MBB->getNumber()].Size;
int RemovedSize = 0;
TII->removeBranch(*MBB, &RemovedSize);
TII->removeBranch(*MBB, /*Indexes=*/nullptr, &RemovedSize);
BBSize -= RemovedSize;
};

Expand Down
68 changes: 20 additions & 48 deletions llvm/lib/CodeGen/MachineBasicBlock.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,7 +681,7 @@ static int findJumpTableIndex(const MachineBasicBlock &MBB) {
}

void MachineBasicBlock::updateTerminator(
MachineBasicBlock *PreviousLayoutSuccessor) {
MachineBasicBlock *PreviousLayoutSuccessor, SlotIndexes *Indexes) {
LLVM_DEBUG(dbgs() << "Updating terminators on " << printMBBReference(*this)
<< "\n");

Expand All @@ -693,15 +693,15 @@ void MachineBasicBlock::updateTerminator(
MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
SmallVector<MachineOperand, 4> Cond;
DebugLoc DL = findBranchDebugLoc();
bool B = TII->analyzeBranch(*this, TBB, FBB, Cond);
bool B = TII->analyzeBranch(*this, TBB, FBB, Cond, /*AllowModify=*/false);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is NFC, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is NFC, but I felt it was appropriate to add an explicit configuration of AllowModify=false here as SlotIndexes will be invalidated if analyzeBranch modifies the instructions.

(void) B;
assert(!B && "UpdateTerminators requires analyzable predecessors!");
if (Cond.empty()) {
if (TBB) {
// The block has an unconditional branch. If its successor is now its
// layout successor, delete the branch.
if (isLayoutSuccessor(TBB))
TII->removeBranch(*this);
TII->removeBranch(*this, Indexes);
} else {
// The block has an unconditional fallthrough, or the end of the block is
// unreachable.
Expand All @@ -717,7 +717,8 @@ void MachineBasicBlock::updateTerminator(
// If the unconditional successor block is not the current layout
// successor, insert a branch to jump to it.
if (!isLayoutSuccessor(PreviousLayoutSuccessor))
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL,
Indexes);
}
return;
}
Expand All @@ -729,11 +730,11 @@ void MachineBasicBlock::updateTerminator(
if (isLayoutSuccessor(TBB)) {
if (TII->reverseBranchCondition(Cond))
return;
TII->removeBranch(*this);
TII->insertBranch(*this, FBB, nullptr, Cond, DL);
TII->removeBranch(*this, Indexes);
TII->insertBranch(*this, FBB, nullptr, Cond, DL, Indexes);
} else if (isLayoutSuccessor(FBB)) {
TII->removeBranch(*this);
TII->insertBranch(*this, TBB, nullptr, Cond, DL);
TII->removeBranch(*this, Indexes);
TII->insertBranch(*this, TBB, nullptr, Cond, DL, Indexes);
}
return;
}
Expand All @@ -747,10 +748,10 @@ void MachineBasicBlock::updateTerminator(
// We had a fallthrough to the same basic block as the conditional jump
// targets. Remove the conditional jump, leaving an unconditional
// fallthrough or an unconditional jump.
TII->removeBranch(*this);
TII->removeBranch(*this, Indexes);
if (!isLayoutSuccessor(TBB)) {
Cond.clear();
TII->insertBranch(*this, TBB, nullptr, Cond, DL);
TII->insertBranch(*this, TBB, nullptr, Cond, DL, Indexes);
}
return;
}
Expand All @@ -760,14 +761,16 @@ void MachineBasicBlock::updateTerminator(
if (TII->reverseBranchCondition(Cond)) {
// We can't reverse the condition, add an unconditional branch.
Cond.clear();
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL,
Indexes);
return;
}
TII->removeBranch(*this);
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL);
TII->removeBranch(*this, Indexes);
TII->insertBranch(*this, PreviousLayoutSuccessor, nullptr, Cond, DL,
Indexes);
} else if (!isLayoutSuccessor(PreviousLayoutSuccessor)) {
TII->removeBranch(*this);
TII->insertBranch(*this, TBB, PreviousLayoutSuccessor, Cond, DL);
TII->removeBranch(*this, Indexes);
TII->insertBranch(*this, TBB, PreviousLayoutSuccessor, Cond, DL, Indexes);
}
}

Expand Down Expand Up @@ -1166,51 +1169,20 @@ 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)
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);
}
}
updateTerminator(PrevFallthrough, Indexes);

// Insert unconditional "jump Succ" instruction in NMBB if necessary.
NMBB->addSuccessor(Succ);
if (!NMBB->isLayoutSuccessor(Succ)) {
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);
}
}
TII->insertBranch(*NMBB, Succ, nullptr, Cond, DL, Indexes);
}

// Fix PHI nodes in Succ so they refer to NMBB instead of this.
Expand Down
44 changes: 32 additions & 12 deletions llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@
//
//===----------------------------------------------------------------------===//

#include "AArch64ExpandImm.h"
#include "AArch64InstrInfo.h"
#include "AArch64ExpandImm.h"
#include "AArch64FrameLowering.h"
#include "AArch64MachineFunctionInfo.h"
#include "AArch64Subtarget.h"
Expand All @@ -31,6 +31,7 @@
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegisterScavenging.h"
#include "llvm/CodeGen/SlotIndexes.h"
#include "llvm/CodeGen/StackMaps.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/CodeGen/TargetSubtargetInfo.h"
Expand Down Expand Up @@ -534,6 +535,7 @@ bool AArch64InstrInfo::reverseBranchCondition(
}

unsigned AArch64InstrInfo::removeBranch(MachineBasicBlock &MBB,
SlotIndexes *Indexes,
int *BytesRemoved) const {
MachineBasicBlock::iterator I = MBB.getLastNonDebugInstr();
if (I == MBB.end())
Expand All @@ -544,6 +546,8 @@ unsigned AArch64InstrInfo::removeBranch(MachineBasicBlock &MBB,
return 0;

// Remove the branch.
if (Indexes)
Indexes->removeMachineInstrFromMaps(*I);
I->eraseFromParent();

I = MBB.end();
Expand All @@ -561,19 +565,27 @@ unsigned AArch64InstrInfo::removeBranch(MachineBasicBlock &MBB,
}

// Remove the branch.
if (Indexes)
Indexes->removeMachineInstrFromMaps(*I);
I->eraseFromParent();
if (BytesRemoved)
*BytesRemoved = 8;

return 2;
}

void AArch64InstrInfo::instantiateCondBranch(
MachineBasicBlock &MBB, const DebugLoc &DL, MachineBasicBlock *TBB,
ArrayRef<MachineOperand> Cond) const {
void AArch64InstrInfo::instantiateCondBranch(MachineBasicBlock &MBB,
const DebugLoc &DL,
MachineBasicBlock *TBB,
ArrayRef<MachineOperand> Cond,
SlotIndexes *Indexes) const {
if (Cond[0].getImm() != -1) {
// Regular Bcc
BuildMI(&MBB, DL, get(AArch64::Bcc)).addImm(Cond[0].getImm()).addMBB(TBB);
MachineInstr *MI = BuildMI(&MBB, DL, get(AArch64::Bcc))
.addImm(Cond[0].getImm())
.addMBB(TBB);
if (Indexes)
Indexes->insertMachineInstrInMaps(*MI);
} else {
// Folded compare-and-branch
// Note that we use addOperand instead of addReg to keep the flags.
Expand All @@ -582,20 +594,25 @@ void AArch64InstrInfo::instantiateCondBranch(
if (Cond.size() > 3)
MIB.addImm(Cond[3].getImm());
MIB.addMBB(TBB);
if (Indexes)
Indexes->insertMachineInstrInMaps(*MIB);
}
}

unsigned AArch64InstrInfo::insertBranch(
MachineBasicBlock &MBB, MachineBasicBlock *TBB, MachineBasicBlock *FBB,
ArrayRef<MachineOperand> Cond, const DebugLoc &DL, int *BytesAdded) const {
ArrayRef<MachineOperand> Cond, const DebugLoc &DL, SlotIndexes *Indexes,
int *BytesAdded) const {
// Shouldn't be a fall through.
assert(TBB && "insertBranch must not be told to insert a fallthrough");

if (!FBB) {
if (Cond.empty()) // Unconditional branch?
BuildMI(&MBB, DL, get(AArch64::B)).addMBB(TBB);
else
instantiateCondBranch(MBB, DL, TBB, Cond);
if (Cond.empty()) { // Unconditional branch?
MachineInstr *MI = BuildMI(&MBB, DL, get(AArch64::B)).addMBB(TBB);
if (Indexes)
Indexes->insertMachineInstrInMaps(*MI);
} else
instantiateCondBranch(MBB, DL, TBB, Cond, Indexes);

if (BytesAdded)
*BytesAdded = 4;
Expand All @@ -604,8 +621,11 @@ unsigned AArch64InstrInfo::insertBranch(
}

// Two-way conditional branch.
instantiateCondBranch(MBB, DL, TBB, Cond);
BuildMI(&MBB, DL, get(AArch64::B)).addMBB(FBB);
instantiateCondBranch(MBB, DL, TBB, Cond, Indexes);
MachineInstr *MI = BuildMI(&MBB, DL, get(AArch64::B)).addMBB(FBB);

if (Indexes)
Indexes->insertMachineInstrInMaps(*MI);

if (BytesAdded)
*BytesAdded = 8;
Expand Down
7 changes: 4 additions & 3 deletions llvm/lib/Target/AArch64/AArch64InstrInfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,11 +225,11 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
bool analyzeBranchPredicate(MachineBasicBlock &MBB,
MachineBranchPredicate &MBP,
bool AllowModify) const override;
unsigned removeBranch(MachineBasicBlock &MBB,
unsigned removeBranch(MachineBasicBlock &MBB, SlotIndexes *Indexes = nullptr,
int *BytesRemoved = nullptr) const override;
unsigned insertBranch(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
MachineBasicBlock *FBB, ArrayRef<MachineOperand> Cond,
const DebugLoc &DL,
const DebugLoc &DL, SlotIndexes *Indexes = nullptr,
int *BytesAdded = nullptr) const override;
bool
reverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const override;
Expand Down Expand Up @@ -378,7 +378,8 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {

void instantiateCondBranch(MachineBasicBlock &MBB, const DebugLoc &DL,
MachineBasicBlock *TBB,
ArrayRef<MachineOperand> Cond) const;
ArrayRef<MachineOperand> Cond,
SlotIndexes *Indexes) const;
bool substituteCmpToZero(MachineInstr &CmpInstr, unsigned SrcReg,
const MachineRegisterInfo &MRI) const;
bool removeCmpToZeroOrOne(MachineInstr &CmpInstr, unsigned SrcReg,
Expand Down
Loading