[llvm] [SandboxIR] SetUse callback (PR #126985)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Feb 13 12:47:01 PST 2025
https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/126985
>From cab8cf73a09ea0a0fb8044dce2e1d217874bf3c0 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 | 68 ++++++++++++++-
llvm/unittests/SandboxIR/RegionTest.cpp | 107 ++++++++++++++++++++++++
3 files changed, 189 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index c1195141cb54c..66eda34ea2edb 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 = "sandboxaux";
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 dbb000e5dd925..ede738c35f032 100644
--- a/llvm/lib/SandboxIR/Region.cpp
+++ b/llvm/lib/SandboxIR/Region.cpp
@@ -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));
+ 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.
@@ -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 {
@@ -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);
}
- 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 1ee72d127daa4..f1ba7ec54bbf3 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, !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
+}
+
+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