[llvm] [SandboxIR] Add tracking for ShuffleVectorInst::commute. (PR #106644)

[Jorge Gorbe Moya via llvm-commits]

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

Track it as an operand swap + a `setShuffleMask` and delegate to the `llvm::ShuffleVectorInst` implementation.

>From d516f76eec9e28e2cbd87eeb5f76e276b2232339 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Thu, 29 Aug 2024 16:49:27 -0700
Subject: [PATCH] [SandboxIR] Add tracking for ShuffleVectorInst::commute.

Track it as an operand swap + a setShuffleMask and delegate to the
llvm::ShuffleVectorInst implementation.
---
 llvm/include/llvm/SandboxIR/SandboxIR.h  |  2 +-
 llvm/lib/SandboxIR/SandboxIR.cpp         |  7 +++++++
 llvm/unittests/SandboxIR/TrackerTest.cpp | 20 ++++++++++++++++----
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index d4c2efca4ecfa0..f16e900a67b763 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -1057,7 +1057,7 @@ class ShuffleVectorInst final
 
   /// Swap the operands and adjust the mask to preserve the semantics of the
   /// instruction.
-  void commute() { cast<llvm::ShuffleVectorInst>(Val)->commute(); }
+  void commute();
 
   /// Return true if a shufflevector instruction can be formed with the
   /// specified operands.
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index b5786cdafd6307..9d42deef2eba32 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -2134,6 +2134,13 @@ void ShuffleVectorInst::setShuffleMask(ArrayRef<int> Mask) {
   cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
 }
 
+void ShuffleVectorInst::commute() {
+  Ctx.getTracker().emplaceIfTracking<ShuffleVectorSetMask>(this);
+  Ctx.getTracker().emplaceIfTracking<UseSwap>(getOperandUse(0),
+                                              getOperandUse(1));
+  cast<llvm::ShuffleVectorInst>(Val)->commute();
+}
+
 Constant *ShuffleVectorInst::getShuffleMaskForBitcode() const {
   return Ctx.getOrCreateConstant(
       cast<llvm::ShuffleVectorInst>(Val)->getShuffleMaskForBitcode());
diff --git a/llvm/unittests/SandboxIR/TrackerTest.cpp b/llvm/unittests/SandboxIR/TrackerTest.cpp
index ca6effb727bf37..398fe3e9d4553a 100644
--- a/llvm/unittests/SandboxIR/TrackerTest.cpp
+++ b/llvm/unittests/SandboxIR/TrackerTest.cpp
@@ -963,7 +963,7 @@ define void @foo(i32 %cond0, i32 %cond1) {
   EXPECT_EQ(Switch->findCaseDest(BB1), One);
 }
 
-TEST_F(TrackerTest, ShuffleVectorInstSetters) {
+TEST_F(TrackerTest, ShuffleVectorInst) {
   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>
@@ -982,10 +982,22 @@ define void @foo(<2 x i8> %v1, <2 x i8> %v2) {
   SmallVector<int, 2> OrigMask(SVI->getShuffleMask());
   Ctx.save();
   SVI->setShuffleMask(ArrayRef<int>({0, 0}));
-  EXPECT_THAT(SVI->getShuffleMask(),
-              testing::Not(testing::ElementsAreArray(OrigMask)));
+  EXPECT_NE(SVI->getShuffleMask(), ArrayRef<int>(OrigMask));
   Ctx.revert();
-  EXPECT_THAT(SVI->getShuffleMask(), testing::ElementsAreArray(OrigMask));
+  EXPECT_EQ(SVI->getShuffleMask(), ArrayRef<int>(OrigMask));
+
+  // Check commute.
+  auto *Op0 = SVI->getOperand(0);
+  auto *Op1 = SVI->getOperand(1);
+  Ctx.save();
+  SVI->commute();
+  EXPECT_EQ(SVI->getOperand(0), Op1);
+  EXPECT_EQ(SVI->getOperand(1), Op0);
+  EXPECT_NE(SVI->getShuffleMask(), ArrayRef<int>(OrigMask));
+  Ctx.revert();
+  EXPECT_EQ(SVI->getOperand(0), Op0);
+  EXPECT_EQ(SVI->getOperand(1), Op1);
+  EXPECT_EQ(SVI->getShuffleMask(), ArrayRef<int>(OrigMask));
 }
 
 TEST_F(TrackerTest, PossiblyDisjointInstSetters) {