[SandboxIR][Region] Auxiliary vector metadata now requires a region

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));
+}