Skip to content

[SandboxIR][Region] Implement an auxiliary vector in Region #126376

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
merged 1 commit into from
Feb 13, 2025
Merged

[SandboxIR][Region] Implement an auxiliary vector in Region #126376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
17 changes: 17 additions & 0 deletions llvm/include/llvm/SandboxIR/Region.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,16 +87,23 @@ class ScoreBoard {
// | |
// |Rgn3| -> Transform1 -> ... -> TransformN -> Check Cost
// +----+
//
// The region can also hold an ordered sequence of "auxiliary" instructions.
// This can be used to pass auxiliary information across region passes, like for
// example the initial seed slice used by the bottom-up vectorizer.

class Region {
/// All the instructions in the Region. Only new instructions generated during
/// vectorization are part of the Region.
SetVector<Instruction *> Insts;
/// An auxiliary sequence of Instruction-Index pairs.
SmallVector<Instruction *> Aux;

/// MDNode that we'll use to mark instructions as being part of the region.
MDNode *RegionMDN;
static constexpr const char *MDKind = "sandboxvec";
static constexpr const char *RegionStr = "sandboxregion";
static constexpr const char *AuxMDKind = "sandboxaux";

Context &Ctx;
/// Keeps track of cost of instructions added and removed.
Expand All @@ -110,6 +117,10 @@ class Region {
// TODO: Add cost modeling.
// TODO: Add a way to encode/decode region info to/from metadata.

/// Set \p I as the \p Idx'th element in the auxiliary vector.
/// NOTE: This is for internal use, it does not set the metadata.
void setAux(unsigned Idx, Instruction *I);

public:
Region(Context &Ctx, TargetTransformInfo &TTI);
~Region();
Expand All @@ -124,6 +135,12 @@ class Region {
bool contains(Instruction *I) const { return Insts.contains(I); }
/// Returns true if the Region has no instructions.
bool empty() const { return Insts.empty(); }
/// Set the auxiliary vector.
void setAux(ArrayRef<Instruction *> Aux);
/// \Returns the auxiliary vector.
const SmallVector<Instruction *> &getAux() const { return Aux; }
/// Clears all auxiliary data.
void clearAux();

using iterator = decltype(Insts.begin());
iterator begin() { return Insts.begin(); }
Expand Down
68 changes: 65 additions & 3 deletions llvm/lib/SandboxIR/Region.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,41 @@ void Region::add(Instruction *I) {
Scoreboard.add(I);
}

void Region::setAux(ArrayRef<Instruction *> Aux) {
this->Aux = SmallVector<Instruction *>(Aux);
auto &LLVMCtx = Ctx.LLVMCtx;
for (auto [Idx, I] : enumerate(Aux)) {
llvm::ConstantInt *IdxC =
llvm::ConstantInt::get(LLVMCtx, llvm::APInt(32, Idx, false));
Copy link
Collaborator

Choose a reason for hiding this comment

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

https://discourse.llvm.org/t/second-stage-of-build-on-windows-fails-in-sandboxir/84841
@vporpo - could you replace the APInt with the ConstantInt::get(Type *Ty, uint64_t V, bool isSigned) variant?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@RKSimon Thanks for letting me know about the issue. I just created a PR with your suggestion: #129082

assert(cast<llvm::Instruction>(I->Val)->getMetadata(AuxMDKind) == nullptr &&
"Instruction already in Aux!");
cast<llvm::Instruction>(I->Val)->setMetadata(
AuxMDKind, MDNode::get(LLVMCtx, ConstantAsMetadata::get(IdxC)));
}
}

void Region::setAux(unsigned Idx, Instruction *I) {
assert((Idx >= Aux.size() || Aux[Idx] == nullptr) &&
"There is already an Instruction at Idx in Aux!");
unsigned ExpectedSz = Idx + 1;
if (Aux.size() < ExpectedSz) {
auto SzBefore = Aux.size();
Aux.resize(ExpectedSz);
// Initialize the gap with nullptr.
for (unsigned Idx = SzBefore; Idx + 1 < ExpectedSz; ++Idx)
Aux[Idx] = nullptr;
}
Aux[Idx] = I;
}

void Region::clearAux() {
for (unsigned Idx : seq<unsigned>(0, Aux.size())) {
auto *LLVMI = cast<llvm::Instruction>(Aux[Idx]->Val);
LLVMI->setMetadata(AuxMDKind, nullptr);
}
Aux.clear();
}

void Region::remove(Instruction *I) {
// Keep track of the instruction cost. This need to be done *before* we remove
// `I` from the region.
Expand All @@ -78,6 +113,15 @@ bool Region::operator==(const Region &Other) const {
void Region::dump(raw_ostream &OS) const {
for (auto *I : Insts)
OS << *I << "\n";
if (!Aux.empty()) {
OS << "\nAux:\n";
for (auto *I : Aux) {
if (I == nullptr)
OS << "NULL\n";
else
OS << *I << "\n";
}
}
}

void Region::dump() const {
Expand All @@ -93,16 +137,34 @@ Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) {
auto &Ctx = F.getContext();
for (BasicBlock &BB : F) {
for (Instruction &Inst : BB) {
if (auto *MDN = cast<llvm::Instruction>(Inst.Val)->getMetadata(MDKind)) {
auto *LLVMI = cast<llvm::Instruction>(Inst.Val);
if (auto *MDN = LLVMI->getMetadata(MDKind)) {
Region *R = nullptr;
auto [It, Inserted] = MDNToRegion.try_emplace(MDN);
if (Inserted) {
Regions.push_back(std::make_unique<Region>(Ctx, TTI));
It->second = Regions.back().get();
R = Regions.back().get();
It->second = R;
} else {
R = It->second;
}
R->add(&Inst);

if (auto *AuxMDN = LLVMI->getMetadata(AuxMDKind)) {
llvm::Constant *IdxC =
dyn_cast<ConstantAsMetadata>(AuxMDN->getOperand(0))->getValue();
auto Idx = cast<llvm::ConstantInt>(IdxC)->getSExtValue();
R->setAux(Idx, &Inst);
Copy link
Collaborator

Choose a reason for hiding this comment

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

Here we could have a check that we don't have two different values whose metadata claims the same index in the aux vector.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I added a check in setAux(Idx, I) but this function is only called internally so using the same index shouldn't be possible unless this function is broken.

Copy link
Collaborator

@slackito slackito Feb 13, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

You could have incorrect metadata, right? Something like

define void @foo(i8 %v) {
  %t0 = add i8 %v, 0, !sandboxvec !0, !sandboxaux !1
  %t1 = add i8 %v, 1, !sandboxvec !0, !sandboxaux !1
  ret void
}
!0 = distinct !{!"sandboxregion"}
!1 = !{i32 0}

Copy link
Contributor Author

Choose a reason for hiding this comment

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

You are right, I think I should be able to add a test for it too.

}
It->second->add(&Inst);
}
}
}
#ifndef NDEBUG
// Check that there are no gaps in the Aux vector.
for (auto &RPtr : Regions)
for (auto *I : RPtr->getAux())
assert(I != nullptr && "Gap in Aux!");
#endif
return Regions;
}

Expand Down
128 changes: 128 additions & 0 deletions llvm/unittests/SandboxIR/RegionTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,3 +292,131 @@ define void @foo(i8 %v0, i8 %v1, i8 %v2) {
EXPECT_EQ(SB.getBeforeCost(), GetCost(LLVMAdd2));
EXPECT_EQ(SB.getAfterCost(), GetCost(LLVMAdd1));
}

TEST_F(RegionTest, Aux) {
parseIR(C, R"IR(
define void @foo(i8 %v) {
%t0 = add i8 %v, 0, !sandboxvec !0, !sandboxaux !2
%t1 = add i8 %v, 1, !sandboxvec !0, !sandboxaux !3
%t2 = add i8 %v, 2, !sandboxvec !1
%t3 = add i8 %v, 3, !sandboxvec !1, !sandboxaux !2
%t4 = add i8 %v, 4, !sandboxvec !1, !sandboxaux !4
%t5 = add i8 %v, 5, !sandboxvec !1, !sandboxaux !3
ret void
}

!0 = distinct !{!"sandboxregion"}
!1 = distinct !{!"sandboxregion"}

!2 = !{i32 0}
!3 = !{i32 1}
!4 = !{i32 2}
)IR");
llvm::Function *LLVMF = &*M->getFunction("foo");
auto *LLVMBB = &*LLVMF->begin();
auto LLVMIt = LLVMBB->begin();
auto *LLVMI0 = &*LLVMIt++;
auto *LLVMI1 = &*LLVMIt++;
sandboxir::Context Ctx(C);
auto *F = Ctx.createFunction(LLVMF);
auto *BB = &*F->begin();
auto It = BB->begin();
auto *T0 = cast<sandboxir::Instruction>(&*It++);
auto *T1 = cast<sandboxir::Instruction>(&*It++);
auto *T2 = cast<sandboxir::Instruction>(&*It++);
auto *T3 = cast<sandboxir::Instruction>(&*It++);
auto *T4 = cast<sandboxir::Instruction>(&*It++);
auto *T5 = cast<sandboxir::Instruction>(&*It++);

SmallVector<std::unique_ptr<sandboxir::Region>> Regions =
sandboxir::Region::createRegionsFromMD(*F, *TTI);
// Check that the regions are correct.
EXPECT_THAT(Regions[0]->insts(), testing::UnorderedElementsAre(T0, T1));
EXPECT_THAT(Regions[1]->insts(),
testing::UnorderedElementsAre(T2, T3, T4, T5));
// Check aux.
EXPECT_THAT(Regions[0]->getAux(), testing::ElementsAre(T0, T1));
EXPECT_THAT(Regions[1]->getAux(), testing::ElementsAre(T3, T5, T4));
// Check clearAux().
EXPECT_TRUE(LLVMI0->getMetadata("sandboxaux"));
EXPECT_TRUE(LLVMI1->getMetadata("sandboxaux"));
Regions[0]->clearAux();
EXPECT_TRUE(Regions[0]->getAux().empty());
EXPECT_FALSE(LLVMI0->getMetadata("sandboxaux"));
EXPECT_FALSE(LLVMI1->getMetadata("sandboxaux"));
}

// Check that Aux is well-formed.
TEST_F(RegionTest, AuxVerify) {
parseIR(C, R"IR(
define void @foo(i8 %v) {
%t0 = add i8 %v, 0, !sandboxvec !0, !sandboxaux !2
%t1 = add i8 %v, 1, !sandboxvec !0, !sandboxaux !3
ret void
}

!0 = distinct !{!"sandboxregion"}
!2 = !{i32 0}
!3 = !{i32 2}
)IR");
llvm::Function *LLVMF = &*M->getFunction("foo");
sandboxir::Context Ctx(C);
auto *F = Ctx.createFunction(LLVMF);
#ifndef NDEBUG
EXPECT_DEATH(sandboxir::Region::createRegionsFromMD(*F, *TTI), ".*Gap*");
#endif
}

// Check that we get an assertion failure if we try to set the same index more
// than once.
TEST_F(RegionTest, AuxSameIndex) {
parseIR(C, R"IR(
define void @foo(i8 %v) {
%t0 = add i8 %v, 0, !sandboxvec !0, !sandboxaux !2
%t1 = add i8 %v, 1, !sandboxvec !0, !sandboxaux !2
ret void
}

!0 = distinct !{!"sandboxregion"}
!2 = !{i32 0}
)IR");
llvm::Function *LLVMF = &*M->getFunction("foo");
sandboxir::Context Ctx(C);
auto *F = Ctx.createFunction(LLVMF);
#ifndef NDEBUG
EXPECT_DEATH(sandboxir::Region::createRegionsFromMD(*F, *TTI), ".*already.*");
#endif // NDEBUG
}

TEST_F(RegionTest, AuxRoundTrip) {
parseIR(C, R"IR(
define i8 @foo(i8 %v0, i8 %v1) {
%t0 = add i8 %v0, 1
%t1 = add i8 %t0, %v1
ret i8 %t1
}
)IR");
llvm::Function *LLVMF = &*M->getFunction("foo");
sandboxir::Context Ctx(C);
auto *F = Ctx.createFunction(LLVMF);
auto *BB = &*F->begin();
auto It = BB->begin();
auto *T0 = cast<sandboxir::Instruction>(&*It++);
auto *T1 = cast<sandboxir::Instruction>(&*It++);

sandboxir::Region Rgn(Ctx, *TTI);
Rgn.add(T0);
Rgn.add(T1);
#ifndef NDEBUG
EXPECT_DEATH(Rgn.setAux({T0, T0}), ".*already.*");
#endif
Rgn.setAux({T1, T0});

SmallVector<std::unique_ptr<sandboxir::Region>> Regions =
sandboxir::Region::createRegionsFromMD(*F, *TTI);
ASSERT_EQ(1U, Regions.size());
#ifndef NDEBUG
EXPECT_EQ(Rgn, *Regions[0].get());
#endif
EXPECT_THAT(Rgn.getAux(), testing::ElementsAre(T1, T0));
}