[ORC] Add CallViaEPC and CallSPSViaEPC utilities. (#170464)

These utilities simplify making typed async calls via
ExecutorProcessControl::callWrapperAsync.

CallViaEPC.h provides utilities for making typed async calls using a
given Serializer to serialize arguments and deserialize results.

callViaEPC takes a result handler function object (accepting
Expected<T>), an EPC reference, a Serializer, a target function address,
and arguments. The return type T is inferred from the handler's argument
type using CallableTraitsHelper.

EPCCaller wraps an ExecutorProcessControl& and Serializer, simplifying
repeated calls with the same serialization.

EPCCall wraps an EPCCaller and target function address, simplifying
repeated calls to a specific wrapper function.

CallSPSViaEPC.h provides utilities for using CallViaEPC with SPS
serialization.

SPSCallSerializer is a CallViaEPC-compatible serializer using SPS.

SPSEPCCaller takes an SPS function signature as its template argument
and applies SPS serialization to arguments and return values.

SPSEPCCall wraps an SPSCaller and a target function address, simplifying
repeated calls to a specific SPS wrapper function.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/CallSPSViaEPC.h

@@ -0,0 +1,79 @@
+//===---- CallSPSViaEPC.h - EPCCalls using SPS serialization ----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// EPCCalls using SimplePackedSerialization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_CALLSPSVIAEPC_H
+#define LLVM_EXECUTIONENGINE_ORC_CALLSPSVIAEPC_H
+
+#include "llvm/ExecutionEngine/Orc/CallViaEPC.h"
+#include "llvm/ExecutionEngine/Orc/CallableTraitsHelper.h"
+#include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
+
+namespace llvm::orc {
+
+namespace detail {
+template <typename RetT, typename... ArgTs> struct SPSCallSerializationImpl {
+  using RetSerialization = shared::SPSArgList<RetT>;
+  using ArgSerialization = shared::SPSArgList<ArgTs...>;
+};
+} // namespace detail
+
+template <typename SPSSig>
+struct SPSCallSerialization
+    : public CallableTraitsHelper<detail::SPSCallSerializationImpl, SPSSig> {};
+
+template <typename SPSSig> class SPSCallSerializer {
+public:
+  template <typename... ArgTs>
+  Expected<shared::WrapperFunctionResult> serialize(ArgTs &&...Args) {
+    using ArgSerialization =
+        typename SPSCallSerialization<SPSSig>::ArgSerialization;
+    auto Buffer = shared::WrapperFunctionResult::allocate(
+        ArgSerialization::size(Args...));
+    shared::SPSOutputBuffer OB(Buffer.data(), Buffer.size());
+    if (!ArgSerialization::serialize(OB, Args...))
+      return make_error<StringError>("Could not serialize arguments",
+                                     inconvertibleErrorCode());
+    return std::move(Buffer);
+  }
+
+  template <typename RetT>
+  Expected<RetT> deserialize(shared::WrapperFunctionResult ResultBytes) {
+    using RetDeserialization =
+        typename SPSCallSerialization<SPSSig>::RetSerialization;
+    shared::SPSInputBuffer IB(ResultBytes.data(), ResultBytes.size());
+    RetT ReturnValue;
+    if (!RetDeserialization::deserialize(IB, ReturnValue))
+      return make_error<StringError>("Could not deserialize return value",
+                                     inconvertibleErrorCode());
+    return ReturnValue;
+  }
+};
+
+template <typename SPSSig>
+class SPSEPCCaller : public EPCCaller<SPSCallSerializer<SPSSig>> {
+public:
+  SPSEPCCaller(ExecutorProcessControl &EPC)
+      : EPCCaller<SPSCallSerializer<SPSSig>>(EPC, SPSCallSerializer<SPSSig>()) {
+  }
+};
+
+template <typename SPSSig>
+class SPSEPCCall : public EPCCall<SPSCallSerializer<SPSSig>> {
+public:
+  SPSEPCCall(ExecutorProcessControl &EPC, ExecutorSymbolDef Fn)
+      : EPCCall<SPSCallSerializer<SPSSig>>(EPC, SPSCallSerializer<SPSSig>(),
+                                           std::move(Fn)) {}
+};
+
+} // namespace llvm::orc
+
+#endif // LLVM_EXECUTIONENGINE_ORC_CALLSPSVIAEPC_H

llvm/include/llvm/ExecutionEngine/Orc/CallViaEPC.h

@@ -0,0 +1,102 @@
+//===------ CallViaEPC.h - Call wrapper functions via EPC -------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Call executor functions with common signatures via
+// ExecutorProcessControl::callWrapperAsync.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_CALLVIAEPC_H
+#define LLVM_EXECUTIONENGINE_ORC_CALLVIAEPC_H
+
+#include "llvm/ExecutionEngine/Orc/CallableTraitsHelper.h"
+#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+
+namespace llvm::orc {
+
+namespace detail {
+
+// Helper to extract the Expected<T> argument type from a handler callable.
+template <typename HandlerT> struct HandlerTraits {
+  using ArgInfo = CallableArgInfo<HandlerT>;
+  using ArgsTuple = typename ArgInfo::ArgsTupleType;
+  static_assert(std::tuple_size_v<ArgsTuple> == 1,
+                "Handler must take exactly one argument");
+  using ExpectedArgType = std::tuple_element_t<0, ArgsTuple>;
+  using RetT = typename std::remove_cv_t<
+      std::remove_reference_t<ExpectedArgType>>::value_type;
+};
+
+} // namespace detail
+
+/// Call a wrapper function via EPC.
+template <typename HandlerT, typename Serializer, typename... ArgTs>
+void callViaEPC(HandlerT &&H, ExecutorProcessControl &EPC, Serializer S,
+                ExecutorSymbolDef Fn, ArgTs &&...Args) {
+  using RetT = typename detail::HandlerTraits<HandlerT>::RetT;
+
+  if (auto ArgBytes = S.serialize(std::forward<ArgTs>(Args)...))
+    EPC.callWrapperAsync(
+        Fn.getAddress(),
+        [S = std::move(S), H = std::forward<HandlerT>(H)](
+            shared::WrapperFunctionResult R) mutable {
+          if (const char *ErrMsg = R.getOutOfBandError())
+            H(make_error<StringError>(ErrMsg, inconvertibleErrorCode()));
+          else
+            H(S.template deserialize<RetT>(std::move(R)));
+        },
+        {ArgBytes->data(), ArgBytes->size()});
+  else
+    H(ArgBytes.takeError());
+}
+
+/// Encapsulates calls via EPC to any function that's compatible with the given
+/// serialization scheme.
+template <typename Serializer> class EPCCaller {
+public:
+  EPCCaller(ExecutorProcessControl &EPC, Serializer &&S)
+      : EPC(EPC), S(std::move(S)) {}
+
+  // TODO: Add an ExecutionSession constructor once ExecutionSession has been
+  //       moved to its own header.
+
+  // Async call version.
+  template <typename HandlerT, typename... ArgTs>
+  void operator()(HandlerT &&H, ExecutorSymbolDef Fn, ArgTs &&...Args) {
+    callViaEPC(std::forward<HandlerT>(H), EPC, S, Fn,
+               std::forward<ArgTs>(Args)...);
+  }
+
+private:
+  ExecutorProcessControl &EPC;
+  Serializer S;
+};
+
+/// Encapsulates calls via EPC to a specific function, using the given
+/// serialization scheme.
+template <typename Serializer> class EPCCall {
+public:
+  EPCCall(ExecutorProcessControl &EPC, Serializer &&S, ExecutorSymbolDef Fn)
+      : Caller(EPC, std::move(S)), Fn(std::move(Fn)) {}
+
+  // TODO: Add an ExecutionSession constructor once ExecutionSession has been
+  //       moved to its own header.
+
+  template <typename HandlerT, typename... ArgTs>
+  void operator()(HandlerT &&H, ArgTs &&...Args) {
+    Caller(std::forward<HandlerT>(H), Fn, std::forward<ArgTs>(Args)...);
+  }
+
+private:
+  EPCCaller<Serializer> Caller;
+  ExecutorSymbolDef Fn;
+};
+
+} // namespace llvm::orc
+
+#endif // LLVM_EXECUTIONENGINE_ORC_CALLVIAEPC_H

llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

@@ -19,6 +19,7 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest(OrcJITTests
   CallableTraitsHelperTest.cpp
+  CallSPSViaEPCTest.cpp
   CoreAPIsTest.cpp
   ExecutorAddressTest.cpp
   ExecutionSessionWrapperFunctionCallsTest.cpp

llvm/unittests/ExecutionEngine/Orc/CallSPSViaEPCTest.cpp

@@ -0,0 +1,90 @@
+//===----------- CallSPSViaEPC.cpp - Test CallSPSViaEPC.h APIs ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/CallSPSViaEPC.h"
+#include "llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h"
+
+#include "llvm/Testing/Support/Error.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::orc::shared;
+
+static CWrapperFunctionResult mainWrapper(const char *ArgData, size_t ArgSize) {
+  return WrapperFunction<int32_t(SPSSequence<SPSString>)>::handle(
+             ArgData, ArgSize,
+             [](std::vector<std::string> Args) -> int32_t {
+               return Args.size();
+             })
+      .release();
+}
+
+TEST(CallSPSViaEPCTest, CallMainViaCaller) {
+  auto EPC = cantFail(SelfExecutorProcessControl::Create());
+  SPSEPCCaller<int32_t(SPSSequence<SPSString>)> C(*EPC);
+  std::vector<std::string> Args;
+
+  std::optional<Expected<int32_t>> R1;
+  C([&](Expected<int32_t> R) { R1 = std::move(R); },
+    ExecutorSymbolDef::fromPtr(mainWrapper), Args);
+  ASSERT_THAT_EXPECTED(*R1, Succeeded());
+  EXPECT_EQ(**R1, 0);
+
+  Args.push_back("foo");
+  std::optional<Expected<int32_t>> R2;
+  C([&](Expected<int32_t> R) { R2 = std::move(R); },
+    ExecutorSymbolDef::fromPtr(mainWrapper), Args);
+  ASSERT_THAT_EXPECTED(*R2, Succeeded());
+  EXPECT_EQ(**R2, 1);
+
+  Args.push_back("foo");
+  std::optional<Expected<int32_t>> R3;
+  C([&](Expected<int32_t> R) { R3 = std::move(R); },
+    ExecutorSymbolDef::fromPtr(mainWrapper), Args);
+  ASSERT_THAT_EXPECTED(*R3, Succeeded());
+  EXPECT_EQ(**R3, 2);
+
+  Args.clear();
+  std::optional<Expected<int32_t>> R4;
+  C([&](Expected<int32_t> R) { R4 = std::move(R); },
+    ExecutorSymbolDef::fromPtr(mainWrapper), Args);
+  ASSERT_THAT_EXPECTED(*R4, Succeeded());
+  EXPECT_EQ(**R4, 0);
+}
+
+TEST(CallSPSViaEPCTest, CallMainViaGenericCall) {
+  auto EPC = cantFail(SelfExecutorProcessControl::Create());
+  SPSEPCCall<int32_t(SPSSequence<SPSString>)> C(
+      *EPC, ExecutorSymbolDef::fromPtr(mainWrapper));
+  std::vector<std::string> Args;
+
+  std::optional<Expected<int32_t>> R1;
+  C([&](Expected<int32_t> R) { R1 = std::move(R); }, Args);
+  ASSERT_THAT_EXPECTED(*R1, Succeeded());
+  EXPECT_EQ(**R1, 0);
+
+  Args.push_back("foo");
+  std::optional<Expected<int32_t>> R2;
+  C([&](Expected<int32_t> R) { R2 = std::move(R); }, Args);
+  ASSERT_THAT_EXPECTED(*R2, Succeeded());
+  EXPECT_EQ(**R2, 1);
+
+  Args.push_back("foo");
+  std::optional<Expected<int32_t>> R3;
+  C([&](Expected<int32_t> R) { R3 = std::move(R); }, Args);
+  ASSERT_THAT_EXPECTED(*R3, Succeeded());
+  EXPECT_EQ(**R3, 2);
+
+  Args.clear();
+  std::optional<Expected<int32_t>> R4;
+  C([&](Expected<int32_t> R) { R4 = std::move(R); }, Args);
+  ASSERT_THAT_EXPECTED(*R4, Succeeded());
+  EXPECT_EQ(**R4, 0);
+}