[llvm] [orc-rt] Add allocation-action execution support. (PR #157244)
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Sat Sep 6 01:20:00 PDT 2025
https://github.com/lhames created https://github.com/llvm/llvm-project/pull/157244
This commit contains executor-side support for ORC allocation actions (see e50aea58d59).
An AllocAction is a function pointer with type
orc_rt_WrapperFunctionBuffer (*)(const char *ArgData, size_t ArgSize), along with an associated blob of argument bytes.
An AllocActionPair is a pair of AllocActions, one to be run at memory finalization time and another to be run at deallocation time.
The runFinalizeActions function can be used to run all non-null finalize actions in a sequence of AllocActionPairs, returning the corresponding sequence of deallocation actions on success.
The runDeallocActions function can be used to run a sequence of dealloc actions returned by runFinalizeActions.
>From 4de22b5a9c80031fb31dade0b77124a743bad1d1 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Sat, 6 Sep 2025 17:34:11 +1000
Subject: [PATCH] [orc-rt] Add allocation-action execution support.
This commit contains executor-side support for ORC allocation actions (see
e50aea58d59).
An AllocAction is a function pointer with type
orc_rt_WrapperFunctionBuffer (*)(const char *ArgData, size_t ArgSize),
along with an associated blob of argument bytes.
An AllocActionPair is a pair of AllocActions, one to be run at memory
finalization time and another to be run at deallocation time.
The runFinalizeActions function can be used to run all non-null finalize
actions in a sequence of AllocActionPairs, returning the corresponding
sequence of deallocation actions on success.
The runDeallocActions function can be used to run a sequence of
dealloc actions returned by runFinalizeActions.
---
orc-rt/include/CMakeLists.txt | 2 +
orc-rt/include/orc-rt/AllocAction.h | 105 ++++++++++++++++
orc-rt/include/orc-rt/SPSAllocAction.h | 46 +++++++
orc-rt/lib/executor/AllocAction.cpp | 55 +++++++++
orc-rt/lib/executor/CMakeLists.txt | 1 +
orc-rt/unittests/AllocActionTest.cpp | 165 +++++++++++++++++++++++++
orc-rt/unittests/CMakeLists.txt | 1 +
7 files changed, 375 insertions(+)
create mode 100644 orc-rt/include/orc-rt/AllocAction.h
create mode 100644 orc-rt/include/orc-rt/SPSAllocAction.h
create mode 100644 orc-rt/lib/executor/AllocAction.cpp
create mode 100644 orc-rt/unittests/AllocActionTest.cpp
diff --git a/orc-rt/include/CMakeLists.txt b/orc-rt/include/CMakeLists.txt
index c3295b3c781bf..bcbee39e647ea 100644
--- a/orc-rt/include/CMakeLists.txt
+++ b/orc-rt/include/CMakeLists.txt
@@ -3,6 +3,7 @@ set(ORC_RT_HEADERS
orc-rt-c/ExternC.h
orc-rt-c/WrapperFunction.h
orc-rt-c/orc-rt.h
+ orc-rt/AllocAction.h
orc-rt/BitmaskEnum.h
orc-rt/Compiler.h
orc-rt/Error.h
@@ -15,6 +16,7 @@ set(ORC_RT_HEADERS
orc-rt/WrapperFunction.h
orc-rt/ScopeExit.h
orc-rt/SimplePackedSerialization.h
+ orc-rt/SPSAllocAction.h
orc-rt/SPSWrapperFunction.h
orc-rt/bind.h
orc-rt/bit.h
diff --git a/orc-rt/include/orc-rt/AllocAction.h b/orc-rt/include/orc-rt/AllocAction.h
new file mode 100644
index 0000000000000..8362148e10fea
--- /dev/null
+++ b/orc-rt/include/orc-rt/AllocAction.h
@@ -0,0 +1,105 @@
+//===---------- AllocAction.h - Allocation action APIs ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// AllocAction and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_ALLOCACTION_H
+#define ORC_RT_ALLOCACTION_H
+
+#include "orc-rt/Error.h"
+#include "orc-rt/WrapperFunction.h"
+
+#include <vector>
+
+namespace orc_rt {
+namespace detail {
+
+template <typename Handler>
+struct AAHandlerTraits
+ : public AAHandlerTraits<
+ decltype(&std::remove_cv_t<std::remove_reference_t<Handler>>::
+ operator())> {};
+
+template <typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {
+ typedef std::tuple<ArgTs...> ArgTuple;
+};
+
+template <typename Class, typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...)>
+ : public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
+
+template <typename Class, typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...) const>
+ : public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
+
+} // namespace detail
+
+/// An AllocActionFn is a function that takes an argument blob and returns an
+/// empty WrapperFunctionBuffer on success, or an out-of-band error on failure.
+typedef orc_rt_WrapperFunctionBuffer (*AllocActionFn)(const char *ArgData,
+ size_t ArgSize);
+
+struct AllocActionFunction {
+
+ template <typename Deserializer, typename Handler>
+ static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
+ Deserializer &&D, Handler &&H) {
+ typename detail::AAHandlerTraits<Handler>::ArgTuple Args;
+ if (!D.deserialize(ArgData, ArgSize, Args))
+ return WrapperFunctionBuffer::createOutOfBandError(
+ "Could not deserialize allocation action argument buffer");
+
+ return std::apply(std::forward<Handler>(H), std::move(Args)).release();
+ }
+};
+
+/// An AllocAction is a pair of an AllocActionFn and an argument data buffer.
+struct AllocAction {
+ AllocAction() = default;
+ AllocAction(AllocActionFn AA, WrapperFunctionBuffer ArgData)
+ : AA(AA), ArgData(std::move(ArgData)) {}
+
+ [[nodiscard]] WrapperFunctionBuffer operator()() {
+ assert(AA && "Attempt to call null action");
+ return AA(ArgData.data(), ArgData.size());
+ }
+
+ explicit operator bool() const noexcept { return !!AA; }
+
+ AllocActionFn AA = nullptr;
+ WrapperFunctionBuffer ArgData;
+};
+
+/// An AllocActionPair is a pair of a Finalize action and a Dealloc action.
+struct AllocActionPair {
+ AllocAction Finalize;
+ AllocAction Dealloc;
+};
+
+/// Run the finalize actions in the given sequence.
+///
+/// On success, returns the list of deallocation actions to be run in reverse
+/// order at deallocation time.
+///
+/// On failure, runs deallocation actions associated with any previously
+/// successful finalize actions, then returns an error.
+///
+/// Both finalize and dealloc actions are permitted to be null (i.e. have a
+/// null action function) in which case they are ignored.
+[[nodiscard]] Expected<std::vector<AllocAction>>
+runFinalizeActions(std::vector<AllocActionPair> AAPs);
+
+/// Run the given deallocation actions in revwerse order.
+void runDeallocActions(std::vector<AllocAction> DAAs);
+
+} // namespace orc_rt
+
+#endif // ORC_RT_ALLOCACTION_H
diff --git a/orc-rt/include/orc-rt/SPSAllocAction.h b/orc-rt/include/orc-rt/SPSAllocAction.h
new file mode 100644
index 0000000000000..a1f9dc8175452
--- /dev/null
+++ b/orc-rt/include/orc-rt/SPSAllocAction.h
@@ -0,0 +1,46 @@
+//===---- SPSAllocAction.h - SPS-serialized AllocAction utils ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utilities for implementing allocation actions that take an SPS-serialized
+// argument buffer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_SPSALLOCACTION_H
+#define ORC_RT_SPSALLOCACTION_H
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/SimplePackedSerialization.h"
+
+namespace orc_rt {
+
+template <typename... SPSArgTs> struct AllocActionSPSDeserializer {
+ template <typename... ArgTs>
+ bool deserialize(const char *ArgData, size_t ArgSize, ArgTs &...Args) {
+ SPSInputBuffer IB(ArgData, ArgSize);
+ return SPSArgList<SPSArgTs...>::deserialize(IB, Args...);
+ }
+};
+
+/// Provides call and handle utilities to simplify writing and invocation of
+/// wrapper functions that use SimplePackedSerialization to serialize and
+/// deserialize their arguments and return values.
+template <typename... SPSArgTs> struct SPSAllocActionFunction {
+
+ template <typename Handler>
+ static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
+ Handler &&H) {
+ return AllocActionFunction::handle(
+ ArgData, ArgSize, AllocActionSPSDeserializer<SPSTuple<SPSArgTs...>>(),
+ std::forward<Handler>(H));
+ }
+};
+
+} // namespace orc_rt
+
+#endif // ORC_RT_SPSWRAPPERFUNCTION_H
diff --git a/orc-rt/lib/executor/AllocAction.cpp b/orc-rt/lib/executor/AllocAction.cpp
new file mode 100644
index 0000000000000..a4c7064bbfbb0
--- /dev/null
+++ b/orc-rt/lib/executor/AllocAction.cpp
@@ -0,0 +1,55 @@
+//===- AllocAction.cpp ----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// AllocAction and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/ScopeExit.h"
+
+namespace orc_rt {
+
+Expected<std::vector<AllocAction>>
+runFinalizeActions(std::vector<AllocActionPair> AAPs) {
+ std::vector<AllocAction> DeallocActions;
+ auto RunDeallocActions = make_scope_exit([&]() {
+ while (!DeallocActions.empty()) {
+ // TODO: Log errors from cleanup dealloc actions.
+ {
+ [[maybe_unused]] auto B = DeallocActions.back()();
+ }
+ DeallocActions.pop_back();
+ }
+ });
+
+ for (auto &AAP : AAPs) {
+ if (AAP.Finalize) {
+ auto B = AAP.Finalize();
+ if (const char *ErrMsg = B.getOutOfBandError())
+ return make_error<StringError>(ErrMsg);
+ }
+ if (AAP.Dealloc)
+ DeallocActions.push_back(std::move(AAP.Dealloc));
+ }
+
+ RunDeallocActions.release();
+ return DeallocActions;
+}
+
+void runDeallocActions(std::vector<AllocAction> DAAs) {
+ while (!DAAs.empty()) {
+ // TODO: Log errors from cleanup dealloc actions.
+ {
+ [[maybe_unused]] auto B = DAAs.back()();
+ }
+ DAAs.pop_back();
+ }
+}
+
+} // namespace orc_rt
diff --git a/orc-rt/lib/executor/CMakeLists.txt b/orc-rt/lib/executor/CMakeLists.txt
index 331954c2c556f..9213522e66a27 100644
--- a/orc-rt/lib/executor/CMakeLists.txt
+++ b/orc-rt/lib/executor/CMakeLists.txt
@@ -1,5 +1,6 @@
set(files
orc-rt-executor.cpp
+ AllocAction.cpp
RTTI.cpp
)
diff --git a/orc-rt/unittests/AllocActionTest.cpp b/orc-rt/unittests/AllocActionTest.cpp
new file mode 100644
index 0000000000000..46a3fdb6d79cc
--- /dev/null
+++ b/orc-rt/unittests/AllocActionTest.cpp
@@ -0,0 +1,165 @@
+//===- AllocActionTest.cpp ------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for orc-rt's AllocAction.h APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/ExecutorAddress.h"
+#include "orc-rt/SPSAllocAction.h"
+
+#include "SimplePackedSerializationTestUtils.h"
+#include "gtest/gtest.h"
+
+using namespace orc_rt;
+
+TEST(AllocActionTest, DefaultConstruct) {
+ AllocAction AA;
+ EXPECT_FALSE(AA);
+}
+
+static orc_rt_WrapperFunctionBuffer noopAction(const char *ArgData,
+ size_t ArgSize) {
+ return WrapperFunctionBuffer().release();
+}
+
+TEST(AllocActionTest, ConstructWithAction) {
+ AllocAction AA(noopAction, WrapperFunctionBuffer());
+ EXPECT_TRUE(AA);
+}
+
+// Increments int via pointer.
+static orc_rt_WrapperFunctionBuffer
+increment_sps_allocaction(const char *ArgData, size_t ArgSize) {
+ return SPSAllocActionFunction<SPSExecutorAddr>::handle(
+ ArgData, ArgSize,
+ [](ExecutorAddr IntPtr) {
+ *IntPtr.toPtr<int *>() += 1;
+ return WrapperFunctionBuffer();
+ })
+ .release();
+}
+
+// Increments int via pointer.
+static orc_rt_WrapperFunctionBuffer
+decrement_sps_allocaction(const char *ArgData, size_t ArgSize) {
+ return SPSAllocActionFunction<SPSExecutorAddr>::handle(
+ ArgData, ArgSize,
+ [](ExecutorAddr IntPtr) {
+ *IntPtr.toPtr<int *>() -= 1;
+ return WrapperFunctionBuffer();
+ })
+ .release();
+}
+
+template <typename T>
+static WrapperFunctionBuffer makeExecutorAddrBuffer(T *P) {
+ return *spsSerialize<SPSArgList<SPSExecutorAddr>>(ExecutorAddr::fromPtr(P));
+}
+
+TEST(AllocActionTest, RunBasicAction) {
+ int Val = 0;
+ AllocAction IncVal(increment_sps_allocaction, makeExecutorAddrBuffer(&Val));
+ EXPECT_TRUE(IncVal);
+ auto B = IncVal();
+ EXPECT_TRUE(B.empty());
+ EXPECT_EQ(Val, 1);
+}
+
+TEST(AllocActionTest, RunFinalizationActionsComplete) {
+ int Val = 0;
+
+ std::vector<AllocActionPair> InitialActions;
+
+ auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+
+ auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+ EXPECT_EQ(Val, 2);
+
+ runDeallocActions(std::move(DeallocActions));
+
+ EXPECT_EQ(Val, 0);
+}
+
+static orc_rt_WrapperFunctionBuffer fail_sps_allocaction(const char *ArgData,
+ size_t ArgSize) {
+ return WrapperFunctionBuffer::createOutOfBandError("failed").release();
+}
+
+TEST(AllocActionTest, RunFinalizeActionsFail) {
+ int Val = 0;
+
+ std::vector<AllocActionPair> InitialActions;
+
+ auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+ InitialActions.push_back({{fail_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+
+ auto DeallocActions = runFinalizeActions(std::move(InitialActions));
+
+ if (DeallocActions) {
+ ADD_FAILURE() << "Failed to report error from runFinalizeActions";
+ return;
+ }
+
+ EXPECT_EQ(toString(DeallocActions.takeError()), std::string("failed"));
+
+ // Check that we ran the decrement corresponding to the first increment.
+ EXPECT_EQ(Val, 0);
+}
+
+TEST(AllocActionTest, RunFinalizeActionsNullFinalize) {
+ int Val = 0;
+
+ std::vector<AllocActionPair> InitialActions;
+
+ auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+ InitialActions.push_back({{nullptr, WrapperFunctionBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+
+ auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+ // Both dealloc actions should be included in the returned list, despite one
+ // of them having a null finalize action.
+ EXPECT_EQ(DeallocActions.size(), 2U);
+
+ runDeallocActions(std::move(DeallocActions));
+
+ EXPECT_EQ(Val, -1);
+}
+
+TEST(AllocActionTest, RunFinalizeActionsNullDealloc) {
+ int Val = 0;
+
+ std::vector<AllocActionPair> InitialActions;
+
+ auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {decrement_sps_allocaction, MakeArgBuffer()}});
+ InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+ {nullptr, WrapperFunctionBuffer()}});
+
+ auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+ // Null dealloc actions should be filtered out of the returned list.
+ EXPECT_EQ(DeallocActions.size(), 1U);
+
+ runDeallocActions(std::move(DeallocActions));
+
+ EXPECT_EQ(Val, 1);
+}
diff --git a/orc-rt/unittests/CMakeLists.txt b/orc-rt/unittests/CMakeLists.txt
index 4ebc5277a5f15..f8556401a9350 100644
--- a/orc-rt/unittests/CMakeLists.txt
+++ b/orc-rt/unittests/CMakeLists.txt
@@ -12,6 +12,7 @@ function(add_orc_rt_unittest test_dirname)
endfunction()
add_orc_rt_unittest(CoreTests
+ AllocActionTest.cpp
BitmaskEnumTest.cpp
CommonTestUtils.cpp
ErrorTest.cpp
More information about the llvm-commits
mailing list