Skip to content

Commit 88a5429

Browse files
lhameslhames
authored
[orc-rt] Add allocation-action execution support. (#157244)
This commit contains executor-side support for ORC allocation actions (see e50aea5). 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.
1 parent 609bcd5 commit 88a5429

File tree

7 files changed

+375
-0
lines changed

7 files changed

+375
-0
lines changed

‎orc-rt/include/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ set(ORC_RT_HEADERS
33
orc-rt-c/ExternC.h
44
orc-rt-c/WrapperFunction.h
55
orc-rt-c/orc-rt.h
6+
orc-rt/AllocAction.h
67
orc-rt/BitmaskEnum.h
78
orc-rt/Compiler.h
89
orc-rt/Error.h
@@ -15,6 +16,7 @@ set(ORC_RT_HEADERS
1516
orc-rt/WrapperFunction.h
1617
orc-rt/ScopeExit.h
1718
orc-rt/SimplePackedSerialization.h
19+
orc-rt/SPSAllocAction.h
1820
orc-rt/SPSWrapperFunction.h
1921
orc-rt/bind.h
2022
orc-rt/bit.h

‎orc-rt/include/orc-rt/AllocAction.h‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
//===---------- AllocAction.h - Allocation action APIs ----------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// AllocAction and related APIs.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef ORC_RT_ALLOCACTION_H
14+
#define ORC_RT_ALLOCACTION_H
15+
16+
#include "orc-rt/Error.h"
17+
#include "orc-rt/WrapperFunction.h"
18+
19+
#include <vector>
20+
21+
namespace orc_rt {
22+
namespace detail {
23+
24+
template <typename Handler>
25+
struct AAHandlerTraits
26+
: public AAHandlerTraits<
27+
decltype(&std::remove_cv_t<std::remove_reference_t<Handler>>::
28+
operator())> {};
29+
30+
template <typename... ArgTs>
31+
struct AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {
32+
typedef std::tuple<ArgTs...> ArgTuple;
33+
};
34+
35+
template <typename Class, typename... ArgTs>
36+
struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...)>
37+
: public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
38+
39+
template <typename Class, typename... ArgTs>
40+
struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...) const>
41+
: public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
42+
43+
} // namespace detail
44+
45+
/// An AllocActionFn is a function that takes an argument blob and returns an
46+
/// empty WrapperFunctionBuffer on success, or an out-of-band error on failure.
47+
typedef orc_rt_WrapperFunctionBuffer (*AllocActionFn)(const char *ArgData,
48+
size_t ArgSize);
49+
50+
struct AllocActionFunction {
51+
52+
template <typename Deserializer, typename Handler>
53+
static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
54+
Deserializer &&D, Handler &&H) {
55+
typename detail::AAHandlerTraits<Handler>::ArgTuple Args;
56+
if (!D.deserialize(ArgData, ArgSize, Args))
57+
return WrapperFunctionBuffer::createOutOfBandError(
58+
"Could not deserialize allocation action argument buffer");
59+
60+
return std::apply(std::forward<Handler>(H), std::move(Args)).release();
61+
}
62+
};
63+
64+
/// An AllocAction is a pair of an AllocActionFn and an argument data buffer.
65+
struct AllocAction {
66+
AllocAction() = default;
67+
AllocAction(AllocActionFn AA, WrapperFunctionBuffer ArgData)
68+
: AA(AA), ArgData(std::move(ArgData)) {}
69+
70+
[[nodiscard]] WrapperFunctionBuffer operator()() {
71+
assert(AA && "Attempt to call null action");
72+
return AA(ArgData.data(), ArgData.size());
73+
}
74+
75+
explicit operator bool() const noexcept { return !!AA; }
76+
77+
AllocActionFn AA = nullptr;
78+
WrapperFunctionBuffer ArgData;
79+
};
80+
81+
/// An AllocActionPair is a pair of a Finalize action and a Dealloc action.
82+
struct AllocActionPair {
83+
AllocAction Finalize;
84+
AllocAction Dealloc;
85+
};
86+
87+
/// Run the finalize actions in the given sequence.
88+
///
89+
/// On success, returns the list of deallocation actions to be run in reverse
90+
/// order at deallocation time.
91+
///
92+
/// On failure, runs deallocation actions associated with any previously
93+
/// successful finalize actions, then returns an error.
94+
///
95+
/// Both finalize and dealloc actions are permitted to be null (i.e. have a
96+
/// null action function) in which case they are ignored.
97+
[[nodiscard]] Expected<std::vector<AllocAction>>
98+
runFinalizeActions(std::vector<AllocActionPair> AAPs);
99+
100+
/// Run the given deallocation actions in revwerse order.
101+
void runDeallocActions(std::vector<AllocAction> DAAs);
102+
103+
} // namespace orc_rt
104+
105+
#endif // ORC_RT_ALLOCACTION_H

‎orc-rt/include/orc-rt/SPSAllocAction.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
//===---- SPSAllocAction.h - SPS-serialized AllocAction utils ---*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Utilities for implementing allocation actions that take an SPS-serialized
10+
// argument buffer.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
14+
#ifndef ORC_RT_SPSALLOCACTION_H
15+
#define ORC_RT_SPSALLOCACTION_H
16+
17+
#include "orc-rt/AllocAction.h"
18+
#include "orc-rt/SimplePackedSerialization.h"
19+
20+
namespace orc_rt {
21+
22+
template <typename... SPSArgTs> struct AllocActionSPSDeserializer {
23+
template <typename... ArgTs>
24+
bool deserialize(const char *ArgData, size_t ArgSize, ArgTs &...Args) {
25+
SPSInputBuffer IB(ArgData, ArgSize);
26+
return SPSArgList<SPSArgTs...>::deserialize(IB, Args...);
27+
}
28+
};
29+
30+
/// Provides call and handle utilities to simplify writing and invocation of
31+
/// wrapper functions that use SimplePackedSerialization to serialize and
32+
/// deserialize their arguments and return values.
33+
template <typename... SPSArgTs> struct SPSAllocActionFunction {
34+
35+
template <typename Handler>
36+
static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
37+
Handler &&H) {
38+
return AllocActionFunction::handle(
39+
ArgData, ArgSize, AllocActionSPSDeserializer<SPSTuple<SPSArgTs...>>(),
40+
std::forward<Handler>(H));
41+
}
42+
};
43+
44+
} // namespace orc_rt
45+
46+
#endif // ORC_RT_SPSWRAPPERFUNCTION_H

‎orc-rt/lib/executor/AllocAction.cpp‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
//===- AllocAction.cpp ----------------------------------------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// AllocAction and related APIs.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#include "orc-rt/AllocAction.h"
14+
#include "orc-rt/ScopeExit.h"
15+
16+
namespace orc_rt {
17+
18+
Expected<std::vector<AllocAction>>
19+
runFinalizeActions(std::vector<AllocActionPair> AAPs) {
20+
std::vector<AllocAction> DeallocActions;
21+
auto RunDeallocActions = make_scope_exit([&]() {
22+
while (!DeallocActions.empty()) {
23+
// TODO: Log errors from cleanup dealloc actions.
24+
{
25+
[[maybe_unused]] auto B = DeallocActions.back()();
26+
}
27+
DeallocActions.pop_back();
28+
}
29+
});
30+
31+
for (auto &AAP : AAPs) {
32+
if (AAP.Finalize) {
33+
auto B = AAP.Finalize();
34+
if (const char *ErrMsg = B.getOutOfBandError())
35+
return make_error<StringError>(ErrMsg);
36+
}
37+
if (AAP.Dealloc)
38+
DeallocActions.push_back(std::move(AAP.Dealloc));
39+
}
40+
41+
RunDeallocActions.release();
42+
return DeallocActions;
43+
}
44+
45+
void runDeallocActions(std::vector<AllocAction> DAAs) {
46+
while (!DAAs.empty()) {
47+
// TODO: Log errors from cleanup dealloc actions.
48+
{
49+
[[maybe_unused]] auto B = DAAs.back()();
50+
}
51+
DAAs.pop_back();
52+
}
53+
}
54+
55+
} // namespace orc_rt

‎orc-rt/lib/executor/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
set(files
22
orc-rt-executor.cpp
3+
AllocAction.cpp
34
RTTI.cpp
45
)
56

‎orc-rt/unittests/AllocActionTest.cpp‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
//===- AllocActionTest.cpp ------------------------------------------------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// Tests for orc-rt's AllocAction.h APIs.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#include "orc-rt/AllocAction.h"
14+
#include "orc-rt/ExecutorAddress.h"
15+
#include "orc-rt/SPSAllocAction.h"
16+
17+
#include "SimplePackedSerializationTestUtils.h"
18+
#include "gtest/gtest.h"
19+
20+
using namespace orc_rt;
21+
22+
TEST(AllocActionTest, DefaultConstruct) {
23+
AllocAction AA;
24+
EXPECT_FALSE(AA);
25+
}
26+
27+
static orc_rt_WrapperFunctionBuffer noopAction(const char *ArgData,
28+
size_t ArgSize) {
29+
return WrapperFunctionBuffer().release();
30+
}
31+
32+
TEST(AllocActionTest, ConstructWithAction) {
33+
AllocAction AA(noopAction, WrapperFunctionBuffer());
34+
EXPECT_TRUE(AA);
35+
}
36+
37+
// Increments int via pointer.
38+
static orc_rt_WrapperFunctionBuffer
39+
increment_sps_allocaction(const char *ArgData, size_t ArgSize) {
40+
return SPSAllocActionFunction<SPSExecutorAddr>::handle(
41+
ArgData, ArgSize,
42+
[](ExecutorAddr IntPtr) {
43+
*IntPtr.toPtr<int *>() += 1;
44+
return WrapperFunctionBuffer();
45+
})
46+
.release();
47+
}
48+
49+
// Increments int via pointer.
50+
static orc_rt_WrapperFunctionBuffer
51+
decrement_sps_allocaction(const char *ArgData, size_t ArgSize) {
52+
return SPSAllocActionFunction<SPSExecutorAddr>::handle(
53+
ArgData, ArgSize,
54+
[](ExecutorAddr IntPtr) {
55+
*IntPtr.toPtr<int *>() -= 1;
56+
return WrapperFunctionBuffer();
57+
})
58+
.release();
59+
}
60+
61+
template <typename T>
62+
static WrapperFunctionBuffer makeExecutorAddrBuffer(T *P) {
63+
return *spsSerialize<SPSArgList<SPSExecutorAddr>>(ExecutorAddr::fromPtr(P));
64+
}
65+
66+
TEST(AllocActionTest, RunBasicAction) {
67+
int Val = 0;
68+
AllocAction IncVal(increment_sps_allocaction, makeExecutorAddrBuffer(&Val));
69+
EXPECT_TRUE(IncVal);
70+
auto B = IncVal();
71+
EXPECT_TRUE(B.empty());
72+
EXPECT_EQ(Val, 1);
73+
}
74+
75+
TEST(AllocActionTest, RunFinalizationActionsComplete) {
76+
int Val = 0;
77+
78+
std::vector<AllocActionPair> InitialActions;
79+
80+
auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
81+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
82+
{decrement_sps_allocaction, MakeArgBuffer()}});
83+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
84+
{decrement_sps_allocaction, MakeArgBuffer()}});
85+
86+
auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
87+
88+
EXPECT_EQ(Val, 2);
89+
90+
runDeallocActions(std::move(DeallocActions));
91+
92+
EXPECT_EQ(Val, 0);
93+
}
94+
95+
static orc_rt_WrapperFunctionBuffer fail_sps_allocaction(const char *ArgData,
96+
size_t ArgSize) {
97+
return WrapperFunctionBuffer::createOutOfBandError("failed").release();
98+
}
99+
100+
TEST(AllocActionTest, RunFinalizeActionsFail) {
101+
int Val = 0;
102+
103+
std::vector<AllocActionPair> InitialActions;
104+
105+
auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
106+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
107+
{decrement_sps_allocaction, MakeArgBuffer()}});
108+
InitialActions.push_back({{fail_sps_allocaction, MakeArgBuffer()},
109+
{decrement_sps_allocaction, MakeArgBuffer()}});
110+
111+
auto DeallocActions = runFinalizeActions(std::move(InitialActions));
112+
113+
if (DeallocActions) {
114+
ADD_FAILURE() << "Failed to report error from runFinalizeActions";
115+
return;
116+
}
117+
118+
EXPECT_EQ(toString(DeallocActions.takeError()), std::string("failed"));
119+
120+
// Check that we ran the decrement corresponding to the first increment.
121+
EXPECT_EQ(Val, 0);
122+
}
123+
124+
TEST(AllocActionTest, RunFinalizeActionsNullFinalize) {
125+
int Val = 0;
126+
127+
std::vector<AllocActionPair> InitialActions;
128+
129+
auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
130+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
131+
{decrement_sps_allocaction, MakeArgBuffer()}});
132+
InitialActions.push_back({{nullptr, WrapperFunctionBuffer()},
133+
{decrement_sps_allocaction, MakeArgBuffer()}});
134+
135+
auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
136+
137+
// Both dealloc actions should be included in the returned list, despite one
138+
// of them having a null finalize action.
139+
EXPECT_EQ(DeallocActions.size(), 2U);
140+
141+
runDeallocActions(std::move(DeallocActions));
142+
143+
EXPECT_EQ(Val, -1);
144+
}
145+
146+
TEST(AllocActionTest, RunFinalizeActionsNullDealloc) {
147+
int Val = 0;
148+
149+
std::vector<AllocActionPair> InitialActions;
150+
151+
auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
152+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
153+
{decrement_sps_allocaction, MakeArgBuffer()}});
154+
InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
155+
{nullptr, WrapperFunctionBuffer()}});
156+
157+
auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
158+
159+
// Null dealloc actions should be filtered out of the returned list.
160+
EXPECT_EQ(DeallocActions.size(), 1U);
161+
162+
runDeallocActions(std::move(DeallocActions));
163+
164+
EXPECT_EQ(Val, 1);
165+
}

0 commit comments

Comments
 (0)