[llvm] [SandboxIR][Region] Implement an auxiliary vector in Region (PR #126376)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Feb 8 08:39:41 PST 2025
https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/126376
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.
>From 7226517fb7a2dd5044b2d0e4b06dc2116cbfd594 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 | 15 ++++
llvm/lib/SandboxIR/Region.cpp | 58 ++++++++++++++-
llvm/unittests/SandboxIR/RegionTest.cpp | 96 +++++++++++++++++++++++++
3 files changed, 166 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index c1195141cb54ce3..c81ae3c6b69ebb8 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(uint32_t Idx, Instruction *I);
+
public:
Region(Context &Ctx, TargetTransformInfo &TTI);
~Region();
@@ -124,6 +135,10 @@ 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; }
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..9f4437b7633dc62 100644
--- a/llvm/lib/SandboxIR/Region.cpp
+++ b/llvm/lib/SandboxIR/Region.cpp
@@ -57,6 +57,31 @@ 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(uint32_t Idx, Instruction *I) {
+ uint32_t 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::remove(Instruction *I) {
// Keep track of the instruction cost. This need to be done *before* we remove
// `I` from the region.
@@ -78,6 +103,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 +127,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..31dbf9661daa0be 100644
--- a/llvm/unittests/SandboxIR/RegionTest.cpp
+++ b/llvm/unittests/SandboxIR/RegionTest.cpp
@@ -292,3 +292,99 @@ 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");
+ 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 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