[llvm] [SandboxIR] IR Tracker (PR #99238)

via llvm-commits llvm-commits at lists.llvm.org
Wed Jul 17 19:49:49 PDT 2024


https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/99238

>From 83dd6166cc853da5a6eff62c0b5e7f38858858a4 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Fri, 12 Jul 2024 10:24:55 -0700
Subject: [PATCH 1/6] [SandboxIR] IR Tracker

This is the first patch in a series of patches for the IR change tracking
component of SandboxIR.
The tracker collects changes in a vector of `IRChangeBase` objects and
provides a `save()`/`accept()`/`revert()` API.

Each type of IR changing event is captured by a dedicated subclass of
`IRChangeBase`. This patch implements only one of them, that for updating
a `sandboxir::Use` source value, named `UseSet`.
---
 llvm/docs/SandboxIR.md                        |  11 ++
 llvm/include/llvm/SandboxIR/SandboxIR.h       |  12 ++
 .../include/llvm/SandboxIR/SandboxIRTracker.h | 181 ++++++++++++++++++
 llvm/include/llvm/SandboxIR/Use.h             |   1 +
 llvm/lib/SandboxIR/CMakeLists.txt             |   1 +
 llvm/lib/SandboxIR/SandboxIR.cpp              |  26 ++-
 llvm/lib/SandboxIR/SandboxIRTracker.cpp       |  84 ++++++++
 llvm/unittests/SandboxIR/CMakeLists.txt       |   1 +
 .../SandboxIR/SandboxIRTrackerTest.cpp        | 154 +++++++++++++++
 9 files changed, 470 insertions(+), 1 deletion(-)
 create mode 100644 llvm/include/llvm/SandboxIR/SandboxIRTracker.h
 create mode 100644 llvm/lib/SandboxIR/SandboxIRTracker.cpp
 create mode 100644 llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp

diff --git a/llvm/docs/SandboxIR.md b/llvm/docs/SandboxIR.md
index 8f8752f102c76..29f5e5ea9346f 100644
--- a/llvm/docs/SandboxIR.md
+++ b/llvm/docs/SandboxIR.md
@@ -51,3 +51,14 @@ For example, for `sandboxir::User::setOperand(OpIdx, sandboxir::Value *Op)`:
 - We get the corresponding LLVM User: `llvm::User *LLVMU = cast<llvm::User>(Val)`
 - Next we get the corresponding LLVM Operand: `llvm::Value *LLVMOp = Op->Val`
 - Finally we modify `LLVMU`'s operand: `LLVMU->setOperand(OpIdx, LLVMOp)
+
+## IR Change Tracking
+Sandbox IR's state can be saved and restored.
+This is done with the help of the tracker component that is tightly coupled to the public Sandbox IR API functions.
+
+To save the state and enable tracking the user needs to call `sandboxir::Context::save()`.
+From this point on any change made to the Sandbox IR state will automatically create a change object and register it with the tracker, without any intervention from the user.
+The changes are accumulated in a vector within the tracker.
+
+To rollback to the saved state the user needs to call `sandboxir::Context::revert()`.
+Reverting back to the saved state is a matter of going over all the accumulated states in reverse and undoing each individual change.
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 473bd93aea7c1..3f1528ffa0355 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -61,6 +61,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
+#include "llvm/SandboxIR/SandboxIRTracker.h"
 #include "llvm/SandboxIR/Use.h"
 #include "llvm/Support/raw_ostream.h"
 #include <iterator>
@@ -171,6 +172,7 @@ class Value {
 
   friend class Context; // For getting `Val`.
   friend class User;    // For getting `Val`.
+  friend class Use;     // For getting `Val`.
 
   /// All values point to the context.
   Context &Ctx;
@@ -641,6 +643,8 @@ class BasicBlock : public Value {
 class Context {
 protected:
   LLVMContext &LLVMCtx;
+  SandboxIRTracker IRTracker;
+
   /// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
   /// SandboxIR objects.
   DenseMap<llvm::Value *, std::unique_ptr<sandboxir::Value>>
@@ -680,6 +684,14 @@ class Context {
 public:
   Context(LLVMContext &LLVMCtx) : LLVMCtx(LLVMCtx) {}
 
+  SandboxIRTracker &getTracker() { return IRTracker; }
+  /// Convenience function for `getTracker().save()`
+  void save() { IRTracker.save(); }
+  /// Convenience function for `getTracker().revert()`
+  void revert() { IRTracker.revert(); }
+  /// Convenience function for `getTracker().accept()`
+  void accept() { IRTracker.accept(); }
+
   sandboxir::Value *getValue(llvm::Value *V) const;
   const sandboxir::Value *getValue(const llvm::Value *V) const {
     return getValue(const_cast<llvm::Value *>(V));
diff --git a/llvm/include/llvm/SandboxIR/SandboxIRTracker.h b/llvm/include/llvm/SandboxIR/SandboxIRTracker.h
new file mode 100644
index 0000000000000..8a819c578a156
--- /dev/null
+++ b/llvm/include/llvm/SandboxIR/SandboxIRTracker.h
@@ -0,0 +1,181 @@
+//===- SandboxIRTracker.h ---------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is the component of SandboxIR that tracks all changes made to its
+// state, such that we can revert the state when needed.
+//
+// Tracking changes
+// ----------------
+// The user needs to call `SandboxIRTracker::save()` to enable tracking changes
+// made to SandboxIR. From that point on, any change made to SandboxIR, will
+// automatically create a change tracking object and register it with the
+// tracker. IR-change objects are subclasses of `IRChangeBase` and get
+// registered with the `SandboxIRTracker::track()` function. The change objects
+// are saved in the order they are registered with the tracker and are stored in
+// the `SandboxIRTracker::Changes` vector. All of this is done transparently to
+// the user.
+//
+// Reverting changes
+// -----------------
+// Calling `SandboxIRTracker::revert()` will restore the state saved when
+// `SandboxIRTracker::save()` was called. Internally this goes through the
+// change objects in `SandboxIRTracker::Changes` in reverse order, calling their
+// `IRChangeBase::revert()` function one by one.
+//
+// Accepting changes
+// -----------------
+// The user needs to either revert or accept changes before the tracker object
+// is destroyed, or else the tracker destructor will cause a crash.
+// This is the job of `SandboxIRTracker::accept()`. Internally this will go
+// through the change objects in `SandboxIRTracker::Changes` in order, calling
+// `IRChangeBase::accept()`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+#define LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+#include "llvm/SandboxIR/Use.h"
+#include "llvm/Support/Debug.h"
+#include <memory>
+#include <regex>
+
+namespace llvm::sandboxir {
+
+class BasicBlock;
+
+/// Each IR change type has an ID.
+enum class TrackID {
+  UseSet,
+};
+
+#ifndef NDEBUG
+static const char *trackIDToStr(TrackID ID) {
+  switch (ID) {
+  case TrackID::UseSet:
+    return "UseSet";
+  }
+  llvm_unreachable("Unimplemented ID");
+}
+#endif // NDEBUG
+
+class SandboxIRTracker;
+
+/// The base class for IR Change classes.
+class IRChangeBase {
+protected:
+#ifndef NDEBUG
+  unsigned Idx = 0;
+#endif
+  const TrackID ID;
+  SandboxIRTracker &Parent;
+
+public:
+  IRChangeBase(TrackID ID, SandboxIRTracker &Parent);
+  TrackID getTrackID() const { return ID; }
+  /// This runs when changes get reverted.
+  virtual void revert() = 0;
+  /// This runs when changes get accepted.
+  virtual void accept() = 0;
+  virtual ~IRChangeBase() = default;
+#ifndef NDEBUG
+  void dumpCommon(raw_ostream &OS) const {
+    OS << Idx << ". " << trackIDToStr(ID);
+  }
+  virtual void dump(raw_ostream &OS) const = 0;
+  LLVM_DUMP_METHOD virtual void dump() const = 0;
+#endif
+};
+
+/// Change the source Value of a sandboxir::Use.
+class UseSet : public IRChangeBase {
+  Use U;
+  Value *OrigV = nullptr;
+
+public:
+  UseSet(const Use &U, SandboxIRTracker &Tracker)
+      : IRChangeBase(TrackID::UseSet, Tracker), U(U), OrigV(U.get()) {}
+  // For isa<> etc.
+  static bool classof(const IRChangeBase *Other) {
+    return Other->getTrackID() == TrackID::UseSet;
+  }
+  void revert() final { U.set(OrigV); }
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final { dumpCommon(OS); }
+  LLVM_DUMP_METHOD void dump() const final;
+  friend raw_ostream &operator<<(raw_ostream &OS, const UseSet &C) {
+    C.dump(OS);
+    return OS;
+  }
+#endif
+};
+
+/// The tracker collects all the change objects and implements the main API for
+/// saving / reverting / accepting.
+class SandboxIRTracker {
+public:
+  enum class TrackerState {
+    Disabled, ///> Tracking is disabled
+    Record,   ///> Tracking changes
+    Revert,   ///> Undoing changes
+    Accept,   ///> Accepting changes
+  };
+
+private:
+  /// The list of changes that are being tracked.
+  SmallVector<std::unique_ptr<IRChangeBase>> Changes;
+  /// The current state of the tracker.
+  TrackerState State = TrackerState::Disabled;
+
+public:
+#ifndef NDEBUG
+  /// Helps catch bugs where we are creating new change objects while in the
+  /// middle of creating other change objects.
+  bool InMiddleOfCreatingChange = false;
+#endif // NDEBUG
+
+  SandboxIRTracker() = default;
+  ~SandboxIRTracker();
+  /// Record \p Change and take ownership. This is the main function used to
+  /// track Sandbox IR changes.
+  void track(std::unique_ptr<IRChangeBase> &&Change);
+  /// \Returns true if the tracker is recording changes.
+  bool tracking() const { return State == TrackerState::Record; }
+  /// \Returns the current state of the tracker.
+  TrackerState getState() const { return State; }
+  /// Turns on IR tracking.
+  void save();
+  /// Stops tracking and accept changes.
+  void accept();
+  /// Stops tracking and reverts to saved state.
+  void revert();
+  /// \Returns the number of change entries recorded so far.
+  unsigned size() const { return Changes.size(); }
+  /// \Returns true if there are no change entries recorded so far.
+  bool empty() const { return Changes.empty(); }
+
+#ifndef NDEBUG
+  /// \Returns the \p Idx'th change. This is used for testing.
+  IRChangeBase *getChange(unsigned Idx) const { return Changes[Idx].get(); }
+  void dump(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+  friend raw_ostream &operator<<(raw_ostream &OS, const SandboxIRTracker &C) {
+    C.dump(OS);
+    return OS;
+  }
+#endif // NDEBUG
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
diff --git a/llvm/include/llvm/SandboxIR/Use.h b/llvm/include/llvm/SandboxIR/Use.h
index 33afb54c1ff29..d77b4568d0fab 100644
--- a/llvm/include/llvm/SandboxIR/Use.h
+++ b/llvm/include/llvm/SandboxIR/Use.h
@@ -44,6 +44,7 @@ class Use {
 public:
   operator Value *() const { return get(); }
   Value *get() const;
+  void set(Value *V);
   class User *getUser() const { return Usr; }
   unsigned getOperandNo() const;
   Context *getContext() const { return Ctx; }
diff --git a/llvm/lib/SandboxIR/CMakeLists.txt b/llvm/lib/SandboxIR/CMakeLists.txt
index 225eca0cadd1a..74b31fe869aed 100644
--- a/llvm/lib/SandboxIR/CMakeLists.txt
+++ b/llvm/lib/SandboxIR/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_llvm_component_library(LLVMSandboxIR
   SandboxIR.cpp
+  SandboxIRTracker.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/SandboxIR
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 2984c6eaccd64..08365048074f2 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -16,6 +16,8 @@ using namespace llvm::sandboxir;
 
 Value *Use::get() const { return Ctx->getValue(LLVMUse->get()); }
 
+void Use::set(Value *V) { LLVMUse->set(V->Val); }
+
 unsigned Use::getOperandNo() const { return Usr->getUseOperandNo(*this); }
 
 #ifndef NDEBUG
@@ -115,13 +117,24 @@ void Value::replaceUsesWithIf(
         User *DstU = cast_or_null<User>(Ctx.getValue(LLVMUse.getUser()));
         if (DstU == nullptr)
           return false;
-        return ShouldReplace(Use(&LLVMUse, DstU, Ctx));
+        Use UseToReplace(&LLVMUse, DstU, Ctx);
+        if (!ShouldReplace(UseToReplace))
+          return false;
+        auto &Tracker = Ctx.getTracker();
+        if (Tracker.tracking())
+          Tracker.track(std::make_unique<UseSet>(UseToReplace, Tracker));
+        return true;
       });
 }
 
 void Value::replaceAllUsesWith(Value *Other) {
   assert(getType() == Other->getType() &&
          "Replacing with Value of different type!");
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.tracking()) {
+    for (auto Use : uses())
+      Tracker.track(std::make_unique<UseSet>(Use, Tracker));
+  }
   // We are delegating RAUW to LLVM IR's RAUW.
   Val->replaceAllUsesWith(Other->Val);
 }
@@ -212,11 +225,22 @@ bool User::classof(const Value *From) {
 
 void User::setOperand(unsigned OperandIdx, Value *Operand) {
   assert(isa<llvm::User>(Val) && "No operands!");
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.tracking())
+    Tracker.track(std::make_unique<UseSet>(getOperandUse(OperandIdx), Tracker));
   // We are delegating to llvm::User::setOperand().
   cast<llvm::User>(Val)->setOperand(OperandIdx, Operand->Val);
 }
 
 bool User::replaceUsesOfWith(Value *FromV, Value *ToV) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.tracking()) {
+    for (auto OpIdx : seq<unsigned>(0, getNumOperands())) {
+      auto Use = getOperandUse(OpIdx);
+      if (Use.get() == FromV)
+        Tracker.track(std::make_unique<UseSet>(Use, Tracker));
+    }
+  }
   // We are delegating RUOW to LLVM IR's RUOW.
   return cast<llvm::User>(Val)->replaceUsesOfWith(FromV->Val, ToV->Val);
 }
diff --git a/llvm/lib/SandboxIR/SandboxIRTracker.cpp b/llvm/lib/SandboxIR/SandboxIRTracker.cpp
new file mode 100644
index 0000000000000..0b62df46c020c
--- /dev/null
+++ b/llvm/lib/SandboxIR/SandboxIRTracker.cpp
@@ -0,0 +1,84 @@
+//===- SandboxIRTracker.cpp -----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/SandboxIR/SandboxIRTracker.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+#include <sstream>
+
+using namespace llvm::sandboxir;
+
+IRChangeBase::IRChangeBase(TrackID ID, SandboxIRTracker &Parent)
+    : ID(ID), Parent(Parent) {
+#ifndef NDEBUG
+  Idx = Parent.size();
+
+  assert(!Parent.InMiddleOfCreatingChange &&
+         "We are in the middle of creating another change!");
+  if (Parent.tracking())
+    Parent.InMiddleOfCreatingChange = true;
+#endif // NDEBUG
+}
+
+#ifndef NDEBUG
+void UseSet::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
+
+SandboxIRTracker::~SandboxIRTracker() {
+  assert(Changes.empty() && "You must accept or revert changes!");
+}
+
+void SandboxIRTracker::track(std::unique_ptr<IRChangeBase> &&Change) {
+#ifndef NDEBUG
+  assert(State != TrackerState::Revert &&
+         "No changes should be tracked during revert()!");
+#endif // NDEBUG
+  Changes.push_back(std::move(Change));
+
+#ifndef NDEBUG
+  InMiddleOfCreatingChange = false;
+#endif
+}
+
+void SandboxIRTracker::save() { State = TrackerState::Record; }
+
+void SandboxIRTracker::revert() {
+  auto SavedState = State;
+  State = TrackerState::Revert;
+  for (auto &Change : reverse(Changes))
+    Change->revert();
+  Changes.clear();
+  State = SavedState;
+}
+
+void SandboxIRTracker::accept() {
+  auto SavedState = State;
+  State = TrackerState::Accept;
+  for (auto &Change : Changes)
+    Change->accept();
+  Changes.clear();
+  State = SavedState;
+}
+
+#ifndef NDEBUG
+void SandboxIRTracker::dump(raw_ostream &OS) const {
+  for (const auto &ChangePtr : Changes) {
+    ChangePtr->dump(OS);
+    OS << "\n";
+  }
+}
+void SandboxIRTracker::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
diff --git a/llvm/unittests/SandboxIR/CMakeLists.txt b/llvm/unittests/SandboxIR/CMakeLists.txt
index 362653bfff965..1bb1a6efbef30 100644
--- a/llvm/unittests/SandboxIR/CMakeLists.txt
+++ b/llvm/unittests/SandboxIR/CMakeLists.txt
@@ -6,4 +6,5 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest(SandboxIRTests
   SandboxIRTest.cpp
+  SandboxIRTrackerTest.cpp
   )
diff --git a/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
new file mode 100644
index 0000000000000..380d5d9ac1fd8
--- /dev/null
+++ b/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
@@ -0,0 +1,154 @@
+//===- SandboxIRTrackerTest.cpp -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+struct SandboxIRTrackerTest : public testing::Test {
+  LLVMContext C;
+  std::unique_ptr<Module> M;
+
+  void parseIR(LLVMContext &C, const char *IR) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(IR, Err, C);
+    if (!M)
+      Err.print("SandboxIRTrackerTest", errs());
+  }
+  BasicBlock *getBasicBlockByName(Function &F, StringRef Name) {
+    for (BasicBlock &BB : F)
+      if (BB.getName() == Name)
+        return &BB;
+    llvm_unreachable("Expected to find basic block!");
+  }
+};
+
+TEST_F(SandboxIRTrackerTest, SetOperand) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %gep0 = getelementptr float, ptr %ptr, i32 0
+  %gep1 = getelementptr float, ptr %ptr, i32 1
+  %ld0 = load float, ptr %gep0
+  store float undef, ptr %gep0
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  auto *F = Ctx.createFunction(&LLVMF);
+  auto *BB = &*F->begin();
+  auto &Tracker = Ctx.getTracker();
+  Tracker.save();
+  auto It = BB->begin();
+  auto *Gep0 = &*It++;
+  auto *Gep1 = &*It++;
+  auto *Ld = &*It++;
+  auto *St = &*It++;
+  St->setOperand(0, Ld);
+  EXPECT_EQ(Tracker.size(), 1u);
+  St->setOperand(1, Gep1);
+  EXPECT_EQ(Tracker.size(), 2u);
+  Ld->setOperand(0, Gep1);
+  EXPECT_EQ(Tracker.size(), 3u);
+  EXPECT_EQ(St->getOperand(0), Ld);
+  EXPECT_EQ(St->getOperand(1), Gep1);
+  EXPECT_EQ(Ld->getOperand(0), Gep1);
+
+  Ctx.getTracker().revert();
+  EXPECT_NE(St->getOperand(0), Ld);
+  EXPECT_EQ(St->getOperand(1), Gep0);
+  EXPECT_EQ(Ld->getOperand(0), Gep0);
+}
+
+TEST_F(SandboxIRTrackerTest, RUWIf_RAUW_RUOW) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %ld0 = load float, ptr %ptr
+  %ld1 = load float, ptr %ptr
+  store float %ld0, ptr %ptr
+  store float %ld0, ptr %ptr
+  ret void
+}
+)IR");
+  llvm::Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  llvm::BasicBlock *LLVMBB = &*LLVMF.begin();
+  auto &Tracker = Ctx.getTracker();
+  Ctx.createFunction(&LLVMF);
+  auto *BB = cast<sandboxir::BasicBlock>(Ctx.getValue(LLVMBB));
+  auto It = BB->begin();
+  sandboxir::Instruction *Ld0 = &*It++;
+  sandboxir::Instruction *Ld1 = &*It++;
+  sandboxir::Instruction *St0 = &*It++;
+  sandboxir::Instruction *St1 = &*It++;
+  Ctx.save();
+  // Check RUWIf when the lambda returns false.
+  Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return false; });
+  EXPECT_TRUE(Tracker.empty());
+
+  // Check RUWIf when the lambda returns true.
+  Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return true; });
+  EXPECT_EQ(Tracker.size(), 2u);
+  EXPECT_EQ(St0->getOperand(0), Ld1);
+  EXPECT_EQ(St1->getOperand(0), Ld1);
+  Ctx.revert();
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+
+  // Check RUWIf user == St0.
+  Ctx.save();
+  Ld0->replaceUsesWithIf(
+      Ld1, [St0](const sandboxir::Use &Use) { return Use.getUser() == St0; });
+  EXPECT_EQ(St0->getOperand(0), Ld1);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+  Ctx.revert();
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+
+  // Check RUWIf user == St1.
+  Ctx.save();
+  Ld0->replaceUsesWithIf(
+      Ld1, [St1](const sandboxir::Use &Use) { return Use.getUser() == St1; });
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld1);
+  Ctx.revert();
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+
+  // Check RAUW.
+  Ctx.save();
+  Ld1->replaceAllUsesWith(Ld0);
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+  Ctx.revert();
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+  EXPECT_EQ(St1->getOperand(0), Ld0);
+
+  // Check RUOW.
+  Ctx.save();
+  St0->replaceUsesOfWith(Ld0, Ld1);
+  EXPECT_EQ(Tracker.size(), 1u);
+  EXPECT_EQ(St0->getOperand(0), Ld1);
+  Ctx.revert();
+  EXPECT_EQ(St0->getOperand(0), Ld0);
+
+  // Check accept().
+  St0->replaceUsesOfWith(Ld0, Ld1);
+  EXPECT_EQ(Tracker.size(), 1u);
+  EXPECT_EQ(St0->getOperand(0), Ld1);
+  Ctx.accept();
+  EXPECT_TRUE(Tracker.empty());
+  EXPECT_EQ(St0->getOperand(0), Ld1);
+}

>From 0124154b1b1321d7cba7027d8c64c4a2e470f709 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Tue, 16 Jul 2024 22:46:58 -0700
Subject: [PATCH 2/6] fixup! [SandboxIR] IR Tracker

---
 llvm/docs/SandboxIR.md                        |  1 +
 llvm/include/llvm/SandboxIR/SandboxIR.h       |  2 +-
 .../{SandboxIRTracker.h => Tracker.h}         | 22 +++++++++----------
 llvm/lib/SandboxIR/CMakeLists.txt             |  2 +-
 llvm/lib/SandboxIR/SandboxIR.cpp              |  8 +++----
 .../{SandboxIRTracker.cpp => Tracker.cpp}     | 16 ++++++--------
 .../SandboxIR/SandboxIRTrackerTest.cpp        |  1 +
 7 files changed, 26 insertions(+), 26 deletions(-)
 rename llvm/include/llvm/SandboxIR/{SandboxIRTracker.h => Tracker.h} (92%)
 rename llvm/lib/SandboxIR/{SandboxIRTracker.cpp => Tracker.cpp} (85%)

diff --git a/llvm/docs/SandboxIR.md b/llvm/docs/SandboxIR.md
index 29f5e5ea9346f..94228e1c7a7b9 100644
--- a/llvm/docs/SandboxIR.md
+++ b/llvm/docs/SandboxIR.md
@@ -55,6 +55,7 @@ For example, for `sandboxir::User::setOperand(OpIdx, sandboxir::Value *Op)`:
 ## IR Change Tracking
 Sandbox IR's state can be saved and restored.
 This is done with the help of the tracker component that is tightly coupled to the public Sandbox IR API functions.
+Please note that nested saves/restores are currently not supported.
 
 To save the state and enable tracking the user needs to call `sandboxir::Context::save()`.
 From this point on any change made to the Sandbox IR state will automatically create a change object and register it with the tracker, without any intervention from the user.
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 3f1528ffa0355..dd15ad73e32b1 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -61,7 +61,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
-#include "llvm/SandboxIR/SandboxIRTracker.h"
+#include "llvm/SandboxIR/Tracker.h"
 #include "llvm/SandboxIR/Use.h"
 #include "llvm/Support/raw_ostream.h"
 #include <iterator>
diff --git a/llvm/include/llvm/SandboxIR/SandboxIRTracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
similarity index 92%
rename from llvm/include/llvm/SandboxIR/SandboxIRTracker.h
rename to llvm/include/llvm/SandboxIR/Tracker.h
index 8a819c578a156..52522d5458bec 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIRTracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -1,4 +1,4 @@
-//===- SandboxIRTracker.h ---------------------------------------*- C++ -*-===//
+//===- Tracker.h ------------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -30,15 +30,15 @@
 // Accepting changes
 // -----------------
 // The user needs to either revert or accept changes before the tracker object
-// is destroyed, or else the tracker destructor will cause a crash.
+// is destroyed. This is enforced in the tracker's destructor.
 // This is the job of `SandboxIRTracker::accept()`. Internally this will go
 // through the change objects in `SandboxIRTracker::Changes` in order, calling
 // `IRChangeBase::accept()`.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
-#define LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+#ifndef LLVM_SANDBOXIR_TRACKER_H
+#define LLVM_SANDBOXIR_TRACKER_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/IRBuilder.h"
@@ -93,10 +93,14 @@ class IRChangeBase {
   }
   virtual void dump(raw_ostream &OS) const = 0;
   LLVM_DUMP_METHOD virtual void dump() const = 0;
+  friend raw_ostream &operator<<(raw_ostream &OS, const IRChangeBase &C) {
+    C.dump(OS);
+    return OS;
+  }
 #endif
 };
 
-/// Change the source Value of a sandboxir::Use.
+/// Tracks the change of the source Value of a sandboxir::Use.
 class UseSet : public IRChangeBase {
   Use U;
   Value *OrigV = nullptr;
@@ -113,10 +117,6 @@ class UseSet : public IRChangeBase {
 #ifndef NDEBUG
   void dump(raw_ostream &OS) const final { dumpCommon(OS); }
   LLVM_DUMP_METHOD void dump() const final;
-  friend raw_ostream &operator<<(raw_ostream &OS, const UseSet &C) {
-    C.dump(OS);
-    return OS;
-  }
 #endif
 };
 
@@ -150,7 +150,7 @@ class SandboxIRTracker {
   /// track Sandbox IR changes.
   void track(std::unique_ptr<IRChangeBase> &&Change);
   /// \Returns true if the tracker is recording changes.
-  bool tracking() const { return State == TrackerState::Record; }
+  bool isTracking() const { return State == TrackerState::Record; }
   /// \Returns the current state of the tracker.
   TrackerState getState() const { return State; }
   /// Turns on IR tracking.
@@ -178,4 +178,4 @@ class SandboxIRTracker {
 
 } // namespace llvm::sandboxir
 
-#endif // LLVM_SANDBOXIR_SANDBOXIRTRACKER_H
+#endif // LLVM_SANDBOXIR_TRACKER_H
diff --git a/llvm/lib/SandboxIR/CMakeLists.txt b/llvm/lib/SandboxIR/CMakeLists.txt
index 74b31fe869aed..6c0666b186b8a 100644
--- a/llvm/lib/SandboxIR/CMakeLists.txt
+++ b/llvm/lib/SandboxIR/CMakeLists.txt
@@ -1,6 +1,6 @@
 add_llvm_component_library(LLVMSandboxIR
   SandboxIR.cpp
-  SandboxIRTracker.cpp
+  Tracker.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/SandboxIR
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 08365048074f2..944869a37989c 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -121,7 +121,7 @@ void Value::replaceUsesWithIf(
         if (!ShouldReplace(UseToReplace))
           return false;
         auto &Tracker = Ctx.getTracker();
-        if (Tracker.tracking())
+        if (Tracker.isTracking())
           Tracker.track(std::make_unique<UseSet>(UseToReplace, Tracker));
         return true;
       });
@@ -131,7 +131,7 @@ void Value::replaceAllUsesWith(Value *Other) {
   assert(getType() == Other->getType() &&
          "Replacing with Value of different type!");
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.tracking()) {
+  if (Tracker.isTracking()) {
     for (auto Use : uses())
       Tracker.track(std::make_unique<UseSet>(Use, Tracker));
   }
@@ -226,7 +226,7 @@ bool User::classof(const Value *From) {
 void User::setOperand(unsigned OperandIdx, Value *Operand) {
   assert(isa<llvm::User>(Val) && "No operands!");
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.tracking())
+  if (Tracker.isTracking())
     Tracker.track(std::make_unique<UseSet>(getOperandUse(OperandIdx), Tracker));
   // We are delegating to llvm::User::setOperand().
   cast<llvm::User>(Val)->setOperand(OperandIdx, Operand->Val);
@@ -234,7 +234,7 @@ void User::setOperand(unsigned OperandIdx, Value *Operand) {
 
 bool User::replaceUsesOfWith(Value *FromV, Value *ToV) {
   auto &Tracker = Ctx.getTracker();
-  if (Tracker.tracking()) {
+  if (Tracker.isTracking()) {
     for (auto OpIdx : seq<unsigned>(0, getNumOperands())) {
       auto Use = getOperandUse(OpIdx);
       if (Use.get() == FromV)
diff --git a/llvm/lib/SandboxIR/SandboxIRTracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
similarity index 85%
rename from llvm/lib/SandboxIR/SandboxIRTracker.cpp
rename to llvm/lib/SandboxIR/Tracker.cpp
index 0b62df46c020c..6e441a52c825a 100644
--- a/llvm/lib/SandboxIR/SandboxIRTracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -1,4 +1,4 @@
-//===- SandboxIRTracker.cpp -----------------------------------------------===//
+//===- Tracker.cpp --------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/SandboxIR/SandboxIRTracker.h"
+#include "llvm/SandboxIR/Tracker.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Instruction.h"
@@ -22,7 +22,7 @@ IRChangeBase::IRChangeBase(TrackID ID, SandboxIRTracker &Parent)
 
   assert(!Parent.InMiddleOfCreatingChange &&
          "We are in the middle of creating another change!");
-  if (Parent.tracking())
+  if (Parent.isTracking())
     Parent.InMiddleOfCreatingChange = true;
 #endif // NDEBUG
 }
@@ -39,10 +39,8 @@ SandboxIRTracker::~SandboxIRTracker() {
 }
 
 void SandboxIRTracker::track(std::unique_ptr<IRChangeBase> &&Change) {
-#ifndef NDEBUG
   assert(State != TrackerState::Revert &&
          "No changes should be tracked during revert()!");
-#endif // NDEBUG
   Changes.push_back(std::move(Change));
 
 #ifndef NDEBUG
@@ -53,21 +51,21 @@ void SandboxIRTracker::track(std::unique_ptr<IRChangeBase> &&Change) {
 void SandboxIRTracker::save() { State = TrackerState::Record; }
 
 void SandboxIRTracker::revert() {
-  auto SavedState = State;
+  assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Revert;
   for (auto &Change : reverse(Changes))
     Change->revert();
   Changes.clear();
-  State = SavedState;
+  State = TrackerState::Disabled;
 }
 
 void SandboxIRTracker::accept() {
-  auto SavedState = State;
+  assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Accept;
   for (auto &Change : Changes)
     Change->accept();
   Changes.clear();
-  State = SavedState;
+  State = TrackerState::Disabled;
 }
 
 #ifndef NDEBUG
diff --git a/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
index 380d5d9ac1fd8..c1e8e9fbfc15a 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
@@ -145,6 +145,7 @@ define void @foo(ptr %ptr) {
   EXPECT_EQ(St0->getOperand(0), Ld0);
 
   // Check accept().
+  Ctx.save();
   St0->replaceUsesOfWith(Ld0, Ld1);
   EXPECT_EQ(Tracker.size(), 1u);
   EXPECT_EQ(St0->getOperand(0), Ld1);

>From 88e93d70aec50be8d26c1f9b8fcd8a59fb0e656d Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 17 Jul 2024 11:22:30 -0700
Subject: [PATCH 3/6] fixup! fixup! [SandboxIR] IR Tracker

---
 llvm/include/llvm/SandboxIR/SandboxIR.h       |  4 +--
 llvm/include/llvm/SandboxIR/Tracker.h         | 34 +++++++++----------
 llvm/lib/SandboxIR/Tracker.cpp                | 16 ++++-----
 llvm/unittests/SandboxIR/CMakeLists.txt       |  2 +-
 ...ndboxIRTrackerTest.cpp => TrackerTest.cpp} | 10 +++---
 5 files changed, 33 insertions(+), 33 deletions(-)
 rename llvm/unittests/SandboxIR/{SandboxIRTrackerTest.cpp => TrackerTest.cpp} (94%)

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index dd15ad73e32b1..c5d59ba47ca31 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -643,7 +643,7 @@ class BasicBlock : public Value {
 class Context {
 protected:
   LLVMContext &LLVMCtx;
-  SandboxIRTracker IRTracker;
+  Tracker IRTracker;
 
   /// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
   /// SandboxIR objects.
@@ -684,7 +684,7 @@ class Context {
 public:
   Context(LLVMContext &LLVMCtx) : LLVMCtx(LLVMCtx) {}
 
-  SandboxIRTracker &getTracker() { return IRTracker; }
+  Tracker &getTracker() { return IRTracker; }
   /// Convenience function for `getTracker().save()`
   void save() { IRTracker.save(); }
   /// Convenience function for `getTracker().revert()`
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 52522d5458bec..dd304d3df1467 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -11,28 +11,28 @@
 //
 // Tracking changes
 // ----------------
-// The user needs to call `SandboxIRTracker::save()` to enable tracking changes
+// The user needs to call `Tracker::save()` to enable tracking changes
 // made to SandboxIR. From that point on, any change made to SandboxIR, will
 // automatically create a change tracking object and register it with the
 // tracker. IR-change objects are subclasses of `IRChangeBase` and get
-// registered with the `SandboxIRTracker::track()` function. The change objects
+// registered with the `Tracker::track()` function. The change objects
 // are saved in the order they are registered with the tracker and are stored in
-// the `SandboxIRTracker::Changes` vector. All of this is done transparently to
+// the `Tracker::Changes` vector. All of this is done transparently to
 // the user.
 //
 // Reverting changes
 // -----------------
-// Calling `SandboxIRTracker::revert()` will restore the state saved when
-// `SandboxIRTracker::save()` was called. Internally this goes through the
-// change objects in `SandboxIRTracker::Changes` in reverse order, calling their
+// Calling `Tracker::revert()` will restore the state saved when
+// `Tracker::save()` was called. Internally this goes through the
+// change objects in `Tracker::Changes` in reverse order, calling their
 // `IRChangeBase::revert()` function one by one.
 //
 // Accepting changes
 // -----------------
 // The user needs to either revert or accept changes before the tracker object
 // is destroyed. This is enforced in the tracker's destructor.
-// This is the job of `SandboxIRTracker::accept()`. Internally this will go
-// through the change objects in `SandboxIRTracker::Changes` in order, calling
+// This is the job of `Tracker::accept()`. Internally this will go
+// through the change objects in `Tracker::Changes` in order, calling
 // `IRChangeBase::accept()`.
 //
 //===----------------------------------------------------------------------===//
@@ -68,7 +68,7 @@ static const char *trackIDToStr(TrackID ID) {
 }
 #endif // NDEBUG
 
-class SandboxIRTracker;
+class Tracker;
 
 /// The base class for IR Change classes.
 class IRChangeBase {
@@ -77,10 +77,10 @@ class IRChangeBase {
   unsigned Idx = 0;
 #endif
   const TrackID ID;
-  SandboxIRTracker &Parent;
+  Tracker &Parent;
 
 public:
-  IRChangeBase(TrackID ID, SandboxIRTracker &Parent);
+  IRChangeBase(TrackID ID, Tracker &Parent);
   TrackID getTrackID() const { return ID; }
   /// This runs when changes get reverted.
   virtual void revert() = 0;
@@ -106,7 +106,7 @@ class UseSet : public IRChangeBase {
   Value *OrigV = nullptr;
 
 public:
-  UseSet(const Use &U, SandboxIRTracker &Tracker)
+  UseSet(const Use &U, Tracker &Tracker)
       : IRChangeBase(TrackID::UseSet, Tracker), U(U), OrigV(U.get()) {}
   // For isa<> etc.
   static bool classof(const IRChangeBase *Other) {
@@ -122,7 +122,7 @@ class UseSet : public IRChangeBase {
 
 /// The tracker collects all the change objects and implements the main API for
 /// saving / reverting / accepting.
-class SandboxIRTracker {
+class Tracker {
 public:
   enum class TrackerState {
     Disabled, ///> Tracking is disabled
@@ -144,8 +144,8 @@ class SandboxIRTracker {
   bool InMiddleOfCreatingChange = false;
 #endif // NDEBUG
 
-  SandboxIRTracker() = default;
-  ~SandboxIRTracker();
+  Tracker() = default;
+  ~Tracker();
   /// Record \p Change and take ownership. This is the main function used to
   /// track Sandbox IR changes.
   void track(std::unique_ptr<IRChangeBase> &&Change);
@@ -169,8 +169,8 @@ class SandboxIRTracker {
   IRChangeBase *getChange(unsigned Idx) const { return Changes[Idx].get(); }
   void dump(raw_ostream &OS) const;
   LLVM_DUMP_METHOD void dump() const;
-  friend raw_ostream &operator<<(raw_ostream &OS, const SandboxIRTracker &C) {
-    C.dump(OS);
+  friend raw_ostream &operator<<(raw_ostream &OS, const Tracker &Tracker) {
+    Tracker.dump(OS);
     return OS;
   }
 #endif // NDEBUG
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 6e441a52c825a..b409e72842695 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -15,7 +15,7 @@
 
 using namespace llvm::sandboxir;
 
-IRChangeBase::IRChangeBase(TrackID ID, SandboxIRTracker &Parent)
+IRChangeBase::IRChangeBase(TrackID ID, Tracker &Parent)
     : ID(ID), Parent(Parent) {
 #ifndef NDEBUG
   Idx = Parent.size();
@@ -34,11 +34,11 @@ void UseSet::dump() const {
 }
 #endif // NDEBUG
 
-SandboxIRTracker::~SandboxIRTracker() {
+Tracker::~Tracker() {
   assert(Changes.empty() && "You must accept or revert changes!");
 }
 
-void SandboxIRTracker::track(std::unique_ptr<IRChangeBase> &&Change) {
+void Tracker::track(std::unique_ptr<IRChangeBase> &&Change) {
   assert(State != TrackerState::Revert &&
          "No changes should be tracked during revert()!");
   Changes.push_back(std::move(Change));
@@ -48,9 +48,9 @@ void SandboxIRTracker::track(std::unique_ptr<IRChangeBase> &&Change) {
 #endif
 }
 
-void SandboxIRTracker::save() { State = TrackerState::Record; }
+void Tracker::save() { State = TrackerState::Record; }
 
-void SandboxIRTracker::revert() {
+void Tracker::revert() {
   assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Revert;
   for (auto &Change : reverse(Changes))
@@ -59,7 +59,7 @@ void SandboxIRTracker::revert() {
   State = TrackerState::Disabled;
 }
 
-void SandboxIRTracker::accept() {
+void Tracker::accept() {
   assert(State == TrackerState::Record && "Forgot to save()!");
   State = TrackerState::Accept;
   for (auto &Change : Changes)
@@ -69,13 +69,13 @@ void SandboxIRTracker::accept() {
 }
 
 #ifndef NDEBUG
-void SandboxIRTracker::dump(raw_ostream &OS) const {
+void Tracker::dump(raw_ostream &OS) const {
   for (const auto &ChangePtr : Changes) {
     ChangePtr->dump(OS);
     OS << "\n";
   }
 }
-void SandboxIRTracker::dump() const {
+void Tracker::dump() const {
   dump(dbgs());
   dbgs() << "\n";
 }
diff --git a/llvm/unittests/SandboxIR/CMakeLists.txt b/llvm/unittests/SandboxIR/CMakeLists.txt
index 1bb1a6efbef30..3f43f6337b919 100644
--- a/llvm/unittests/SandboxIR/CMakeLists.txt
+++ b/llvm/unittests/SandboxIR/CMakeLists.txt
@@ -6,5 +6,5 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest(SandboxIRTests
   SandboxIRTest.cpp
-  SandboxIRTrackerTest.cpp
+  TrackerTest.cpp
   )
diff --git a/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
similarity index 94%
rename from llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
rename to llvm/unittests/SandboxIR/TrackerTest.cpp
index c1e8e9fbfc15a..f8d6b6af58dc1 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -1,4 +1,4 @@
-//===- SandboxIRTrackerTest.cpp -------------------------------------------===//
+//===- TrackerTest.cpp ----------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -17,7 +17,7 @@
 
 using namespace llvm;
 
-struct SandboxIRTrackerTest : public testing::Test {
+struct TrackerTest : public testing::Test {
   LLVMContext C;
   std::unique_ptr<Module> M;
 
@@ -25,7 +25,7 @@ struct SandboxIRTrackerTest : public testing::Test {
     SMDiagnostic Err;
     M = parseAssemblyString(IR, Err, C);
     if (!M)
-      Err.print("SandboxIRTrackerTest", errs());
+      Err.print("TrackerTest", errs());
   }
   BasicBlock *getBasicBlockByName(Function &F, StringRef Name) {
     for (BasicBlock &BB : F)
@@ -35,7 +35,7 @@ struct SandboxIRTrackerTest : public testing::Test {
   }
 };
 
-TEST_F(SandboxIRTrackerTest, SetOperand) {
+TEST_F(TrackerTest, SetOperand) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr) {
   %gep0 = getelementptr float, ptr %ptr, i32 0
@@ -72,7 +72,7 @@ define void @foo(ptr %ptr) {
   EXPECT_EQ(Ld->getOperand(0), Gep0);
 }
 
-TEST_F(SandboxIRTrackerTest, RUWIf_RAUW_RUOW) {
+TEST_F(TrackerTest, RUWIf_RAUW_RUOW) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr) {
   %ld0 = load float, ptr %ptr

>From 9a20d51df8b955f853747798eff89d54df8d11ec Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 17 Jul 2024 13:00:22 -0700
Subject: [PATCH 4/6] fixup! fixup! fixup! [SandboxIR] IR Tracker

---
 llvm/docs/SandboxIR.md                |  8 ++++-
 llvm/include/llvm/SandboxIR/Tracker.h | 44 ++++++++++-----------------
 llvm/lib/SandboxIR/Tracker.cpp        |  8 +++--
 3 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/llvm/docs/SandboxIR.md b/llvm/docs/SandboxIR.md
index 94228e1c7a7b9..86f9e2d2bd7ca 100644
--- a/llvm/docs/SandboxIR.md
+++ b/llvm/docs/SandboxIR.md
@@ -62,4 +62,10 @@ From this point on any change made to the Sandbox IR state will automatically cr
 The changes are accumulated in a vector within the tracker.
 
 To rollback to the saved state the user needs to call `sandboxir::Context::revert()`.
-Reverting back to the saved state is a matter of going over all the accumulated states in reverse and undoing each individual change.
+Reverting back to the saved state is a matter of going over all the accumulated changes in reverse and undoing each individual change.
+
+To accept the changes made to the IR the user needs to call `sandboxir::Context::accept()`.
+Internally this will go through the changes and run any finalization required.
+
+Please note that after a call to `revert()` or `accept()` tracking will stop.
+So the user would need to start it again if needed with a call to `save()`.
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index dd304d3df1467..00625e441fa97 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -52,45 +52,32 @@
 namespace llvm::sandboxir {
 
 class BasicBlock;
-
-/// Each IR change type has an ID.
-enum class TrackID {
-  UseSet,
-};
-
-#ifndef NDEBUG
-static const char *trackIDToStr(TrackID ID) {
-  switch (ID) {
-  case TrackID::UseSet:
-    return "UseSet";
-  }
-  llvm_unreachable("Unimplemented ID");
-}
-#endif // NDEBUG
-
 class Tracker;
 
 /// The base class for IR Change classes.
 class IRChangeBase {
 protected:
 #ifndef NDEBUG
+  /// Index within the `Changes` vector, used for debugging.
   unsigned Idx = 0;
+  /// The name is only used for debugging.
+  const char *Name;
 #endif
-  const TrackID ID;
   Tracker &Parent;
 
 public:
-  IRChangeBase(TrackID ID, Tracker &Parent);
-  TrackID getTrackID() const { return ID; }
+#ifndef NDEBUG
+  IRChangeBase(const char *Name, Tracker &Parent);
+#else
+  IRChangeBase(Tracker &Parent);
+#endif
   /// This runs when changes get reverted.
   virtual void revert() = 0;
   /// This runs when changes get accepted.
   virtual void accept() = 0;
   virtual ~IRChangeBase() = default;
-#ifndef NDEBUG
-  void dumpCommon(raw_ostream &OS) const {
-    OS << Idx << ". " << trackIDToStr(ID);
-  }
+#ifndef NDEBUGn
+  void dumpCommon(raw_ostream &OS) const { OS << Idx << ". " << Name; }
   virtual void dump(raw_ostream &OS) const = 0;
   LLVM_DUMP_METHOD virtual void dump() const = 0;
   friend raw_ostream &operator<<(raw_ostream &OS, const IRChangeBase &C) {
@@ -106,12 +93,13 @@ class UseSet : public IRChangeBase {
   Value *OrigV = nullptr;
 
 public:
+#ifndef NDEBUG
   UseSet(const Use &U, Tracker &Tracker)
-      : IRChangeBase(TrackID::UseSet, Tracker), U(U), OrigV(U.get()) {}
-  // For isa<> etc.
-  static bool classof(const IRChangeBase *Other) {
-    return Other->getTrackID() == TrackID::UseSet;
-  }
+      : IRChangeBase("UseSet", Tracker), U(U), OrigV(U.get()) {}
+#else
+  UseSet(const Use &U, Tracker &Tracker)
+      : IRChangeBase(Tracker), U(U), OrigV(U.get()) {}
+#endif // NDEBUG
   void revert() final { U.set(OrigV); }
   void accept() final {}
 #ifndef NDEBUG
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index b409e72842695..99b598ac84246 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -15,8 +15,12 @@
 
 using namespace llvm::sandboxir;
 
-IRChangeBase::IRChangeBase(TrackID ID, Tracker &Parent)
-    : ID(ID), Parent(Parent) {
+#ifndef NDEBUG
+IRChangeBase::IRChangeBase(const char *Name, Tracker &Parent)
+    : Name(Name), Parent(Parent) {
+#else
+IRChangeBase::IRChangeBase(Tracker &Parent) : Parent(Parent) {
+#endif // NDEBUG
 #ifndef NDEBUG
   Idx = Parent.size();
 

>From f69fea26c86921e196da09143f7524eb1da22b02 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 17 Jul 2024 14:46:01 -0700
Subject: [PATCH 5/6] fixup! fixup! fixup! fixup! [SandboxIR] IR Tracker

---
 llvm/docs/SandboxIR.md                |  2 +-
 llvm/include/llvm/SandboxIR/Tracker.h | 32 ++++++++++-----------------
 llvm/lib/SandboxIR/Tracker.cpp        | 13 +++++------
 3 files changed, 19 insertions(+), 28 deletions(-)

diff --git a/llvm/docs/SandboxIR.md b/llvm/docs/SandboxIR.md
index 86f9e2d2bd7ca..3b792659bb59b 100644
--- a/llvm/docs/SandboxIR.md
+++ b/llvm/docs/SandboxIR.md
@@ -68,4 +68,4 @@ To accept the changes made to the IR the user needs to call `sandboxir::Context:
 Internally this will go through the changes and run any finalization required.
 
 Please note that after a call to `revert()` or `accept()` tracking will stop.
-So the user would need to start it again if needed with a call to `save()`.
+To start tracking again, the user needs to call `save()`.
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 00625e441fa97..9456959492350 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -57,27 +57,20 @@ class Tracker;
 /// The base class for IR Change classes.
 class IRChangeBase {
 protected:
-#ifndef NDEBUG
-  /// Index within the `Changes` vector, used for debugging.
-  unsigned Idx = 0;
-  /// The name is only used for debugging.
-  const char *Name;
-#endif
   Tracker &Parent;
 
 public:
-#ifndef NDEBUG
-  IRChangeBase(const char *Name, Tracker &Parent);
-#else
   IRChangeBase(Tracker &Parent);
-#endif
   /// This runs when changes get reverted.
   virtual void revert() = 0;
   /// This runs when changes get accepted.
   virtual void accept() = 0;
   virtual ~IRChangeBase() = default;
-#ifndef NDEBUGn
-  void dumpCommon(raw_ostream &OS) const { OS << Idx << ". " << Name; }
+#ifndef NDEBUG
+  /// \Returns the index of this change by iterating over all changes in the
+  /// tracker. This is only used for debugging.
+  unsigned getIdx() const;
+  void dumpCommon(raw_ostream &OS) const { OS << getIdx() << ". "; }
   virtual void dump(raw_ostream &OS) const = 0;
   LLVM_DUMP_METHOD virtual void dump() const = 0;
   friend raw_ostream &operator<<(raw_ostream &OS, const IRChangeBase &C) {
@@ -93,17 +86,15 @@ class UseSet : public IRChangeBase {
   Value *OrigV = nullptr;
 
 public:
-#ifndef NDEBUG
-  UseSet(const Use &U, Tracker &Tracker)
-      : IRChangeBase("UseSet", Tracker), U(U), OrigV(U.get()) {}
-#else
   UseSet(const Use &U, Tracker &Tracker)
       : IRChangeBase(Tracker), U(U), OrigV(U.get()) {}
-#endif // NDEBUG
   void revert() final { U.set(OrigV); }
   void accept() final {}
 #ifndef NDEBUG
-  void dump(raw_ostream &OS) const final { dumpCommon(OS); }
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "UseSet";
+  }
   LLVM_DUMP_METHOD void dump() const final;
 #endif
 };
@@ -122,6 +113,9 @@ class Tracker {
 private:
   /// The list of changes that are being tracked.
   SmallVector<std::unique_ptr<IRChangeBase>> Changes;
+#ifndef NDEBUG
+  friend unsigned IRChangeBase::getIdx() const; // For accessing `Changes`.
+#endif
   /// The current state of the tracker.
   TrackerState State = TrackerState::Disabled;
 
@@ -153,8 +147,6 @@ class Tracker {
   bool empty() const { return Changes.empty(); }
 
 #ifndef NDEBUG
-  /// \Returns the \p Idx'th change. This is used for testing.
-  IRChangeBase *getChange(unsigned Idx) const { return Changes[Idx].get(); }
   void dump(raw_ostream &OS) const;
   LLVM_DUMP_METHOD void dump() const;
   friend raw_ostream &operator<<(raw_ostream &OS, const Tracker &Tracker) {
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 99b598ac84246..2c7bff94e5c68 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -15,15 +15,8 @@
 
 using namespace llvm::sandboxir;
 
-#ifndef NDEBUG
-IRChangeBase::IRChangeBase(const char *Name, Tracker &Parent)
-    : Name(Name), Parent(Parent) {
-#else
 IRChangeBase::IRChangeBase(Tracker &Parent) : Parent(Parent) {
-#endif // NDEBUG
 #ifndef NDEBUG
-  Idx = Parent.size();
-
   assert(!Parent.InMiddleOfCreatingChange &&
          "We are in the middle of creating another change!");
   if (Parent.isTracking())
@@ -32,6 +25,12 @@ IRChangeBase::IRChangeBase(Tracker &Parent) : Parent(Parent) {
 }
 
 #ifndef NDEBUG
+unsigned IRChangeBase::getIdx() const {
+  auto It =
+      find_if(Parent.Changes, [this](auto &Ptr) { return Ptr.get() == this; });
+  return It - Parent.Changes.begin();
+}
+
 void UseSet::dump() const {
   dump(dbgs());
   dbgs() << "\n";

>From 23621814594c669800b77bbfd392294a6444e399 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Wed, 17 Jul 2024 16:57:56 -0700
Subject: [PATCH 6/6] fixup! fixup! fixup! fixup! fixup! [SandboxIR] IR Tracker

---
 llvm/include/llvm/SandboxIR/Tracker.h    | 6 ------
 llvm/lib/SandboxIR/Tracker.cpp           | 9 +++------
 llvm/unittests/SandboxIR/TrackerTest.cpp | 9 ---------
 3 files changed, 3 insertions(+), 21 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 9456959492350..2d0904f5665b1 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -106,8 +106,6 @@ class Tracker {
   enum class TrackerState {
     Disabled, ///> Tracking is disabled
     Record,   ///> Tracking changes
-    Revert,   ///> Undoing changes
-    Accept,   ///> Accepting changes
   };
 
 private:
@@ -141,10 +139,6 @@ class Tracker {
   void accept();
   /// Stops tracking and reverts to saved state.
   void revert();
-  /// \Returns the number of change entries recorded so far.
-  unsigned size() const { return Changes.size(); }
-  /// \Returns true if there are no change entries recorded so far.
-  bool empty() const { return Changes.empty(); }
 
 #ifndef NDEBUG
   void dump(raw_ostream &OS) const;
diff --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 2c7bff94e5c68..1182f5c55d10b 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -42,8 +42,7 @@ Tracker::~Tracker() {
 }
 
 void Tracker::track(std::unique_ptr<IRChangeBase> &&Change) {
-  assert(State != TrackerState::Revert &&
-         "No changes should be tracked during revert()!");
+  assert(State == TrackerState::Record && "The tracker should be tracking!");
   Changes.push_back(std::move(Change));
 
 #ifndef NDEBUG
@@ -55,20 +54,18 @@ void Tracker::save() { State = TrackerState::Record; }
 
 void Tracker::revert() {
   assert(State == TrackerState::Record && "Forgot to save()!");
-  State = TrackerState::Revert;
+  State = TrackerState::Disabled;
   for (auto &Change : reverse(Changes))
     Change->revert();
   Changes.clear();
-  State = TrackerState::Disabled;
 }
 
 void Tracker::accept() {
   assert(State == TrackerState::Record && "Forgot to save()!");
-  State = TrackerState::Accept;
+  State = TrackerState::Disabled;
   for (auto &Change : Changes)
     Change->accept();
   Changes.clear();
-  State = TrackerState::Disabled;
 }
 
 #ifndef NDEBUG
diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index f8d6b6af58dc1..0902a5f7c8d15 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -57,11 +57,8 @@ define void @foo(ptr %ptr) {
   auto *Ld = &*It++;
   auto *St = &*It++;
   St->setOperand(0, Ld);
-  EXPECT_EQ(Tracker.size(), 1u);
   St->setOperand(1, Gep1);
-  EXPECT_EQ(Tracker.size(), 2u);
   Ld->setOperand(0, Gep1);
-  EXPECT_EQ(Tracker.size(), 3u);
   EXPECT_EQ(St->getOperand(0), Ld);
   EXPECT_EQ(St->getOperand(1), Gep1);
   EXPECT_EQ(Ld->getOperand(0), Gep1);
@@ -85,7 +82,6 @@ define void @foo(ptr %ptr) {
   llvm::Function &LLVMF = *M->getFunction("foo");
   sandboxir::Context Ctx(C);
   llvm::BasicBlock *LLVMBB = &*LLVMF.begin();
-  auto &Tracker = Ctx.getTracker();
   Ctx.createFunction(&LLVMF);
   auto *BB = cast<sandboxir::BasicBlock>(Ctx.getValue(LLVMBB));
   auto It = BB->begin();
@@ -96,11 +92,9 @@ define void @foo(ptr %ptr) {
   Ctx.save();
   // Check RUWIf when the lambda returns false.
   Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return false; });
-  EXPECT_TRUE(Tracker.empty());
 
   // Check RUWIf when the lambda returns true.
   Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return true; });
-  EXPECT_EQ(Tracker.size(), 2u);
   EXPECT_EQ(St0->getOperand(0), Ld1);
   EXPECT_EQ(St1->getOperand(0), Ld1);
   Ctx.revert();
@@ -139,7 +133,6 @@ define void @foo(ptr %ptr) {
   // Check RUOW.
   Ctx.save();
   St0->replaceUsesOfWith(Ld0, Ld1);
-  EXPECT_EQ(Tracker.size(), 1u);
   EXPECT_EQ(St0->getOperand(0), Ld1);
   Ctx.revert();
   EXPECT_EQ(St0->getOperand(0), Ld0);
@@ -147,9 +140,7 @@ define void @foo(ptr %ptr) {
   // Check accept().
   Ctx.save();
   St0->replaceUsesOfWith(Ld0, Ld1);
-  EXPECT_EQ(Tracker.size(), 1u);
   EXPECT_EQ(St0->getOperand(0), Ld1);
   Ctx.accept();
-  EXPECT_TRUE(Tracker.empty());
   EXPECT_EQ(St0->getOperand(0), Ld1);
 }



More information about the llvm-commits mailing list