[llvm] [SandboxVec] Tag insts in a Region with metadata. (PR #109353)
Jorge Gorbe Moya via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 20 14:45:07 PDT 2024
https://github.com/slackito updated https://github.com/llvm/llvm-project/pull/109353
>From 56c52ba22f0b1d869560d1269865f045cfc81a47 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Thu, 19 Sep 2024 16:27:17 -0700
Subject: [PATCH 1/5] [SandboxVec] Tag insts in a Region with metadata.
For each region, we create a metadata node with the Region ID. Then
when an instruction is added to the Region, it gets tagged with the
metadata node for that region. In the following example, we have a
Region that contains only the `%t0` instruction.
```
define i8 @foo(i8 %v0, i8 %v1) {
%t0 = add i8 %v0, 1, !sbvec !0
%t1 = add i8 %t0, %v1
ret i8 %t1
}
!0 = !{!"region", i32 0}
```
This commit also adds a function to create regions from metadata already
present in a Function.
This metadata can be used for debugging: if we dump IR before a Region
pass, the IR will contain enough info to re-create the Region and run
the pass by itself in a later invocation.
---
llvm/include/llvm/SandboxIR/SandboxIR.h | 4 ++
.../Vectorize/SandboxVectorizer/Region.h | 9 +++++
.../Vectorize/SandboxVectorizer/Region.cpp | 39 ++++++++++++++++++-
.../SandboxVectorizer/RegionTest.cpp | 31 ++++++++++++++-
4 files changed, 79 insertions(+), 4 deletions(-)
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index c01516aa9d31ac..a16ba6832f42fc 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -338,6 +338,9 @@ class Value {
friend class GlobalIFunc; // For `Val`.
friend class GlobalVariable; // For `Val`.
friend class GlobalAlias; // For `Val`.
+ // Region needs to manipulate metadata in the underlying LLVM Value, we don't
+ // expose metadata in sandboxir.
+ friend class Region;
/// All values point to the context.
Context &Ctx;
@@ -4337,6 +4340,7 @@ class Context {
friend class IntegerType; // For LLVMCtx.
friend class StructType; // For LLVMCtx.
friend class TargetExtType; // For LLVMCtx.
+ friend class Region; // For LLVMCtx.
Tracker IRTracker;
/// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
index 2f893bac213a01..2663de823188ad 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
@@ -9,6 +9,8 @@
#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_REGION_H
#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_REGION_H
+#include <memory>
+
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/SandboxIR/SandboxIR.h"
@@ -58,6 +60,11 @@ class Region {
/// A unique ID, used for debugging.
unsigned RegionID = 0;
+ /// MDNode that we'll use to mark instructions as being part of the region.
+ MDNode *RegionMDN;
+ static constexpr const char *MDKind = "sbvec";
+ static constexpr const char *RegionStr = "region";
+
Context &Ctx;
// TODO: Add cost modeling.
@@ -85,6 +92,8 @@ class Region {
iterator end() { return Insts.end(); }
iterator_range<iterator> insts() { return make_range(begin(), end()); }
+ static SmallVector<std::unique_ptr<Region>> createRegionsFromMD(Function &F);
+
#ifndef NDEBUG
/// This is an expensive check, meant for testing.
bool operator==(const Region &Other) const;
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
index 34aa9f3786f34c..4d26b6c4f7c552 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
@@ -13,13 +13,25 @@ namespace llvm::sandboxir {
Region::Region(Context &Ctx) : Ctx(Ctx) {
static unsigned StaticRegionID;
RegionID = StaticRegionID++;
+
+ LLVMContext &LLVMCtx = Ctx.LLVMCtx;
+ auto *RegionStrMD = MDString::get(LLVMCtx, RegionStr);
+ auto *RegionIDMD = llvm::ConstantAsMetadata::get(
+ llvm::ConstantInt::get(LLVMCtx, APInt(sizeof(RegionID) * 8, RegionID)));
+ RegionMDN = MDNode::get(LLVMCtx, {RegionStrMD, RegionIDMD});
}
Region::~Region() {}
-void Region::add(Instruction *I) { Insts.insert(I); }
+void Region::add(Instruction *I) {
+ Insts.insert(I);
+ cast<llvm::Instruction>(I->Val)->setMetadata(MDKind, RegionMDN);
+}
-void Region::remove(Instruction *I) { Insts.remove(I); }
+void Region::remove(Instruction *I) {
+ Insts.remove(I);
+ cast<llvm::Instruction>(I->Val)->setMetadata(MDKind, nullptr);
+}
#ifndef NDEBUG
bool Region::operator==(const Region &Other) const {
@@ -42,4 +54,27 @@ void Region::dump() const {
}
#endif // NDEBUG
+SmallVector<std::unique_ptr<Region>> Region::createRegionsFromMD(Function &F) {
+ SmallVector<std::unique_ptr<Region>> Regions;
+ DenseMap<MDNode *, Region *> MDNToRegion;
+ auto &Ctx = F.getContext();
+ for (BasicBlock &BB : F) {
+ for (Instruction &Inst : BB) {
+ if (auto *MDN = cast<llvm::Instruction>(Inst.Val)->getMetadata(MDKind)) {
+ Region *R = nullptr;
+ auto It = MDNToRegion.find(MDN);
+ if (It == MDNToRegion.end()) {
+ Regions.push_back(std::make_unique<Region>(Ctx));
+ R = Regions.back().get();
+ MDNToRegion[MDN] = R;
+ } else {
+ R = It->second;
+ }
+ R->add(&Inst);
+ }
+ }
+ }
+ return Regions;
+}
+
} // namespace llvm::sandboxir
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
index 2c7390c515f114..dcd3bd0975014f 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
@@ -45,9 +45,8 @@ define i8 @foo(i8 %v0, i8 %v1) {
auto *Ret = cast<sandboxir::Instruction>(&*It++);
sandboxir::Region Rgn(Ctx);
- // Check getters
+ // Check getContext.
EXPECT_EQ(&Ctx, &Rgn.getContext());
- EXPECT_EQ(0U, Rgn.getID());
// Check add / remove / empty.
EXPECT_TRUE(Rgn.empty());
@@ -79,3 +78,31 @@ define i8 @foo(i8 %v0, i8 %v1) {
EXPECT_EQ(Rgn, Other);
#endif
}
+
+TEST_F(RegionTest, MetadataRoundTrip) {
+ 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);
+ Rgn.add(T0);
+ Rgn.add(T1);
+
+ SmallVector<std::unique_ptr<sandboxir::Region>> Regions =
+ sandboxir::Region::createRegionsFromMD(*F);
+ ASSERT_EQ(1U, Regions.size());
+#ifndef NDEBUG
+ EXPECT_EQ(Rgn, *Regions[0].get());
+#endif
+}
>From 87636bbead29671ca787d4241cd013d04fda89d8 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Fri, 20 Sep 2024 14:12:33 -0700
Subject: [PATCH 2/5] Address review feedback
- Added new test case that creates regions from explicit metadata.
- As suggested offline by @aeubanks, removed region IDs in favor of
using `distinct` metadata nodes to identify each region.
Also removed a stale include for "InstructionCost" left over from the
original version of PR108899. All cost-related stuff was removed.
---
.../Vectorize/SandboxVectorizer/Region.h | 6 -----
.../Vectorize/SandboxVectorizer/Region.cpp | 8 +-----
.../SandboxVectorizer/RegionTest.cpp | 27 +++++++++++++++++++
3 files changed, 28 insertions(+), 13 deletions(-)
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
index 2663de823188ad..03280bb552ec5f 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h
@@ -14,7 +14,6 @@
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/SandboxIR/SandboxIR.h"
-#include "llvm/Support/InstructionCost.h"
#include "llvm/Support/raw_ostream.h"
namespace llvm::sandboxir {
@@ -57,9 +56,6 @@ class Region {
/// vectorization are part of the Region.
SetVector<Instruction *> Insts;
- /// A unique ID, used for debugging.
- unsigned RegionID = 0;
-
/// MDNode that we'll use to mark instructions as being part of the region.
MDNode *RegionMDN;
static constexpr const char *MDKind = "sbvec";
@@ -75,8 +71,6 @@ class Region {
~Region();
Context &getContext() const { return Ctx; }
- /// Returns the region's unique ID.
- unsigned getID() const { return RegionID; }
/// Adds I to the set.
void add(Instruction *I);
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
index 4d26b6c4f7c552..55f08f424655b0 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
@@ -11,14 +11,9 @@
namespace llvm::sandboxir {
Region::Region(Context &Ctx) : Ctx(Ctx) {
- static unsigned StaticRegionID;
- RegionID = StaticRegionID++;
-
LLVMContext &LLVMCtx = Ctx.LLVMCtx;
auto *RegionStrMD = MDString::get(LLVMCtx, RegionStr);
- auto *RegionIDMD = llvm::ConstantAsMetadata::get(
- llvm::ConstantInt::get(LLVMCtx, APInt(sizeof(RegionID) * 8, RegionID)));
- RegionMDN = MDNode::get(LLVMCtx, {RegionStrMD, RegionIDMD});
+ RegionMDN = MDNode::getDistinct(LLVMCtx, {RegionStrMD});
}
Region::~Region() {}
@@ -43,7 +38,6 @@ bool Region::operator==(const Region &Other) const {
}
void Region::dump(raw_ostream &OS) const {
- OS << "RegionID: " << getID() << "\n";
for (auto *I : Insts)
OS << *I << "\n";
}
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
index dcd3bd0975014f..1626a5d2ce567a 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp
@@ -79,6 +79,33 @@ define i8 @foo(i8 %v0, i8 %v1) {
#endif
}
+TEST_F(RegionTest, MetadataFromIR) {
+ parseIR(C, R"IR(
+define i8 @foo(i8 %v0, i8 %v1) {
+ %t0 = add i8 %v0, 1, !sbvec !0
+ %t1 = add i8 %t0, %v1, !sbvec !1
+ %t2 = add i8 %t1, %v1, !sbvec !1
+ ret i8 %t2
+}
+
+!0 = distinct !{!"region"}
+!1 = distinct !{!"region"}
+)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++);
+
+ SmallVector<std::unique_ptr<sandboxir::Region>> Regions =
+ sandboxir::Region::createRegionsFromMD(*F);
+ EXPECT_THAT(Regions[0]->insts(), testing::UnorderedElementsAre(T0));
+ EXPECT_THAT(Regions[1]->insts(), testing::UnorderedElementsAre(T1, T2));
+}
+
TEST_F(RegionTest, MetadataRoundTrip) {
parseIR(C, R"IR(
define i8 @foo(i8 %v0, i8 %v1) {
>From d155e23e476aa9bf06592ff2e92f1f54191689b2 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Fri, 20 Sep 2024 14:21:13 -0700
Subject: [PATCH 3/5] Add requested TODO
---
llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp | 1 +
1 file changed, 1 insertion(+)
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
index 55f08f424655b0..5f2c28484f62bc 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp
@@ -20,6 +20,7 @@ Region::~Region() {}
void Region::add(Instruction *I) {
Insts.insert(I);
+ // TODO: Consider tagging instructions lazily.
cast<llvm::Instruction>(I->Val)->setMetadata(MDKind, RegionMDN);
}
>From 5b06288da0913ec392fab4a047ed59930683535a Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Fri, 20 Sep 2024 14:38:02 -0700
Subject: [PATCH 4/5] removed extra whitespace at end of line
---
llvm/include/llvm/SandboxIR/SandboxIR.h | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 3af1af5e7718eb..7c89ea3db4f8cd 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -4402,7 +4402,7 @@ class Context {
friend class StructType; // For LLVMCtx.
friend class ::llvm::TargetExtType; // For LLVMCtx.
friend class Region; // For LLVMCtx.
-
+
Tracker IRTracker;
/// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
>From fafe0926ba2e26e6f1b537fb07300d4afdbb0491 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe at google.com>
Date: Fri, 20 Sep 2024 14:44:50 -0700
Subject: [PATCH 5/5] clang-format
---
llvm/include/llvm/SandboxIR/SandboxIR.h | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 7c89ea3db4f8cd..126a2ae4e47740 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -4398,10 +4398,10 @@ class Context {
friend class PointerType; // For LLVMCtx.
friend class CmpInst; // For LLVMCtx. TODO: cleanup when sandboxir::VectorType
// is complete
- friend class IntegerType; // For LLVMCtx.
- friend class StructType; // For LLVMCtx.
+ friend class IntegerType; // For LLVMCtx.
+ friend class StructType; // For LLVMCtx.
friend class ::llvm::TargetExtType; // For LLVMCtx.
- friend class Region; // For LLVMCtx.
+ friend class Region; // For LLVMCtx.
Tracker IRTracker;
More information about the llvm-commits
mailing list