[llvm] c5432d3 - [SandboxIR][Tracker] Track eraseFromParent() (#99431)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 18 11:06:36 PDT 2024
Author: vporpo
Date: 2024-07-18T11:06:33-07:00
New Revision: c5432d31cb339262451215f6cf9c356a514a1770
URL: https://github.com/llvm/llvm-project/commit/c5432d31cb339262451215f6cf9c356a514a1770
DIFF: https://github.com/llvm/llvm-project/commit/c5432d31cb339262451215f6cf9c356a514a1770.diff
LOG: [SandboxIR][Tracker] Track eraseFromParent() (#99431)
This patch adds tracking support for Instruction::eraseFromParent(). The
Instruction is not actually being erased, but instead it is detached
from the instruction list and drops its Use edges. The original
instruction position and Use edges are saved in the `EraseFromParent`
change object, and are being used during `revert()` to restore the
original state.
Added:
Modified:
llvm/include/llvm/SandboxIR/SandboxIR.h
llvm/include/llvm/SandboxIR/Tracker.h
llvm/lib/SandboxIR/SandboxIR.cpp
llvm/lib/SandboxIR/Tracker.cpp
llvm/unittests/SandboxIR/TrackerTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index c5d59ba47ca31..a9f0177eb9338 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -493,6 +493,7 @@ class Instruction : public sandboxir::User {
/// \Returns the LLVM IR Instructions that this SandboxIR maps to in program
/// order.
virtual SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const = 0;
+ friend class EraseFromParent; // For getLLVMInstrs().
public:
static const char *getOpcodeName(Opcode Opc);
@@ -658,6 +659,7 @@ class Context {
friend void Instruction::eraseFromParent(); // For detach().
/// Take ownership of VPtr and store it in `LLVMValueToValueMap`.
Value *registerValue(std::unique_ptr<Value> &&VPtr);
+ friend class EraseFromParent; // For registerValue().
/// This is the actual function that creates sandboxir values for \p V,
/// and among others handles all instruction types.
Value *getOrCreateValueInternal(llvm::Value *V, llvm::User *U = nullptr);
@@ -682,7 +684,7 @@ class Context {
friend class BasicBlock; // For getOrCreateValue().
public:
- Context(LLVMContext &LLVMCtx) : LLVMCtx(LLVMCtx) {}
+ Context(LLVMContext &LLVMCtx) : LLVMCtx(LLVMCtx), IRTracker(*this) {}
Tracker &getTracker() { return IRTracker; }
/// Convenience function for `getTracker().save()`
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 2d0904f5665b1..a91b9f178b8aa 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -40,6 +40,7 @@
#ifndef LLVM_SANDBOXIR_TRACKER_H
#define LLVM_SANDBOXIR_TRACKER_H
+#include "llvm/ADT/PointerUnion.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
@@ -99,6 +100,41 @@ class UseSet : public IRChangeBase {
#endif
};
+class EraseFromParent : public IRChangeBase {
+ /// Contains all the data we need to restore an "erased" (i.e., detached)
+ /// instruction: the instruction itself and its operands in order.
+ struct InstrAndOperands {
+ /// The operands that got dropped.
+ SmallVector<llvm::Value *> Operands;
+ /// The instruction that got "erased".
+ llvm::Instruction *LLVMI;
+ };
+ /// The instruction data is in reverse program order, which helps create the
+ /// original program order during revert().
+ SmallVector<InstrAndOperands> InstrData;
+ /// This is either the next Instruction in the stream, or the parent
+ /// BasicBlock if at the end of the BB.
+ PointerUnion<llvm::Instruction *, llvm::BasicBlock *> NextLLVMIOrBB;
+ /// We take ownership of the "erased" instruction.
+ std::unique_ptr<sandboxir::Value> ErasedIPtr;
+
+public:
+ EraseFromParent(std::unique_ptr<sandboxir::Value> &&IPtr, Tracker &Tracker);
+ void revert() final;
+ void accept() final;
+#ifndef NDEBUG
+ void dump(raw_ostream &OS) const final {
+ dumpCommon(OS);
+ OS << "EraseFromParent";
+ }
+ LLVM_DUMP_METHOD void dump() const final;
+ friend raw_ostream &operator<<(raw_ostream &OS, const EraseFromParent &C) {
+ C.dump(OS);
+ return OS;
+ }
+#endif
+};
+
/// The tracker collects all the change objects and implements the main API for
/// saving / reverting / accepting.
class Tracker {
@@ -116,6 +152,7 @@ class Tracker {
#endif
/// The current state of the tracker.
TrackerState State = TrackerState::Disabled;
+ Context &Ctx;
public:
#ifndef NDEBUG
@@ -124,8 +161,9 @@ class Tracker {
bool InMiddleOfCreatingChange = false;
#endif // NDEBUG
- Tracker() = default;
+ explicit Tracker(Context &Ctx) : Ctx(Ctx) {}
~Tracker();
+ Context &getContext() const { return Ctx; }
/// Record \p Change and take ownership. This is the main function used to
/// track Sandbox IR changes.
void track(std::unique_ptr<IRChangeBase> &&Change);
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 944869a37989c..c6daf1a586546 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -341,11 +341,26 @@ void Instruction::removeFromParent() {
void Instruction::eraseFromParent() {
assert(users().empty() && "Still connected to users, can't erase!");
std::unique_ptr<Value> Detached = Ctx.detach(this);
- // We don't have Tracking yet, so just erase the LLVM IR instructions.
- // Erase in reverse to avoid erasing nstructions with attached uses.
- auto Instrs = getLLVMInstrs();
- for (llvm::Instruction *I : reverse(Instrs))
- I->eraseFromParent();
+ auto LLVMInstrs = getLLVMInstrs();
+
+ auto &Tracker = Ctx.getTracker();
+ if (Tracker.isTracking()) {
+ Tracker.track(
+ std::make_unique<EraseFromParent>(std::move(Detached), Tracker));
+ // We don't actually delete the IR instruction, because then it would be
+ // impossible to bring it back from the dead at the same memory location.
+ // Instead we remove it from its BB and track its current location.
+ for (llvm::Instruction *I : LLVMInstrs)
+ I->removeFromParent();
+ // TODO: Multi-instructions need special treatment because some of the
+ // references are internal to the instruction.
+ for (llvm::Instruction *I : LLVMInstrs)
+ I->dropAllReferences();
+ } else {
+ // Erase in reverse to avoid erasing nstructions with attached uses.
+ for (llvm::Instruction *I : reverse(LLVMInstrs))
+ I->eraseFromParent();
+ }
}
void Instruction::moveBefore(BasicBlock &BB, const BBIterator &WhereIt) {
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 1182f5c55d10b..2336a0067abbd 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -41,6 +41,65 @@ Tracker::~Tracker() {
assert(Changes.empty() && "You must accept or revert changes!");
}
+EraseFromParent::EraseFromParent(std::unique_ptr<sandboxir::Value> &&ErasedIPtr,
+ Tracker &Tracker)
+ : IRChangeBase(Tracker), ErasedIPtr(std::move(ErasedIPtr)) {
+ auto *I = cast<Instruction>(this->ErasedIPtr.get());
+ auto LLVMInstrs = I->getLLVMInstrs();
+ // Iterate in reverse program order.
+ for (auto *LLVMI : reverse(LLVMInstrs)) {
+ SmallVector<llvm::Value *> Operands;
+ Operands.reserve(LLVMI->getNumOperands());
+ for (auto [OpNum, Use] : enumerate(LLVMI->operands()))
+ Operands.push_back(Use.get());
+ InstrData.push_back({Operands, LLVMI});
+ }
+ assert(is_sorted(InstrData,
+ [](const auto &D0, const auto &D1) {
+ return D0.LLVMI->comesBefore(D1.LLVMI);
+ }) &&
+ "Expected reverse program order!");
+ auto *BotLLVMI = cast<llvm::Instruction>(I->Val);
+ if (BotLLVMI->getNextNode() != nullptr)
+ NextLLVMIOrBB = BotLLVMI->getNextNode();
+ else
+ NextLLVMIOrBB = BotLLVMI->getParent();
+}
+
+void EraseFromParent::accept() {
+ for (const auto &IData : InstrData)
+ IData.LLVMI->deleteValue();
+}
+
+void EraseFromParent::revert() {
+ // Place the bottom-most instruction first.
+ auto [Operands, BotLLVMI] = InstrData[0];
+ if (auto *NextLLVMI = NextLLVMIOrBB.dyn_cast<llvm::Instruction *>()) {
+ BotLLVMI->insertBefore(NextLLVMI);
+ } else {
+ auto *LLVMBB = NextLLVMIOrBB.get<llvm::BasicBlock *>();
+ BotLLVMI->insertInto(LLVMBB, LLVMBB->end());
+ }
+ for (auto [OpNum, Op] : enumerate(Operands))
+ BotLLVMI->setOperand(OpNum, Op);
+
+ // Go over the rest of the instructions and stack them on top.
+ for (auto [Operands, LLVMI] : drop_begin(InstrData)) {
+ LLVMI->insertBefore(BotLLVMI);
+ for (auto [OpNum, Op] : enumerate(Operands))
+ LLVMI->setOperand(OpNum, Op);
+ BotLLVMI = LLVMI;
+ }
+ Parent.getContext().registerValue(std::move(ErasedIPtr));
+}
+
+#ifndef NDEBUG
+void EraseFromParent::dump() const {
+ dump(dbgs());
+ dbgs() << "\n";
+}
+#endif
+
void Tracker::track(std::unique_ptr<IRChangeBase> &&Change) {
assert(State == TrackerState::Record && "The tracker should be tracking!");
Changes.push_back(std::move(Change));
diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index f090dc521c32b..4e5dccf3a231f 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -146,3 +146,55 @@ define void @foo(ptr %ptr) {
Ctx.accept();
EXPECT_EQ(St0->getOperand(0), Ld1);
}
+
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, EraseFromParent) {
+ parseIR(C, R"IR(
+define void @foo(i32 %v1) {
+ %add0 = add i32 %v1, %v1
+ %add1 = add i32 %add0, %v1
+ ret void
+}
+)IR");
+ Function &LLVMF = *M->getFunction("foo");
+ sandboxir::Context Ctx(C);
+
+ auto *F = Ctx.createFunction(&LLVMF);
+ auto *BB = &*F->begin();
+ auto It = BB->begin();
+ sandboxir::Instruction *Add0 = &*It++;
+ sandboxir::Instruction *Add1 = &*It++;
+ sandboxir::Instruction *Ret = &*It++;
+
+ Ctx.save();
+ // Check erase.
+ Add1->eraseFromParent();
+ It = BB->begin();
+ EXPECT_EQ(&*It++, Add0);
+ EXPECT_EQ(&*It++, Ret);
+ EXPECT_EQ(It, BB->end());
+ EXPECT_EQ(Add0->getNumUses(), 0u);
+
+ // Check revert().
+ Ctx.revert();
+ It = BB->begin();
+ EXPECT_EQ(&*It++, Add0);
+ EXPECT_EQ(&*It++, Add1);
+ EXPECT_EQ(&*It++, Ret);
+ EXPECT_EQ(It, BB->end());
+ EXPECT_EQ(Add1->getOperand(0), Add0);
+
+ // Same for the last instruction in the block.
+ Ctx.save();
+ Ret->eraseFromParent();
+ It = BB->begin();
+ EXPECT_EQ(&*It++, Add0);
+ EXPECT_EQ(&*It++, Add1);
+ EXPECT_EQ(It, BB->end());
+ Ctx.revert();
+ It = BB->begin();
+ EXPECT_EQ(&*It++, Add0);
+ EXPECT_EQ(&*It++, Add1);
+ EXPECT_EQ(&*It++, Ret);
+ EXPECT_EQ(It, BB->end());
+}
More information about the llvm-commits
mailing list