[llvm] [SandboxVec] Tag insts in a Region with metadata. (PR #109353)

Jorge Gorbe Moya via llvm-commits llvm-commits at lists.llvm.org
Thu Sep 19 16:52:37 PDT 2024


https://github.com/slackito created https://github.com/llvm/llvm-project/pull/109353

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.

>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] [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
+}



More information about the llvm-commits mailing list