Skip to content

[SandboxIR] Add tracking for ShuffleVectorInst::setShuffleMask. #105590

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slackito merged 2 commits into from
Aug 21, 2024
Merged

[SandboxIR] Add tracking for ShuffleVectorInst::setShuffleMask. #105590

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
42c8093
[SandboxIR] Add tracking for ShuffleVectorInst::setShuffleMask
slackito Aug 21, 2024
4d26519
formatting fixes
slackito Aug 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions llvm/include/llvm/SandboxIR/SandboxIR.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
15 changes: 15 additions & 0 deletions llvm/include/llvm/SandboxIR/Tracker.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ class AllocaInst;
class CatchSwitchInst;
class SwitchInst;
class ConstantInt;
class ShuffleVectorInst;

/// The base class for IR Change classes.
class IRChangeBase {
Expand Down Expand Up @@ -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 {
Expand Down
5 changes: 5 additions & 0 deletions llvm/lib/SandboxIR/SandboxIR.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Copy link
Contributor

@vporpo vporpo Aug 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm do you thing there is a way to use the GenericSetter somehow? It won't work as is because it will try to use ArrayRef<int> as its SavedValT OrigVal, but we need a SmallVector instead.

Copy link
Collaborator Author

@slackito slackito Aug 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I created a new class specifically because I thought GenericSetter wouldn't work with ArrayRef because we need a copy. I'm not sure if we can make it work with ArrayRef (maybe adding some template specialization that always copies if the getter returns an ArrayRef? not sure if that's actually possible).

We can always revisit this if we see more instances of getters that return ArrayRef that we could generalize. For now, this is a one-off so I think it's okay to keep it as is.

cast<llvm::ShuffleVectorInst>(Val)->setShuffleMask(Mask);
}

Constant *ShuffleVectorInst::getShuffleMaskForBitcode() const {
return Ctx.getOrCreateConstant(
cast<llvm::ShuffleVectorInst>(Val)->getShuffleMaskForBitcode());
Expand Down
14 changes: 14 additions & 0 deletions llvm/lib/SandboxIR/Tracker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
26 changes: 26 additions & 0 deletions llvm/unittests/SandboxIR/TrackerTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down