From 4de22b5a9c80031fb31dade0b77124a743bad1d1 Mon Sep 17 00:00:00 2001 From: Lang Hames 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 + +namespace orc_rt { +namespace detail { + +template +struct AAHandlerTraits + : public AAHandlerTraits< + decltype(&std::remove_cv_t>:: + operator())> {}; + +template +struct AAHandlerTraits { + typedef std::tuple ArgTuple; +}; + +template +struct AAHandlerTraits + : public AAHandlerTraits {}; + +template +struct AAHandlerTraits + : public AAHandlerTraits {}; + +} // 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 + static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize, + Deserializer &&D, Handler &&H) { + typename detail::AAHandlerTraits::ArgTuple Args; + if (!D.deserialize(ArgData, ArgSize, Args)) + return WrapperFunctionBuffer::createOutOfBandError( + "Could not deserialize allocation action argument buffer"); + + return std::apply(std::forward(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> +runFinalizeActions(std::vector AAPs); + +/// Run the given deallocation actions in revwerse order. +void runDeallocActions(std::vector 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 struct AllocActionSPSDeserializer { + template + bool deserialize(const char *ArgData, size_t ArgSize, ArgTs &...Args) { + SPSInputBuffer IB(ArgData, ArgSize); + return SPSArgList::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 struct SPSAllocActionFunction { + + template + static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize, + Handler &&H) { + return AllocActionFunction::handle( + ArgData, ArgSize, AllocActionSPSDeserializer>(), + std::forward(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> +runFinalizeActions(std::vector AAPs) { + std::vector 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(ErrMsg); + } + if (AAP.Dealloc) + DeallocActions.push_back(std::move(AAP.Dealloc)); + } + + RunDeallocActions.release(); + return DeallocActions; +} + +void runDeallocActions(std::vector 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::handle( + ArgData, ArgSize, + [](ExecutorAddr IntPtr) { + *IntPtr.toPtr() += 1; + return WrapperFunctionBuffer(); + }) + .release(); +} + +// Increments int via pointer. +static orc_rt_WrapperFunctionBuffer +decrement_sps_allocaction(const char *ArgData, size_t ArgSize) { + return SPSAllocActionFunction::handle( + ArgData, ArgSize, + [](ExecutorAddr IntPtr) { + *IntPtr.toPtr() -= 1; + return WrapperFunctionBuffer(); + }) + .release(); +} + +template +static WrapperFunctionBuffer makeExecutorAddrBuffer(T *P) { + return *spsSerialize>(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 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 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 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 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