[llvm] 0959e58 - [SandboxIR][Tracker] Track creation of instructions (#102013)

[via llvm-commits]

Author: vporpo
Date: 2024-08-05T10:23:36-07:00
New Revision: 0959e58fbe08a5333333100b367587cc04673eb8

URL: https://github.com/llvm/llvm-project/commit/0959e58fbe08a5333333100b367587cc04673eb8
DIFF: https://github.com/llvm/llvm-project/commit/0959e58fbe08a5333333100b367587cc04673eb8.diff

LOG: [SandboxIR][Tracker] Track creation of instructions (#102013)

This patch adds tracking support for the creation of SandboxIR
instructions. This allows us to undo the creation of instructions during
revert().

Please note that we don't allow creation of detached instructions,
meaning that the instructions need to be inserted into a block upon
creation. This is why the tracker class combines creation and insertion.

Added: 
    

Modified: 
    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/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index f660ae94c8f7f..3c491bf0b3f58 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -328,6 +328,23 @@ class InsertIntoBB final : public IRChangeBase {
 #endif // NDEBUG
 };
 
+class CreateAndInsertInst final : public IRChangeBase {
+  Instruction *NewI = nullptr;
+
+public:
+  CreateAndInsertInst(Instruction *NewI, Tracker &Tracker)
+      : IRChangeBase(Tracker), NewI(NewI) {}
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "CreateAndInsertInst";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
 /// The tracker collects all the change objects and implements the main API for
 /// saving / reverting / accepting.
 class Tracker {

diff  --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 00abee48821ff..235c4574ead13 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -1359,6 +1359,16 @@ std::unique_ptr<Value> Context::detach(Value *V) {
 Value *Context::registerValue(std::unique_ptr<Value> &&VPtr) {
   assert(VPtr->getSubclassID() != Value::ClassID::User &&
          "Can't register a user!");
+
+  // Track creation of instructions.
+  // Please note that we don't allow the creation of detached instructions,
+  // meaning that the instructions need to be inserted into a block upon
+  // creation. This is why the tracker class combines creation and insertion.
+  auto &Tracker = getTracker();
+  if (Tracker.isTracking())
+    if (auto *I = dyn_cast<Instruction>(VPtr.get()))
+      Tracker.track(std::make_unique<CreateAndInsertInst>(I, Tracker));
+
   Value *V = VPtr.get();
   [[maybe_unused]] auto Pair =
       LLVMValueToValueMap.insert({VPtr->Val, std::move(VPtr)});

diff  --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 4ab24123cafc5..aa18c21cd03d6 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -300,6 +300,15 @@ void InsertIntoBB::dump() const {
 }
 #endif
 
+void CreateAndInsertInst::revert() { NewI->eraseFromParent(); }
+
+#ifndef NDEBUG
+void CreateAndInsertInst::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 76cee1105cbf0..e61258392a388 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -505,6 +505,42 @@ define void @foo(i32 %arg) {
   Add0->insertBefore(Ret);
 }
 
+// TODO: Test multi-instruction patterns.
+TEST_F(TrackerTest, CreateAndInsertInst) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %ld = load i8, ptr %ptr, align 64
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *Ptr = F->getArg(0);
+  auto *BB = &*F->begin();
+  auto It = BB->begin();
+  auto *Ld = cast<sandboxir::LoadInst>(&*It++);
+  auto *Ret = &*It++;
+
+  Ctx.save();
+  // Check create(InsertBefore) with tracking enabled.
+  sandboxir::LoadInst *NewLd =
+      sandboxir::LoadInst::create(Ld->getType(), Ptr, Align(8),
+                                  /*InsertBefore=*/Ld, Ctx, "NewLd");
+  It = BB->begin();
+  EXPECT_EQ(&*It++, NewLd);
+  EXPECT_EQ(&*It++, Ld);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+  // Check revert().
+  Ctx.revert();
+  It = BB->begin();
+  EXPECT_EQ(&*It++, Ld);
+  EXPECT_EQ(&*It++, Ret);
+  EXPECT_EQ(It, BB->end());
+}
+
 TEST_F(TrackerTest, CallBaseSetters) {
   parseIR(C, R"IR(
 declare void @bar1(i8)