[llvm] [PeepholeOptimizer] Fix use-after-free of recycled MachineInstr in LocalMIs (PR #201285)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Jun 3 00:54:47 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-x86
Author: lijinpei-amd
<details>
<summary>Changes</summary>
Peephole scans a block keeping LocalMIs = the instructions seen so far, and the following call sequence corrupts it:
run()
optimizeCmpInstr(RedundantCmp)
TII->optimizeCompareInstr(RedundantCmp)
RedundantCmp->eraseFromParent() // frees RedundantCmp's storage
foldLoadInto(FlagProducer) // fold the load into the reused
TII->optimizeLoadInstr(FlagProducer) // flag producer
BuildMI(...) -> FoldedCmp // recycler hands back RedundantCmp's
// storage: FoldedCmp == old ptr
LocalMIs.insert(FoldedCmp)
LocalMIs.erase(RedundantCmp) // stale ptr aliases FoldedCmp, so this
// drops the live FoldedCmp from the set
...
optimizeExtInstr(Ext) // Ext extends FoldedCmp's source reg
!LocalMIs.count(FoldedCmp) // absent -> treated as "after Ext" ->
// rewrite FoldedCmp's source operand to
// a sub-reg of Ext's result: a use of
// Ext before its def
LiveIntervals -> "Use not jointly dominated by defs"
Maintain LocalMIs through the MachineFunction::Delegate callbacks instead. This relies on list edits reaching the delegate; splice/moveBefore would bypass it, but no peephole-reachable transform repositions that way.
Fixes #<!-- -->199237.
---
Full diff: https://github.com/llvm/llvm-project/pull/201285.diff
2 Files Affected:
- (modified) llvm/lib/CodeGen/PeepholeOptimizer.cpp (+37-41)
- (added) llvm/test/CodeGen/X86/peephole-ext-fold-recycle.mir (+39)
``````````diff
diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index ec8a0336a2105..50da0469092aa 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -91,6 +91,7 @@
#include "llvm/Pass.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
+#include "llvm/Support/SaveAndRestore.h"
#include "llvm/Support/raw_ostream.h"
#include <cassert>
#include <cstdint>
@@ -445,17 +446,16 @@ class PeepholeOptimizer : private MachineFunction::Delegate {
private:
bool optimizeCmpInstr(MachineInstr &MI, MachineFunction &MF,
- SmallPtrSet<MachineInstr *, 16> &LocalMIs);
+ const SmallPtrSet<MachineInstr *, 16> &LocalMIs);
bool optimizeExtInstr(MachineInstr &MI, MachineBasicBlock &MBB,
- SmallPtrSetImpl<MachineInstr *> &LocalMIs);
+ const SmallPtrSetImpl<MachineInstr *> &LocalMIs);
bool optimizeSelect(MachineInstr &MI,
SmallPtrSetImpl<MachineInstr *> &LocalMIs);
bool optimizeCondBranch(MachineInstr &MI);
bool optimizeCoalescableCopyImpl(Rewriter &&CpyRewriter);
bool optimizeCoalescableCopy(MachineInstr &MI);
- bool optimizeUncoalescableCopy(MachineInstr &MI,
- SmallPtrSetImpl<MachineInstr *> &LocalMIs);
+ bool optimizeUncoalescableCopy(MachineInstr &MI);
bool optimizeRecurrence(MachineInstr &PHI);
bool findNextSource(const TargetRegisterClass *DefRC, unsigned DefSubReg,
RegSubRegPair RegSubReg, RewriteMapTy &RewriteMap);
@@ -496,11 +496,11 @@ class PeepholeOptimizer : private MachineFunction::Delegate {
SmallSet<Register, 16> &FoldAsLoadDefCandidates);
/// Try to fold the load defined by \p FoldReg into \p MI using
- /// TII->optimizeLoadInstr. On success, updates \p LocalMIs, erases the old
- /// instructions, and returns the replacement; returns nullptr otherwise.
+ /// TII->optimizeLoadInstr. On success, erases the old instructions and
+ /// returns the replacement; returns nullptr otherwise. LocalMIs is kept up to
+ /// date by the delegate.
MachineInstr *foldLoadInto(MachineFunction &MF, MachineInstr &MI,
- Register FoldReg,
- SmallPtrSet<MachineInstr *, 16> &LocalMIs);
+ Register FoldReg);
/// Check whether \p MI is understood by the register coalescer
/// but may require some rewriting.
@@ -527,8 +527,14 @@ class PeepholeOptimizer : private MachineFunction::Delegate {
// holds any physreg which requires def tracking.
DenseMap<RegSubRegPair, MachineInstr *> CopySrcMIs;
- // MachineFunction::Delegate implementation. Used to maintain CopySrcMIs.
- void MF_HandleInsertion(MachineInstr &MI) override {}
+ // Points at the LocalMIs set of the block run() is currently scanning, so the
+ // delegate callbacks keep it in sync.
+ SmallPtrSetImpl<MachineInstr *> *CurLocalMIs = nullptr;
+
+ void MF_HandleInsertion(MachineInstr &MI) override {
+ if (CurLocalMIs)
+ CurLocalMIs->insert(&MI);
+ }
bool getCopySrc(MachineInstr &MI, RegSubRegPair &SrcPair) {
if (!MI.isCopy())
@@ -555,7 +561,11 @@ class PeepholeOptimizer : private MachineFunction::Delegate {
CopySrcMIs.erase(It);
}
- void MF_HandleRemoval(MachineInstr &MI) override { deleteChangedCopy(MI); }
+ void MF_HandleRemoval(MachineInstr &MI) override {
+ if (CurLocalMIs)
+ CurLocalMIs->erase(&MI);
+ deleteChangedCopy(MI);
+ }
void MF_HandleChangeDesc(MachineInstr &MI, const MCInstrDesc &TID) override {
deleteChangedCopy(MI);
@@ -782,7 +792,7 @@ INITIALIZE_PASS_END(PeepholeOptimizerLegacy, DEBUG_TYPE,
/// debug uses.
bool PeepholeOptimizer::optimizeExtInstr(
MachineInstr &MI, MachineBasicBlock &MBB,
- SmallPtrSetImpl<MachineInstr *> &LocalMIs) {
+ const SmallPtrSetImpl<MachineInstr *> &LocalMIs) {
Register SrcReg, DstReg;
unsigned SubIdx;
if (!TII->isCoalescableExtInstr(MI, SrcReg, DstReg, SubIdx))
@@ -946,7 +956,7 @@ bool PeepholeOptimizer::optimizeExtInstr(
/// previous instruction.
bool PeepholeOptimizer::optimizeCmpInstr(
MachineInstr &MI, MachineFunction &MF,
- SmallPtrSet<MachineInstr *, 16> &LocalMIs) {
+ const SmallPtrSet<MachineInstr *, 16> &LocalMIs) {
// If this instruction is a comparison against zero and isn't comparing a
// physical register, we can try to optimize it.
Register SrcReg, SrcReg2;
@@ -975,7 +985,7 @@ bool PeepholeOptimizer::optimizeCmpInstr(
make_range(std::next(LoadMI->getIterator()),
FlagProducer->getIterator()),
[](const MachineInstr &I) { return I.isLoadFoldBarrier(); }))
- foldLoadInto(MF, *FlagProducer, SrcReg, LocalMIs);
+ foldLoadInto(MF, *FlagProducer, SrcReg);
}
return true;
@@ -1348,9 +1358,7 @@ MachineInstr &PeepholeOptimizer::rewriteSource(MachineInstr &CopyLike,
/// \pre isUncoalescableCopy(*MI) is true.
/// \return True, when \p MI has been optimized. In that case, \p MI has
/// been removed from its parent.
-/// All COPY instructions created, are inserted in \p LocalMIs.
-bool PeepholeOptimizer::optimizeUncoalescableCopy(
- MachineInstr &MI, SmallPtrSetImpl<MachineInstr *> &LocalMIs) {
+bool PeepholeOptimizer::optimizeUncoalescableCopy(MachineInstr &MI) {
assert(isUncoalescableCopy(MI) && "Invalid argument");
UncoalescableRewriter CpyRewriter(MI);
@@ -1381,11 +1389,9 @@ bool PeepholeOptimizer::optimizeUncoalescableCopy(
}
// The change is possible for all defs, do it.
- for (const RegSubRegPair &Def : RewritePairs) {
+ for (const RegSubRegPair &Def : RewritePairs)
// Rewrite the "copy" in a way the register coalescer understands.
- MachineInstr &NewCopy = rewriteSource(MI, Def, RewriteMap);
- LocalMIs.insert(&NewCopy);
- }
+ rewriteSource(MI, Def, RewriteMap);
// MI is now dead.
LLVM_DEBUG(dbgs() << "Deleting uncoalescable copy: " << MI);
@@ -1417,10 +1423,9 @@ bool PeepholeOptimizer::isLoadFoldable(
return false;
}
-MachineInstr *
-PeepholeOptimizer::foldLoadInto(MachineFunction &MF, MachineInstr &MI,
- Register FoldReg,
- SmallPtrSet<MachineInstr *, 16> &LocalMIs) {
+MachineInstr *PeepholeOptimizer::foldLoadInto(MachineFunction &MF,
+ MachineInstr &MI,
+ Register FoldReg) {
Register Reg = FoldReg;
MachineInstr *DefMI = nullptr;
MachineInstr *CopyMI = nullptr;
@@ -1428,11 +1433,6 @@ PeepholeOptimizer::foldLoadInto(MachineFunction &MF, MachineInstr &MI,
if (!FoldMI)
return nullptr;
LLVM_DEBUG(dbgs() << "Replacing: " << MI << " With: " << *FoldMI);
- LocalMIs.erase(&MI);
- LocalMIs.erase(DefMI);
- LocalMIs.insert(FoldMI);
- if (CopyMI)
- LocalMIs.insert(CopyMI);
if (MI.shouldUpdateAdditionalCallInfo())
MF.moveAdditionalCallInfo(&MI, FoldMI);
MI.eraseFromParent();
@@ -1761,9 +1761,11 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
// "given a pointer to an MI in the current BB, is it located before or
// after the current instruction".
// To perform this, the following set keeps track of the MIs already seen
- // during the scan, if a MI is not in the set, it is assumed to be located
- // after. Newly created MIs have to be inserted in the set as well.
+ // during the scan; if a MI is not in the set, it is assumed to be located
+ // after.
SmallPtrSet<MachineInstr *, 16> LocalMIs;
+ SaveAndRestore<SmallPtrSetImpl<MachineInstr *> *> RestoreCurLocalMIs(
+ CurLocalMIs, &LocalMIs);
SmallSet<Register, 4> ImmDefRegs;
DenseMap<Register, MachineInstr *> ImmDefMIs;
SmallSet<Register, 16> FoldAsLoadDefCandidates;
@@ -1843,16 +1845,13 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
}
if (MI->isCompare() && optimizeCmpInstr(*MI, MF, LocalMIs)) {
- LocalMIs.erase(MI);
Changed = true;
continue;
}
- if ((isUncoalescableCopy(*MI) &&
- optimizeUncoalescableCopy(*MI, LocalMIs)) ||
+ if ((isUncoalescableCopy(*MI) && optimizeUncoalescableCopy(*MI)) ||
(MI->isSelect() && optimizeSelect(*MI, LocalMIs))) {
// MI is deleted.
- LocalMIs.erase(MI);
Changed = true;
continue;
}
@@ -1870,7 +1869,6 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
if (MI->isCopy() && (foldRedundantCopy(*MI) ||
foldRedundantNAPhysCopy(*MI, NAPhysToVirtMIs))) {
- LocalMIs.erase(MI);
LLVM_DEBUG(dbgs() << "Deleting redundant copy: " << *MI << "\n");
MI->eraseFromParent();
Changed = true;
@@ -1889,10 +1887,8 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
if (SeenMoveImm) {
bool Deleted;
Changed |= foldImmediate(*MI, ImmDefRegs, ImmDefMIs, Deleted);
- if (Deleted) {
- LocalMIs.erase(MI);
+ if (Deleted)
continue;
- }
}
}
@@ -1918,7 +1914,7 @@ bool PeepholeOptimizer::run(MachineFunction &MF) {
// optimizeCmpInstr can enable folding by converting SUB to CMP.
Register FoldedReg = FoldAsLoadDefReg;
if (MachineInstr *FoldMI =
- foldLoadInto(MF, *MI, FoldAsLoadDefReg, LocalMIs)) {
+ foldLoadInto(MF, *MI, FoldAsLoadDefReg)) {
FoldAsLoadDefCandidates.erase(FoldedReg);
// MI is replaced with FoldMI so we can continue trying to fold
Changed = true;
diff --git a/llvm/test/CodeGen/X86/peephole-ext-fold-recycle.mir b/llvm/test/CodeGen/X86/peephole-ext-fold-recycle.mir
new file mode 100644
index 0000000000000..964e948773b98
--- /dev/null
+++ b/llvm/test/CodeGen/X86/peephole-ext-fold-recycle.mir
@@ -0,0 +1,39 @@
+# RUN: llc -mtriple=x86_64-- -run-pass=peephole-opt -verify-machineinstrs -o - %s | FileCheck %s
+
+# Regression test for https://github.com/llvm/llvm-project/issues/199237
+
+--- |
+ @g13 = global i16 0
+ @g14 = global i16 0
+ @g20 = global i8 0
+ @g29 = global i8 0
+ define void @f10() { ret void }
+...
+---
+name: f10
+tracksRegLiveness: true
+body: |
+ bb.0:
+ ; CHECK-LABEL: bb.0:
+ ; The first compare must use the g13 value taken from the load (%[[G13]]),
+ ; and there must be no second sub_16bit COPY feeding it (the bug rewrote the
+ ; operand to a COPY of the later-defined MOVSX result).
+ ; CHECK: %[[G13:[0-9]+]]:gr16 = COPY %{{[0-9]+}}.sub_16bit
+ ; CHECK-NOT: = COPY %{{[0-9]+}}.sub_16bit
+ ; CHECK: CMP16rm %[[G13]], $rip, 1, $noreg, @g14, $noreg
+ ; The extension is defined before its result is used (no use-before-def).
+ ; CHECK: %[[EXT:[0-9]+]]:gr32 = MOVSX32rr16 %[[G13]]
+ ; CHECK: CMP32ri %[[EXT]], 538
+ %6:gr32 = MOVZX32rm16 $rip, 1, $noreg, @g13, $noreg :: (load (s16) from @g13)
+ %7:gr16 = COPY %6.sub_16bit
+ %8:gr16 = MOV16rm $rip, 1, $noreg, @g14, $noreg :: (load (s16) from @g14)
+ %9:gr16 = SUB16rr %7, %8, implicit-def $eflags
+ SETCCm $rip, 1, $noreg, @g20, $noreg, 13, implicit $eflags :: (store (s8) into @g20)
+ %10:gr16 = SUB16rr %8, %7, implicit-def $eflags
+ %1:gr8 = SETCCr 13, implicit $eflags
+ %11:gr32 = MOVSX32rr16 %7
+ %12:gr32 = SUB32ri %11, 538, implicit-def $eflags
+ SETCCm $rip, 1, $noreg, @g29, $noreg, 12, implicit $eflags :: (store (s8) into @g29)
+ MOV8mr $rip, 1, $noreg, @g20, $noreg, %1 :: (store (s8) into @g20)
+ RET 0
+...
``````````
</details>
https://github.com/llvm/llvm-project/pull/201285
More information about the llvm-commits
mailing list