[SandboxIR][Region] Auxiliary vector metadata now requires a region (#137443)

The auxiliary vector is represented by the !sandboxaux metadata. But
until now adding an instruction to the aux vector would not
automatically add it to the region (i.e., it would not mark it with
!sandboxvec). This behavior was problematic when generating regions from
metadata, because you wouldn't know which region owns the auxiliary
vector.

This patch fixes this issue: We now require that an instruction with
!sandboxaux also defines its region with !sandboxvec.

llvm/include/llvm/SandboxIR/Region.h

@@ -114,8 +114,18 @@ class Region {
   /// ID (for later deregistration) of the "erase instruction" callback.
   Context::CallbackID EraseInstCB;
 
-  /// Adds I to the set.
-  void add(Instruction *I);
+  /// Adds \p I to the set but also don't track the instruction's score if \p
+  /// IgnoreCost is true. Only to be used when adding an instruction to the
+  /// auxiliary vector.
+  /// NOTE: When an instruction is added to the region we track it cost in the
+  /// scoreboard, which currently resides in the region class. However, when we
+  /// add an instruction to the auxiliary vector it does get tagged as being a
+  /// member of the region (for ownership reasons), but its cost does not get
+  /// counted because the instruction hasn't been added in the "normal" way.
+  void addImpl(Instruction *I, bool IgnoreCost);
+  /// Adds I to the set. This is the main API for adding an instruction to the
+  /// region.
+  void add(Instruction *I) { addImpl(I, /*IgnoreCost=*/false); }
   /// Removes I from the set.
   void remove(Instruction *I);
   friend class Context; // The callbacks need to call add() and remove().

llvm/lib/SandboxIR/Region.cpp

@@ -51,12 +51,13 @@ Region::~Region() {
   Ctx.unregisterEraseInstrCallback(EraseInstCB);
 }
 
-void Region::add(Instruction *I) {
+void Region::addImpl(Instruction *I, bool IgnoreCost) {
   Insts.insert(I);
   // TODO: Consider tagging instructions lazily.
   cast<llvm::Instruction>(I->Val)->setMetadata(MDKind, RegionMDN);
-  // Keep track of the instruction cost.
-  Scoreboard.add(I);
+  if (!IgnoreCost)
+    // Keep track of the instruction cost.
+    Scoreboard.add(I);
 }
 
 void Region::setAux(ArrayRef<Instruction *> Aux) {
@@ -69,6 +70,8 @@ void Region::setAux(ArrayRef<Instruction *> Aux) {
            "Instruction already in Aux!");
     cast<llvm::Instruction>(I->Val)->setMetadata(
         AuxMDKind, MDNode::get(LLVMCtx, ConstantAsMetadata::get(IdxC)));
+    // Aux instrs should always be in a region.
+    addImpl(I, /*DontTrackCost=*/true);
   }
 }
 
@@ -84,6 +87,8 @@ void Region::setAux(unsigned Idx, Instruction *I) {
       Aux[Idx] = nullptr;
   }
   Aux[Idx] = I;
+  // Aux instrs should always be in a region.
+  addImpl(I, /*DontTrackCost=*/true);
 }
 
 void Region::dropAuxMetadata(Instruction *I) {
@@ -151,8 +156,8 @@ Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) {
   for (BasicBlock &BB : F) {
     for (Instruction &Inst : BB) {
       auto *LLVMI = cast<llvm::Instruction>(Inst.Val);
+      Region *R = nullptr;
       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));
@@ -161,14 +166,17 @@ Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) {
         } 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);
+        R->addImpl(&Inst, /*IgnoreCost=*/true);
+      }
+      if (auto *AuxMDN = LLVMI->getMetadata(AuxMDKind)) {
+        llvm::Constant *IdxC =
+            dyn_cast<ConstantAsMetadata>(AuxMDN->getOperand(0))->getValue();
+        auto Idx = cast<llvm::ConstantInt>(IdxC)->getSExtValue();
+        if (R == nullptr) {
+          errs() << "No region specified for Aux: '" << *LLVMI << "'\n";
+          exit(1);
         }
+        R->setAux(Idx, &Inst);
       }
     }
   }

llvm/unittests/SandboxIR/RegionTest.cpp

@@ -429,6 +429,22 @@ define void @foo(i8 %v) {
   }
 }
 
+TEST_F(RegionTest, AuxWithoutRegion) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v) {
+  %Add0 = add i8 %v, 0, !sandboxaux !0
+  ret void
+}
+!0 = !{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), "No region.*");
+#endif
+}
+
 TEST_F(RegionTest, AuxRoundTrip) {
   parseIR(C, R"IR(
 define i8 @foo(i8 %v0, i8 %v1) {
@@ -461,3 +477,36 @@ define i8 @foo(i8 %v0, i8 %v1) {
 #endif
   EXPECT_THAT(Rgn.getAux(), testing::ElementsAre(T1, T0));
 }
+
+// Same as before but only add instructions to aux. They should get added too
+// the region too automatically.
+TEST_F(RegionTest, AuxOnlyRoundTrip) {
+  parseIR(C, R"IR(
+define void @foo(i8 %v) {
+  %add0 = add i8 %v, 0
+  %add1 = add i8 %v, 1
+  ret void
+}
+)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 *Add0 = cast<sandboxir::Instruction>(&*It++);
+  auto *Add1 = cast<sandboxir::Instruction>(&*It++);
+
+  sandboxir::Region Rgn(Ctx, *TTI);
+#ifndef NDEBUG
+  EXPECT_DEATH(Rgn.setAux({Add0, Add0}), ".*already.*");
+#endif
+  Rgn.setAux({Add1, Add0});
+
+  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(Add1, Add0));
+}