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

Jorge Gorbe Moya via llvm-commits llvm-commits at lists.llvm.org
Wed Aug 21 15:02:09 PDT 2024


https://github.com/slackito created https://github.com/llvm/llvm-project/pull/105590

None

>From 42c8093cb34413782a4b4665699f02426144a3df Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Wed, 21 Aug 2024 14:51:43 -0700
Subject: [PATCH 1/2] [SandboxIR] Add tracking for
 ShuffleVectorInst::setShuffleMask

---
 llvm/include/llvm/SandboxIR/SandboxIR.h  |  4 +---
 llvm/include/llvm/SandboxIR/Tracker.h    | 15 +++++++++++++
 llvm/lib/SandboxIR/SandboxIR.cpp         |  5 +++++
 llvm/lib/SandboxIR/Tracker.cpp           | 14 ++++++++++++
 llvm/unittests/SandboxIR/TrackerTest.cpp | 28 ++++++++++++++++++++++++
 5 files changed, 63 insertions(+), 3 deletions(-)

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..c0485d47310aee 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,33 @@ 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++);
+
+  SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
+
+  // Check setShuffleMask.
+  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) {

>From 4d26519f2f1fcc866baae023ffb1e3e81cd586be Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Wed, 21 Aug 2024 14:59:14 -0700
Subject: [PATCH 2/2] formatting fixes

---
 llvm/unittests/SandboxIR/TrackerTest.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index c0485d47310aee..a2c3080011f162 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -808,9 +808,8 @@ define void @foo(<2 x i8> %v1, <2 x i8> %v2) {
   auto It = BB->begin();
   auto *SVI = cast<sandboxir::ShuffleVectorInst>(&*It++);
 
-  SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
-
   // Check setShuffleMask.
+  SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
   Ctx.save();
   SVI->setShuffleMask(ArrayRef<int>({0, 0}));
   EXPECT_THAT(SVI->getShuffleMask(),
@@ -819,7 +818,6 @@ define void @foo(<2 x i8> %v1, <2 x i8> %v2) {
   EXPECT_THAT(SVI->getShuffleMask(), testing::ElementsAreArray(OrigMask));
 }
 
-
 TEST_F(TrackerTest, AtomicRMWSetters) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr, i8 %arg) {



More information about the llvm-commits mailing list