[llvm] [SandboxIR][Region] Implement an auxiliary vector in Region (PR #126376)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Feb 10 09:00:37 PST 2025
https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/126376
>From 4ddf85d3d84fe14a03f961fafb30cb664e77b29a Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Tue, 28 Jan 2025 17:01:54 -0800
Subject: [PATCH] [SandboxIR][Region] Implement an auxiliary vector in Region
This patch adds additional functionality to the sandboxir Region.
The Region is used as a way of passing a set of Instructions across region
passes in a way that can be represented in the IR with metadata.
This is a design choice that allows us to test region passes in isolation
with lit tests.
Up until now the region was only used to tag the instructions generated
by the passes. There is a need to represent an ordered set of instructions,
which can be used as a way to represent the initial seeds to the first
vectorization pass. This patch implements this auxiliary vector that
can be used to convey such information.
---
llvm/include/llvm/SandboxIR/Region.h | 17 ++++
llvm/lib/SandboxIR/Region.cpp | 66 ++++++++++++++-
llvm/unittests/SandboxIR/RegionTest.cpp | 107 ++++++++++++++++++++++++
3 files changed, 187 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index c1195141cb54ce3..02370e9e7124371 100644
--- a/llvm/include/llvm/SandboxIR/Region.h
+++ b/llvm/include/llvm/SandboxIR/Region.h
@@ -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 = "sbaux";
Context &Ctx;
/// Keeps track of cost of instructions added and removed.
@@ -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();
@@ -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(); }
diff --git a/llvm/lib/SandboxIR/Region.cpp b/llvm/lib/SandboxIR/Region.cpp
index dbb000e5dd92530..4728c0a34fc3d14 100644
--- a/llvm/lib/SandboxIR/Region.cpp
+++ b/llvm/lib/SandboxIR/Region.cpp
@@ -57,6 +57,39 @@ 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));
+ 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) {
+ 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.
@@ -78,6 +111,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 {
@@ -93,16 +135,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);
}
- 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;
}
diff --git a/llvm/unittests/SandboxIR/RegionTest.cpp b/llvm/unittests/SandboxIR/RegionTest.cpp
index 1ee72d127daa4d5..43ea0c75a9fe9e3 100644
--- a/llvm/unittests/SandboxIR/RegionTest.cpp
+++ b/llvm/unittests/SandboxIR/RegionTest.cpp
@@ -292,3 +292,110 @@ 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, !sbaux !2
+ %t1 = add i8 %v, 1, !sandboxvec !0, !sbaux !3
+ %t2 = add i8 %v, 2, !sandboxvec !1
+ %t3 = add i8 %v, 3, !sandboxvec !1, !sbaux !2
+ %t4 = add i8 %v, 4, !sandboxvec !1, !sbaux !4
+ %t5 = add i8 %v, 5, !sandboxvec !1, !sbaux !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("sbaux"));
+ EXPECT_TRUE(LLVMI1->getMetadata("sbaux"));
+ Regions[0]->clearAux();
+ EXPECT_TRUE(Regions[0]->getAux().empty());
+ EXPECT_FALSE(LLVMI0->getMetadata("sbaux"));
+ EXPECT_FALSE(LLVMI1->getMetadata("sbaux"));
+}
+
+// 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, !sbaux !2
+ %t1 = add i8 %v, 1, !sandboxvec !0, !sbaux !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
+}
+
+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));
+}
More information about the llvm-commits
mailing list