[llvm] 62e6fe8 - Reapply "[SandboxIR] Callback registration now allows specifying order" (#221099)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Sep 5 18:18:50 PDT 2026
Author: vporpo
Date: 2026-09-05T18:18:46-07:00
New Revision: 62e6fe8bbab320f2517c004e6087be489f3ecb88
URL: https://github.com/llvm/llvm-project/commit/62e6fe8bbab320f2517c004e6087be489f3ecb88
DIFF: https://github.com/llvm/llvm-project/commit/62e6fe8bbab320f2517c004e6087be489f3ecb88.diff
LOG: Reapply "[SandboxIR] Callback registration now allows specifying order" (#221099)
This reverts commit f111407c0502851409a887c8e5dea4a3ec190d4e.
Should fix the sanitizer issue.
Added:
Modified:
llvm/include/llvm/SandboxIR/Context.h
llvm/lib/SandboxIR/Context.cpp
llvm/unittests/SandboxIR/SandboxIRTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/SandboxIR/Context.h b/llvm/include/llvm/SandboxIR/Context.h
index b8b47afe65e11..dfc0df0efd56e 100644
--- a/llvm/include/llvm/SandboxIR/Context.h
+++ b/llvm/include/llvm/SandboxIR/Context.h
@@ -62,6 +62,7 @@ class LLVM_ABI Context {
public:
CallbackID() = default;
+ bool operator==(const CallbackID &Other) const { return Val == Other.Val; }
friend class Context;
friend struct DenseMapInfo<CallbackID>;
};
@@ -95,18 +96,18 @@ class LLVM_ABI Context {
/// Type objects.
DenseMap<llvm::Type *, std::unique_ptr<Type, TypeDeleter>> LLVMTypeToTypeMap;
- /// Callbacks called when an IR instruction is about to get erased. Keys are
- /// used as IDs for deregistration.
- MapVector<CallbackID, EraseInstrCallback> EraseInstrCallbacks;
- /// Callbacks called when an IR instruction is about to get created. Keys are
- /// used as IDs for deregistration.
- MapVector<CallbackID, CreateInstrCallback> CreateInstrCallbacks;
- /// Callbacks called when an IR instruction is about to get moved. Keys are
- /// used as IDs for deregistration.
- MapVector<CallbackID, MoveInstrCallback> MoveInstrCallbacks;
- /// Callbacks called when a Use gets its source set. Keys are used as IDs for
- /// deregistration.
- MapVector<CallbackID, SetUseCallback> SetUseCallbacks;
+ /// Callbacks called when an IR instruction is about to get erased. CallbackID
+ /// is used as an identifier for deregistration.
+ SmallVector<std::pair<CallbackID, EraseInstrCallback>> EraseInstrCallbacks;
+ /// Callbacks called when an IR instruction is about to get created.
+ /// CallbackID is used as an identifier for deregistration.
+ SmallVector<std::pair<CallbackID, CreateInstrCallback>> CreateInstrCallbacks;
+ /// Callbacks called when an IR instruction is about to get moved. CallbackID
+ /// is used as an identifier for deregistration.
+ SmallVector<std::pair<CallbackID, MoveInstrCallback>> MoveInstrCallbacks;
+ /// Callbacks called when a Use gets its source set. CallbackID is used as an
+ /// identifier for deregistration.
+ SmallVector<std::pair<CallbackID, SetUseCallback>> SetUseCallbacks;
/// A counter used for assigning callback IDs during registration. The same
/// counter is used for all kinds of callbacks so we can detect mismatched
@@ -277,29 +278,69 @@ class LLVM_ABI Context {
/// \Returns the number of values registered with Context.
size_t getNumValues() const { return LLVMValueToValueMap.size(); }
+private:
+ // An arbitrary limit, to check for accidental misuse. We expect a small
+ // number of callbacks to be registered at a time, but we can increase this
+ // number if we discover we needed more.
+ [[maybe_unused]] static constexpr int MaxRegisteredCallbacks = 16;
+
+ template <typename CBT, typename CBVecT>
+ CallbackID registerCallbackCommon(CBT CB, CBVecT &CBVec,
+ std::optional<CallbackID> BeforeID) {
+ assert(CBVec.size() <= MaxRegisteredCallbacks &&
+ "EraseInstrCallbacks size limit exceeded");
+ auto BeforeIt = CBVec.end();
+ if (BeforeID) {
+ BeforeIt = find_if(CBVec, [BeforeID](const auto &Pair) {
+ return Pair.first == *BeforeID;
+ });
+ assert(BeforeIt != CBVec.end() && "Not found!");
+ }
+ CallbackID ID{NextCallbackID++};
+ CBVec.insert(BeforeIt, {ID, std::move(CB)});
+ return ID;
+ }
+
+public:
/// Register a callback that gets called when a SandboxIR instruction is about
/// to be removed from its parent. Note that this will also be called when
/// reverting the creation of an instruction.
+ /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
+ /// so it will be called first once the callbacks get executed.
/// \Returns a callback ID for later deregistration.
- CallbackID registerEraseInstrCallback(EraseInstrCallback CB);
+ CallbackID
+ registerEraseInstrCallback(EraseInstrCallback CB,
+ std::optional<CallbackID> BeforeID = std::nullopt);
void unregisterEraseInstrCallback(CallbackID ID);
/// Register a callback that gets called right after a SandboxIR instruction
/// is created. Note that this will also be called when reverting the removal
/// of an instruction.
+ /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
+ /// so it will be called first once the callbacks get executed.
/// \Returns a callback ID for later deregistration.
- CallbackID registerCreateInstrCallback(CreateInstrCallback CB);
+ CallbackID registerCreateInstrCallback(
+ CreateInstrCallback CB,
+ std::optional<CallbackID> BeforeID = std::nullopt);
void unregisterCreateInstrCallback(CallbackID ID);
/// Register a callback that gets called when a SandboxIR instruction is about
/// to be moved. Note that this will also be called when reverting a move.
+ /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
+ /// so it will be called first once the callbacks get executed.
/// \Returns a callback ID for later deregistration.
- CallbackID registerMoveInstrCallback(MoveInstrCallback CB);
+ CallbackID
+ registerMoveInstrCallback(MoveInstrCallback CB,
+ std::optional<CallbackID> BeforeID = std::nullopt);
void unregisterMoveInstrCallback(CallbackID ID);
/// Register a callback that gets called when a Use gets set.
+ /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
+ /// so it will be called first once the callbacks get executed.
/// \Returns a callback ID for later deregistration.
- CallbackID registerSetUseCallback(SetUseCallback CB);
+ CallbackID
+ registerSetUseCallback(SetUseCallback CB,
+ std::optional<CallbackID> BeforeID = std::nullopt);
void unregisterSetUseCallback(CallbackID ID);
};
diff --git a/llvm/lib/SandboxIR/Context.cpp b/llvm/lib/SandboxIR/Context.cpp
index bf2ecda344824..20dc3e035d355 100644
--- a/llvm/lib/SandboxIR/Context.cpp
+++ b/llvm/lib/SandboxIR/Context.cpp
@@ -712,8 +712,8 @@ Module *Context::createModule(llvm::Module *LLVMM) {
}
void Context::runEraseInstrCallbacks(Instruction *I) {
- for (const auto &CBEntry : EraseInstrCallbacks)
- CBEntry.second(I);
+ for (const auto &[ID, CB] : EraseInstrCallbacks)
+ CB(I);
}
void Context::runCreateInstrCallbacks(Instruction *I) {
@@ -731,62 +731,56 @@ void Context::runSetUseCallbacks(const Use &U, Value *NewSrc) {
CBEntry.second(U, NewSrc);
}
-// An arbitrary limit, to check for accidental misuse. We expect a small number
-// of callbacks to be registered at a time, but we can increase this number if
-// we discover we needed more.
-[[maybe_unused]] static constexpr int MaxRegisteredCallbacks = 16;
-
-Context::CallbackID Context::registerEraseInstrCallback(EraseInstrCallback CB) {
- assert(EraseInstrCallbacks.size() <= MaxRegisteredCallbacks &&
- "EraseInstrCallbacks size limit exceeded");
- CallbackID ID{NextCallbackID++};
- EraseInstrCallbacks[ID] = std::move(CB);
- return ID;
+Context::CallbackID
+Context::registerEraseInstrCallback(EraseInstrCallback CB,
+ std::optional<CallbackID> BeforeID) {
+ return registerCallbackCommon(CB, EraseInstrCallbacks, BeforeID);
}
void Context::unregisterEraseInstrCallback(CallbackID ID) {
- [[maybe_unused]] bool Erased = EraseInstrCallbacks.erase(ID);
- assert(Erased &&
+ [[maybe_unused]] auto It = find_if(
+ EraseInstrCallbacks, [ID](const auto &Pair) { return Pair.first == ID; });
+ assert(It != EraseInstrCallbacks.end() &&
"Callback ID not found in EraseInstrCallbacks during deregistration");
+ EraseInstrCallbacks.erase(It);
}
Context::CallbackID
-Context::registerCreateInstrCallback(CreateInstrCallback CB) {
- assert(CreateInstrCallbacks.size() <= MaxRegisteredCallbacks &&
- "CreateInstrCallbacks size limit exceeded");
- CallbackID ID{NextCallbackID++};
- CreateInstrCallbacks[ID] = std::move(CB);
- return ID;
+Context::registerCreateInstrCallback(CreateInstrCallback CB,
+ std::optional<CallbackID> BeforeID) {
+ return registerCallbackCommon(CB, CreateInstrCallbacks, BeforeID);
}
void Context::unregisterCreateInstrCallback(CallbackID ID) {
- [[maybe_unused]] bool Erased = CreateInstrCallbacks.erase(ID);
- assert(Erased &&
+ auto It = find_if(CreateInstrCallbacks,
+ [ID](const auto &Pair) { return Pair.first == ID; });
+ assert(It != CreateInstrCallbacks.end() &&
"Callback ID not found in CreateInstrCallbacks during deregistration");
+ CreateInstrCallbacks.erase(It);
}
-Context::CallbackID Context::registerMoveInstrCallback(MoveInstrCallback CB) {
- assert(MoveInstrCallbacks.size() <= MaxRegisteredCallbacks &&
- "MoveInstrCallbacks size limit exceeded");
- CallbackID ID{NextCallbackID++};
- MoveInstrCallbacks[ID] = std::move(CB);
- return ID;
+Context::CallbackID
+Context::registerMoveInstrCallback(MoveInstrCallback CB,
+ std::optional<CallbackID> BeforeID) {
+ return registerCallbackCommon(CB, MoveInstrCallbacks, BeforeID);
}
void Context::unregisterMoveInstrCallback(CallbackID ID) {
- [[maybe_unused]] bool Erased = MoveInstrCallbacks.erase(ID);
- assert(Erased &&
+ [[maybe_unused]] auto It = find_if(
+ MoveInstrCallbacks, [ID](const auto &Pair) { return Pair.first == ID; });
+ assert(It != MoveInstrCallbacks.end() &&
"Callback ID not found in MoveInstrCallbacks during deregistration");
+ MoveInstrCallbacks.erase(It);
}
-Context::CallbackID Context::registerSetUseCallback(SetUseCallback CB) {
- assert(SetUseCallbacks.size() <= MaxRegisteredCallbacks &&
- "SetUseCallbacks size limit exceeded");
- CallbackID ID{NextCallbackID++};
- SetUseCallbacks[ID] = std::move(CB);
- return ID;
+Context::CallbackID
+Context::registerSetUseCallback(SetUseCallback CB,
+ std::optional<CallbackID> BeforeID) {
+ return registerCallbackCommon(CB, SetUseCallbacks, BeforeID);
}
void Context::unregisterSetUseCallback(CallbackID ID) {
- [[maybe_unused]] bool Erased = SetUseCallbacks.erase(ID);
- assert(Erased &&
+ [[maybe_unused]] auto It = find_if(
+ SetUseCallbacks, [ID](const auto &Pair) { return Pair.first == ID; });
+ assert(It != SetUseCallbacks.end() &&
"Callback ID not found in SetUseCallbacks during deregistration");
+ SetUseCallbacks.erase(It);
}
} // namespace llvm::sandboxir
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index 03b3667abc333..4110cd8e2a2b7 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -6343,6 +6343,104 @@ TEST_F(SandboxIRTest, InstructionCallbacks) {
EXPECT_THAT(Moved, testing::IsEmpty());
}
+TEST_F(SandboxIRTest, InstructionCallbacks_BeforeID) {
+ parseIR(C, R"IR(
+ define void @foo(i8 %v0, ptr %ptr) {
+ %add0 = add i8 %v0, %v0
+ ret void
+ }
+ )IR");
+ Function &LLVMF = *M->getFunction("foo");
+ sandboxir::Context Ctx(C);
+
+ auto &F = *Ctx.createFunction(&LLVMF);
+ auto &BB = *F.begin();
+ sandboxir::Argument *Val = F.getArg(0);
+ sandboxir::Argument *Ptr = F.getArg(1);
+ auto It = BB.begin();
+ sandboxir::Instruction *Add0 = &*It++;
+ sandboxir::Instruction *Ret = &*It++;
+ auto *Arg0 = F.getArg(0);
+
+ // Callbacks write to this vector.
+ SmallVector<unsigned> CBs;
+ {
+ // Check EraseInstr callbacks.
+ CBs.clear();
+ // The first callback.
+ auto CB0 = Ctx.registerEraseInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(0); });
+ // This callback is placed after the first.
+ [[maybe_unused]] auto CB1 = Ctx.registerEraseInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(1); });
+ // This should insert this callback before the first.
+ [[maybe_unused]] auto CB2 = Ctx.registerEraseInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(2); },
+ /*BeforeID=*/CB0);
+ Ctx.save();
+ Ret->eraseFromParent();
+ EXPECT_THAT(CBs, testing::ElementsAre(2, 0, 1));
+ Ctx.revert();
+ }
+ {
+ // Check CreateInstr callbacks.
+ CBs.clear();
+ // The first callback.
+ auto CB0 = Ctx.registerCreateInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(0); });
+ // This callback is placed after the first.
+ [[maybe_unused]] auto CB1 = Ctx.registerCreateInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(1); });
+ // This should insert this callback before the first.
+ [[maybe_unused]] auto CB2 = Ctx.registerCreateInstrCallback(
+ [&CBs](sandboxir::Instruction *I) { CBs.push_back(2); },
+ /*BeforeID=*/CB0);
+ Ctx.save();
+ sandboxir::StoreInst::create(Val, Ptr, /*Align=*/std::nullopt,
+ Ret->getIterator(), Ctx);
+ EXPECT_THAT(CBs, testing::ElementsAre(2, 0, 1));
+ Ctx.revert();
+ }
+ {
+ // Check MoveInstr callbacks.
+ CBs.clear();
+ // The first callback.
+ auto CB0 = Ctx.registerMoveInstrCallback(
+ [&CBs](sandboxir::Instruction *I, const sandboxir::BBIterator &Where) {
+ CBs.push_back(10);
+ });
+ // This should insert this callback before the first.
+ [[maybe_unused]] auto CB1 = Ctx.registerMoveInstrCallback(
+ [&CBs](sandboxir::Instruction *I, const sandboxir::BBIterator &Where) {
+ CBs.push_back(11);
+ },
+ /*BeforeID=*/CB0);
+ Ctx.save();
+ Ret->moveBefore(Add0);
+ EXPECT_THAT(CBs, testing::ElementsAre(11, 10));
+ Ctx.revert();
+ }
+ {
+ // Check SetUse callbacks.
+ CBs.clear();
+ // The first callback.
+ auto CB0 = Ctx.registerSetUseCallback(
+ [&CBs](sandboxir::Use U, sandboxir::Value *NewSrc) {
+ CBs.push_back(100);
+ });
+ // This should insert this callback before the first.
+ [[maybe_unused]] auto CB1 = Ctx.registerSetUseCallback(
+ [&CBs](sandboxir::Use U, sandboxir::Value *NewSrc) {
+ CBs.push_back(101);
+ },
+ /*BeforeID=*/CB0);
+ Ctx.save();
+ Add0->setOperand(0, Arg0);
+ EXPECT_THAT(CBs, testing::ElementsAre(101, 100));
+ Ctx.revert();
+ }
+}
+
// Check callbacks when we set a Use.
TEST_F(SandboxIRTest, SetUseCallbacks) {
parseIR(C, R"IR(
More information about the llvm-commits
mailing list