[SandboxIR][Region] Implement an auxiliary vector in Region (#126376)

vporpo · web-flow · commit c9ba2b08d57f · 2025-02-13T13:51:55.000-08:00
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.
diff --git 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
@@ -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
@@ -292,3 +292,131 @@ 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
+}
+
+// Check that we get an assertion failure if we try to set the same index more
+// than once.
+TEST_F(RegionTest, AuxSameIndex) {
+  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 !2
+  ret void
+}
+
+!0 = distinct !{!"sandboxregion"}
+!2 = !{i32 0}
+)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), ".*already.*");
+#endif // NDEBUG
+}
+
+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));
+}
