[SandboxIR] Add setVolatile function to LoadInst and StoreInst (#101759)

This patch adds sandboxir::LoadInst::setVolatile() and sandboxir::StoreInst::setVolatile()
and the corresponding tracking class.

Author: wizardengineer

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -870,7 +870,8 @@ class LoadInst final : public UnaryInstruction {
 public:
   /// Return true if this is a load from a volatile memory location.
   bool isVolatile() const { return cast<llvm::LoadInst>(Val)->isVolatile(); }
-
+  /// Specify whether this is a volatile load or not.
+  void setVolatile(bool V);
   unsigned getUseOperandNo(const Use &Use) const final {
     return getUseOperandNoDefault(Use);
   }
@@ -919,6 +920,8 @@ class StoreInst final : public Instruction {
 public:
   /// Return true if this is a store from a volatile memory location.
   bool isVolatile() const { return cast<llvm::StoreInst>(Val)->isVolatile(); }
+  /// Specify whether this is a volatile store or not.
+  void setVolatile(bool V);
   unsigned getUseOperandNo(const Use &Use) const final {
     return getUseOperandNoDefault(Use);
   }

llvm/include/llvm/SandboxIR/Tracker.h

@@ -54,6 +54,8 @@ namespace llvm::sandboxir {
 
 class BasicBlock;
 class CallBrInst;
+class LoadInst;
+class StoreInst;
 class Instruction;
 class Tracker;
 
@@ -271,6 +273,25 @@ class CallBrInstSetIndirectDest : public IRChangeBase {
 #endif
 };
 
+class SetVolatile : public IRChangeBase {
+  /// This holds the properties of whether LoadInst or StoreInst was volatile
+  bool WasVolatile;
+  /// This could either be StoreInst or LoadInst
+  PointerUnion<StoreInst *, LoadInst *> StoreOrLoad;
+
+public:
+  SetVolatile(Instruction *I, Tracker &Tracker);
+  void revert() final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final {
+    dumpCommon(OS);
+    OS << "SetVolatile";
+  }
+  LLVM_DUMP_METHOD void dump() const final;
+#endif
+};
+
 class MoveInstr : public IRChangeBase {
   /// The instruction that moved.
   Instruction *MovedI;

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -610,6 +610,13 @@ void BranchInst::dump() const {
 }
 #endif // NDEBUG
 
+void LoadInst::setVolatile(bool V) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<SetVolatile>(this, Tracker));
+  cast<llvm::LoadInst>(Val)->setVolatile(V);
+}
+
 LoadInst *LoadInst::create(Type *Ty, Value *Ptr, MaybeAlign Align,
                            Instruction *InsertBefore, Context &Ctx,
                            const Twine &Name) {
@@ -664,6 +671,14 @@ void LoadInst::dump() const {
   dbgs() << "\n";
 }
 #endif // NDEBUG
+
+void StoreInst::setVolatile(bool V) {
+  auto &Tracker = Ctx.getTracker();
+  if (Tracker.isTracking())
+    Tracker.track(std::make_unique<SetVolatile>(this, Tracker));
+  cast<llvm::StoreInst>(Val)->setVolatile(V);
+}
+
 StoreInst *StoreInst::create(Value *V, Value *Ptr, MaybeAlign Align,
                              Instruction *InsertBefore, Context &Ctx) {
   return create(V, Ptr, Align, InsertBefore, /*IsVolatile=*/false, Ctx);

llvm/lib/SandboxIR/Tracker.cpp

@@ -235,6 +235,35 @@ void CallBrInstSetIndirectDest::dump() const {
 }
 #endif
 
+SetVolatile::SetVolatile(Instruction *I, Tracker &Tracker)
+    : IRChangeBase(Tracker) {
+  if (auto *Load = dyn_cast<LoadInst>(I)) {
+    WasVolatile = Load->isVolatile();
+    StoreOrLoad = Load;
+  } else if (auto *Store = dyn_cast<StoreInst>(I)) {
+    WasVolatile = Store->isVolatile();
+    StoreOrLoad = Store;
+  } else {
+    llvm_unreachable("Expected LoadInst or StoreInst");
+  }
+}
+
+void SetVolatile::revert() {
+  if (auto *Load = StoreOrLoad.dyn_cast<LoadInst *>()) {
+    Load->setVolatile(WasVolatile);
+  } else if (auto *Store = StoreOrLoad.dyn_cast<StoreInst *>()) {
+    Store->setVolatile(WasVolatile);
+  } else {
+    llvm_unreachable("Expected LoadInst or StoreInst");
+  }
+}
+#ifndef NDEBUG
+void SetVolatile::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif
+
 MoveInstr::MoveInstr(Instruction *MovedI, Tracker &Tracker)
     : IRChangeBase(Tracker), MovedI(MovedI) {
   if (auto *NextI = MovedI->getNextNode())

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -754,6 +754,7 @@ define void @foo(ptr %arg0, ptr %arg1) {
   EXPECT_TRUE(isa<sandboxir::UnaryInstruction>(Ld));
   auto *VLd = cast<sandboxir::LoadInst>(&*It++);
   auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+  bool OrigVolatileValue;
 
   // Check isVolatile()
   EXPECT_FALSE(Ld->isVolatile());
@@ -768,6 +769,11 @@ define void @foo(ptr %arg0, ptr %arg1) {
       sandboxir::LoadInst::create(Ld->getType(), Arg1, Align(8),
                                   /*InsertBefore=*/Ret, Ctx, "NewLd");
   EXPECT_FALSE(NewLd->isVolatile());
+  OrigVolatileValue = NewLd->isVolatile();
+  NewLd->setVolatile(true);
+  EXPECT_TRUE(NewLd->isVolatile());
+  NewLd->setVolatile(OrigVolatileValue);
+  EXPECT_FALSE(NewLd->isVolatile());
   EXPECT_EQ(NewLd->getType(), Ld->getType());
   EXPECT_EQ(NewLd->getPointerOperand(), Arg1);
   EXPECT_EQ(NewLd->getAlign(), 8);
@@ -778,6 +784,11 @@ define void @foo(ptr %arg0, ptr %arg1) {
                                   /*InsertBefore=*/Ret,
                                   /*IsVolatile=*/true, Ctx, "NewVLd");
 
+  EXPECT_TRUE(NewVLd->isVolatile());
+  OrigVolatileValue = NewVLd->isVolatile();
+  NewVLd->setVolatile(false);
+  EXPECT_FALSE(NewVLd->isVolatile());
+  NewVLd->setVolatile(OrigVolatileValue);
   EXPECT_TRUE(NewVLd->isVolatile());
   EXPECT_EQ(NewVLd->getName(), "NewVLd");
   // Check create(InsertAtEnd)
@@ -823,6 +834,7 @@ define void @foo(i8 %val, ptr %ptr) {
   auto *St = cast<sandboxir::StoreInst>(&*It++);
   auto *VSt = cast<sandboxir::StoreInst>(&*It++);
   auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
+  bool OrigVolatileValue;
 
   // Check that the StoreInst has been created correctly.
   EXPECT_FALSE(St->isVolatile());
@@ -837,6 +849,11 @@ define void @foo(i8 %val, ptr %ptr) {
       sandboxir::StoreInst::create(Val, Ptr, Align(8),
                                    /*InsertBefore=*/Ret, Ctx);
   EXPECT_FALSE(NewSt->isVolatile());
+  OrigVolatileValue = NewSt->isVolatile();
+  NewSt->setVolatile(true);
+  EXPECT_TRUE(NewSt->isVolatile());
+  NewSt->setVolatile(OrigVolatileValue);
+  EXPECT_FALSE(NewSt->isVolatile());
   EXPECT_EQ(NewSt->getType(), St->getType());
   EXPECT_EQ(NewSt->getValueOperand(), Val);
   EXPECT_EQ(NewSt->getPointerOperand(), Ptr);
@@ -848,6 +865,11 @@ define void @foo(i8 %val, ptr %ptr) {
                                    /*InsertBefore=*/Ret,
                                    /*IsVolatile=*/true, Ctx);
   EXPECT_TRUE(NewVSt->isVolatile());
+  OrigVolatileValue = NewVSt->isVolatile();
+  NewVSt->setVolatile(false);
+  EXPECT_FALSE(NewVSt->isVolatile());
+  NewVSt->setVolatile(OrigVolatileValue);
+  EXPECT_TRUE(NewVSt->isVolatile());
   EXPECT_EQ(NewVSt->getType(), VSt->getType());
   EXPECT_EQ(NewVSt->getValueOperand(), Val);
   EXPECT_EQ(NewVSt->getPointerOperand(), Ptr);

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -708,3 +708,35 @@ define void @foo(i8 %arg0, i8 %arg1, i8 %arg2) {
   EXPECT_EQ(PHI->getIncomingBlock(1), BB1);
   EXPECT_EQ(PHI->getIncomingValue(1), Arg1);
 }
+
+TEST_F(TrackerTest, SetVolatile) {
+  parseIR(C, R"IR(
+define void @foo(ptr %arg0, i8 %val) {
+  %ld = load i8, ptr %arg0, align 64
+  store i8 %val, ptr %arg0, align 64
+  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();
+  auto *Load = cast<sandboxir::LoadInst>(&*It++);
+  auto *Store = cast<sandboxir::StoreInst>(&*It++);
+
+  EXPECT_FALSE(Load->isVolatile());
+  Ctx.save();
+  Load->setVolatile(true);
+  EXPECT_TRUE(Load->isVolatile());
+  Ctx.revert();
+  EXPECT_FALSE(Load->isVolatile());
+
+  EXPECT_FALSE(Store->isVolatile());
+  Ctx.save();
+  Store->setVolatile(true);
+  EXPECT_TRUE(Store->isVolatile());
+  Ctx.revert();
+  EXPECT_FALSE(Store->isVolatile());
+}