[llvm] b5ba726 - [SandboxIR] Add tracking for `ShuffleVectorInst::setShuffleMask`. (#105590)

[via llvm-commits]

Author: Jorge Gorbe Moya
Date: 2024-08-21T15:56:55-07:00
New Revision: b5ba726577f7e7af880b62a6352c6208bda4cd0b

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

LOG: [SandboxIR] Add tracking for `ShuffleVectorInst::setShuffleMask`. (#105590)

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 01ef8013ea42a0..278951113aed84 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -1024,9 +1024,7 @@ class ShuffleVectorInst final
   static Constant *convertShuffleMaskForBitcode(ArrayRef<int> Mask,
                                                 Type *ResultTy, Context &Ctx);
 
-  void setShuffleMask(ArrayRef<int> Mask) {
-    cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
-  }
+  void setShuffleMask(ArrayRef<int> Mask);
 
   ArrayRef<int> getShuffleMask() const {
     return cast<llvm::ShuffleVectorInst>(Val)->getShuffleMask();

diff  --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 6f205ae2a075c6..c8a9e99a34341d 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -62,6 +62,7 @@ class AllocaInst;
 class CatchSwitchInst;
 class SwitchInst;
 class ConstantInt;
+class ShuffleVectorInst;
 
 /// The base class for IR Change classes.
 class IRChangeBase {
@@ -355,6 +356,20 @@ class CreateAndInsertInst final : public IRChangeBase {
 #endif
 };
 
+class ShuffleVectorSetMask final : public IRChangeBase {
+  ShuffleVectorInst *SVI;
+  SmallVector<int, 8> PrevMask;
+
+public:
+  ShuffleVectorSetMask(ShuffleVectorInst *SVI);
+  void revert(Tracker &Tracker) final;
+  void accept() final {}
+#ifndef NDEBUG
+  void dump(raw_ostream &OS) const final { OS << "ShuffleVectorSetMask"; }
+  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 a62c879b91e8b9..92054e7cab86ee 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -1868,6 +1868,11 @@ Value *ShuffleVectorInst::create(Value *V1, Value *V2, ArrayRef<int> Mask,
   return Ctx.getOrCreateConstant(cast<llvm::Constant>(NewV));
 }
 
+void ShuffleVectorInst::setShuffleMask(ArrayRef<int> Mask) {
+  Ctx.getTracker().emplaceIfTracking<ShuffleVectorSetMask>(this);
+  cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
+}
+
 Constant *ShuffleVectorInst::getShuffleMaskForBitcode() const {
   return Ctx.getOrCreateConstant(
       cast<llvm::ShuffleVectorInst>(Val)->getShuffleMaskForBitcode());

diff  --git a/llvm/lib/SandboxIR/Tracker.cpp b/llvm/lib/SandboxIR/Tracker.cpp
index 38a1c03556650e..953d4bd51353a9 100644
--- a/llvm/lib/SandboxIR/Tracker.cpp
+++ b/llvm/lib/SandboxIR/Tracker.cpp
@@ -234,6 +234,20 @@ void CreateAndInsertInst::dump() const {
 }
 #endif
 
+ShuffleVectorSetMask::ShuffleVectorSetMask(ShuffleVectorInst *SVI)
+    : SVI(SVI), PrevMask(SVI->getShuffleMask()) {}
+
+void ShuffleVectorSetMask::revert(Tracker &Tracker) {
+  SVI->setShuffleMask(PrevMask);
+}
+
+#ifndef NDEBUG
+void ShuffleVectorSetMask::dump() const {
+  dump(dbgs());
+  dbgs() << "\n";
+}
+#endif
+
 void Tracker::save() { State = TrackerState::Record; }
 
 void Tracker::revert() {

diff  --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index 9f502375204024..a2c3080011f162 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -13,6 +13,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/SandboxIR/SandboxIR.h"
 #include "llvm/Support/SourceMgr.h"
+#include "gmock/gmock-matchers.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -792,6 +793,31 @@ define void @foo(i32 %cond0, i32 %cond1) {
   EXPECT_EQ(Switch->findCaseDest(BB1), One);
 }
 
+TEST_F(TrackerTest, ShuffleVectorInstSetters) {
+  parseIR(C, R"IR(
+define void @foo(<2 x i8> %v1, <2 x i8> %v2) {
+  %shuf = shufflevector <2 x i8> %v1, <2 x i8> %v2, <2 x i32> <i32 1, i32 2>
+  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 *SVI = cast<sandboxir::ShuffleVectorInst>(&*It++);
+
+  // Check setShuffleMask.
+  SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
+  Ctx.save();
+  SVI->setShuffleMask(ArrayRef<int>({0, 0}));
+  EXPECT_THAT(SVI->getShuffleMask(),
+              testing::Not(testing::ElementsAreArray(OrigMask)));
+  Ctx.revert();
+  EXPECT_THAT(SVI->getShuffleMask(), testing::ElementsAreArray(OrigMask));
+}
+
 TEST_F(TrackerTest, AtomicRMWSetters) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %arg) {