[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