[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.

Author: vporpo

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 = "sbaux";
 
   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(); }

llvm/lib/SandboxIR/Region.cpp

@@ -57,6 +57,39 @@ 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) {
+  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 +111,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 +135,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;
 }
 

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, !sbaux !2
+  %t1 = add i8 %v, 1, !sandboxvec !0, !sbaux !3
+  %t2 = add i8 %v, 2, !sandboxvec !1
+  %t3 = add i8 %v, 3, !sandboxvec !1, !sbaux !2
+  %t4 = add i8 %v, 4, !sandboxvec !1, !sbaux !4
+  %t5 = add i8 %v, 5, !sandboxvec !1, !sbaux !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("sbaux"));
+  EXPECT_TRUE(LLVMI1->getMetadata("sbaux"));
+  Regions[0]->clearAux();
+  EXPECT_TRUE(Regions[0]->getAux().empty());
+  EXPECT_FALSE(LLVMI0->getMetadata("sbaux"));
+  EXPECT_FALSE(LLVMI1->getMetadata("sbaux"));
+}
+
+// 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, !sbaux !2
+  %t1 = add i8 %v, 1, !sandboxvec !0, !sbaux !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));
+}