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

via llvm-commits llvm-commits at lists.llvm.org
Thu Sep 19 16:53:11 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Jorge Gorbe Moya (slackito)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/109353.diff


4 Files Affected:

- (modified) llvm/include/llvm/SandboxIR/SandboxIR.h (+4) 
- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Region.h (+9) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Region.cpp (+37-2) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/RegionTest.cpp (+29-2) 


``````````diff
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
+}

``````````

</details>


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


More information about the llvm-commits mailing list