[ORC] Add, and call through, reoptimize function in ReOptimizeLayer. (#173204)

lhames · web-flow · commit 01effcd82dfb · 2025-12-23T17:28:28.000+11:00
ReOptimizeLayer was building LLVM IR to define a precomputed,
SPS-serialized argument buffer, then inserting calls directly to
__orc_rt_jit_dispatch, passing the address of the precomputed buffer and
an __orc_rt_reoptimize_tag defined by the ORC runtime. This design is
non-canonical, requiring the ORC runtime to be loaded (or an extra
definition for __orc_rt_reoptimize_tag to be inserted) while not using
the runtime to perform the serialization.

This commit updates ReOptimizeLayer to instead insert calls to an
__orc_rt_reoptimize function implemented in the ORC runtime. This
function will perform serialization and call __orc_rt_jit_dispatch,
similar to other functions in the ORC runtime.

To maintain support for in-process JITs that don't use the ORC runtime,
this commit adds a ReOptimizeLayer::addOrcRTLiteSupport method which
injects IR to define __orc_rt_reoptimize (calling through to an
orc_rt_lite_reoptimize_helper function defined in LLVM) and
__orc_rt_reoptimize_tag. The ReOptimizeLayerTest is updated to use
addOrcRTLiteSupport.
diff --git a/compiler-rt/lib/orc/CMakeLists.txt b/compiler-rt/lib/orc/CMakeLists.txt
@@ -4,6 +4,7 @@
 set(ORC_COMMON_SOURCES
   debug.cpp
   dlfcn_wrapper.cpp
+  reoptimize.cpp
   rtti.cpp
   log_error_to_stderr.cpp
   run_program_wrapper.cpp
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -30,7 +30,6 @@ using namespace orc_rt;
 using namespace orc_rt::elfnix;
 
 // Declare function tags for functions in the JIT process.
-ORC_RT_JIT_DISPATCH_TAG(__orc_rt_reoptimize_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_push_initializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_symbol_lookup_tag)
 
diff --git a/compiler-rt/lib/orc/reoptimize.cpp b/compiler-rt/lib/orc/reoptimize.cpp
@@ -0,0 +1,27 @@
+//===- reoptimize.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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains code required to load the rest of the ELF-on-*IX runtime.
+//
+//===----------------------------------------------------------------------===//
+
+#include "jit_dispatch.h"
+#include "wrapper_function_utils.h"
+
+using namespace orc_rt;
+
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_reoptimize_tag)
+
+ORC_RT_INTERFACE void __orc_rt_reoptimize(uint64_t MUID, uint32_t CurVersion) {
+  if (auto Err = WrapperFunction<void(uint64_t, uint32_t)>::call(
+          JITDispatch(&__orc_rt_reoptimize_tag), MUID, CurVersion)) {
+    __orc_rt_log_error(toString(std::move(Err)).c_str());
+    // FIXME: Should we abort here? Depending on the error we can't guarantee
+    //        that the JIT'd code is in a consistent state.
+  }
+}
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
@@ -60,6 +60,15 @@ class LLVM_ABI ReOptimizeLayer : public IRLayer, public ResourceManager {
     this->ProfilerFunc = std::move(ProfilerFunc);
   }
 
+  /// Add ORC Runtime-lite support for reoptimization to PlatformJD.
+  ///
+  /// This allows reoptimization to be used without the ORC runtime.
+  ///
+  /// WARNING: For use with in-process JITs only.
+  /// WARNING: Do not use if the ORC runtime is loaded, as this will introduce
+  ///          duplicate definitions.
+  Error addOrcRTLiteSupport(JITDylib &PlatformJD, const DataLayout &DL);
+
   /// Registers reoptimize runtime dispatch handlers to given PlatformJD. The
   /// reoptimization request will not be handled if dispatch handler is not
   /// registered by using this function.
@@ -87,7 +96,8 @@ class LLVM_ABI ReOptimizeLayer : public IRLayer, public ResourceManager {
 
   // Create IR reoptimize request fucntion call.
   static void createReoptimizeCall(Module &M, Instruction &IP,
-                                   GlobalVariable *ArgBuffer);
+                                   ReOptMaterializationUnitID MUID,
+                                   unsigned CurVersion);
 
   Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
   void handleTransferResources(JITDylib &JD, ResourceKey DstK,
@@ -148,10 +158,6 @@ class LLVM_ABI ReOptimizeLayer : public IRLayer, public ResourceManager {
   void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
                      uint32_t CurVersion);
 
-  static Expected<Constant *>
-  createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
-                            uint32_t CurVersion);
-
   ReOptMaterializationUnitState &
   createMaterializationUnitState(const ThreadSafeModule &TSM);
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
@@ -26,6 +26,103 @@ void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
   Reoptimizing = false;
 }
 
+static void orc_rt_lite_reoptimize_helper(
+    shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
+                                                  void *Data, size_t Size),
+    void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
+  // Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
+  using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
+  auto ArgBytes =
+      shared::WrapperFunctionBuffer::allocate(SPSArgs::size(MUID, CurVersion));
+  shared::SPSOutputBuffer OB(ArgBytes.data(), ArgBytes.size());
+  if (!SPSArgs::serialize(OB, MUID, CurVersion)) {
+    errs()
+        << "Reoptimization error: could not serialize reoptimization arguments";
+    abort();
+  }
+  shared::WrapperFunctionBuffer Buf{
+      JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
+
+  if (const char *ErrMsg = Buf.getOutOfBandError()) {
+    errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";
+    abort();
+  }
+}
+
+Error ReOptimizeLayer::addOrcRTLiteSupport(JITDylib &PlatformJD,
+                                           const DataLayout &DL) {
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto Mod = std::make_unique<Module>("orc-rt-lite-reoptimize.ll", *Ctx);
+  Mod->setDataLayout(DL);
+
+  IRBuilder<> Builder(*Ctx);
+
+  // Create basic types portably
+  Type *VoidTy = Type::getVoidTy(*Ctx);
+  Type *Int8Ty = Type::getInt8Ty(*Ctx);
+  Type *Int32Ty = Type::getInt32Ty(*Ctx);
+  Type *Int64Ty = Type::getInt64Ty(*Ctx);
+  Type *VoidPtrTy = PointerType::getUnqual(*Ctx);
+
+  // Helper function type: void (void*, void*, void*, uint64_t, uint32_t)
+  FunctionType *HelperFnTy = FunctionType::get(
+      VoidTy, {VoidPtrTy, VoidPtrTy, VoidPtrTy, Int64Ty, Int32Ty}, false);
+
+  // Define ReoptimizeTag with initializer = 0
+  GlobalVariable *ReoptimizeTag = new GlobalVariable(
+      *Mod, Int8Ty, false, GlobalValue::ExternalLinkage,
+      ConstantInt::get(Int8Ty, 0), "__orc_rt_reoptimize_tag");
+
+  // Define orc_rt_lite_reoptimize function: void (uint64_t, uint32_t)
+  FunctionType *ReOptimizeFnTy =
+      FunctionType::get(VoidTy, {Int64Ty, Int32Ty}, false);
+
+  Function *ReOptimizeFn =
+      Function::Create(ReOptimizeFnTy, Function::ExternalLinkage,
+                       "__orc_rt_reoptimize", Mod.get());
+
+  // Set parameter names
+  auto ArgIt = ReOptimizeFn->arg_begin();
+  Value *MUID = &*ArgIt++;
+  MUID->setName("MUID");
+  Value *CurVersion = &*ArgIt;
+  CurVersion->setName("CurVersion");
+
+  // Build function body
+  BasicBlock *Entry = BasicBlock::Create(*Ctx, "entry", ReOptimizeFn);
+  Builder.SetInsertPoint(Entry);
+
+  // Create absolute address constants
+  auto &JDI = PlatformJD.getExecutionSession()
+                  .getExecutorProcessControl()
+                  .getJITDispatchInfo();
+
+  Type *IntPtrTy = DL.getIntPtrType(*Ctx);
+  Constant *JITDispatchPtr = ConstantExpr::getIntToPtr(
+      ConstantInt::get(IntPtrTy, JDI.JITDispatchFunction.getValue()),
+      VoidPtrTy);
+  Constant *JITDispatchCtxPtr = ConstantExpr::getIntToPtr(
+      ConstantInt::get(IntPtrTy, JDI.JITDispatchContext.getValue()), VoidPtrTy);
+  Constant *HelperFnAddr = ConstantExpr::getIntToPtr(
+      ConstantInt::get(IntPtrTy, reinterpret_cast<uintptr_t>(
+                                     &orc_rt_lite_reoptimize_helper)),
+      PointerType::getUnqual(*Ctx));
+
+  // Cast ReoptimizeTag to void*
+  Value *ReoptimizeTagPtr = Builder.CreatePointerCast(ReoptimizeTag, VoidPtrTy);
+
+  // Call the helper function
+  Builder.CreateCall(
+      HelperFnTy, HelperFnAddr,
+      {JITDispatchPtr, JITDispatchCtxPtr, ReoptimizeTagPtr, MUID, CurVersion});
+
+  // Return void
+  Builder.CreateRetVoid();
+
+  return BaseLayer.add(PlatformJD,
+                       ThreadSafeModule(std::move(Mod), std::move(Ctx)));
+}
+
 Error ReOptimizeLayer::registerRuntimeFunctions(JITDylib &PlatformJD) {
   ExecutionSession::JITDispatchHandlerAssociationMap WFs;
   using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
@@ -88,12 +185,6 @@ Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
     GlobalVariable *Counter = new GlobalVariable(
         M, I64Ty, false, GlobalValue::InternalLinkage,
         Constant::getNullValue(I64Ty), "__orc_reopt_counter");
-    auto ArgBufferConst = createReoptimizeArgBuffer(M, MUID, CurVersion);
-    if (auto Err = ArgBufferConst.takeError())
-      return Err;
-    GlobalVariable *ArgBuffer =
-        new GlobalVariable(M, (*ArgBufferConst)->getType(), true,
-                           GlobalValue::InternalLinkage, (*ArgBufferConst));
     for (auto &F : M) {
       if (F.isDeclaration())
         continue;
@@ -107,7 +198,7 @@ Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
       Value *Added = IRB.CreateAdd(Cnt, ConstantInt::get(I64Ty, 1));
       (void)IRB.CreateStore(Added, Counter);
       Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cmp, IP, false);
-      createReoptimizeCall(M, *SplitTerminator, ArgBuffer);
+      createReoptimizeCall(M, *SplitTerminator, MUID, CurVersion);
     }
     return Error::success();
   });
@@ -195,49 +286,23 @@ void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
   SendResult(Error::success());
 }
 
-Expected<Constant *> ReOptimizeLayer::createReoptimizeArgBuffer(
-    Module &M, ReOptMaterializationUnitID MUID, uint32_t CurVersion) {
-  size_t ArgBufferSize = SPSReoptimizeArgList::size(MUID, CurVersion);
-  std::vector<char> ArgBuffer(ArgBufferSize);
-  shared::SPSOutputBuffer OB(ArgBuffer.data(), ArgBuffer.size());
-  if (!SPSReoptimizeArgList::serialize(OB, MUID, CurVersion))
-    return make_error<StringError>("Could not serealize args list",
-                                   inconvertibleErrorCode());
-  return ConstantDataArray::get(M.getContext(), ArrayRef(ArgBuffer));
-}
-
 void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
-                                           GlobalVariable *ArgBuffer) {
-  GlobalVariable *DispatchCtx =
-      M.getGlobalVariable("__orc_rt_jit_dispatch_ctx");
-  if (!DispatchCtx)
-    DispatchCtx = new GlobalVariable(M, PointerType::get(M.getContext(), 0),
-                                     false, GlobalValue::ExternalLinkage,
-                                     nullptr, "__orc_rt_jit_dispatch_ctx");
-  GlobalVariable *ReoptimizeTag =
-      M.getGlobalVariable("__orc_rt_reoptimize_tag");
-  if (!ReoptimizeTag)
-    ReoptimizeTag = new GlobalVariable(M, PointerType::get(M.getContext(), 0),
-                                       false, GlobalValue::ExternalLinkage,
-                                       nullptr, "__orc_rt_reoptimize_tag");
-  Function *DispatchFunc = M.getFunction("__orc_rt_jit_dispatch");
-  if (!DispatchFunc) {
-    std::vector<Type *> Args = {PointerType::get(M.getContext(), 0),
-                                PointerType::get(M.getContext(), 0),
-                                PointerType::get(M.getContext(), 0),
-                                IntegerType::get(M.getContext(), 64)};
+                                           ReOptMaterializationUnitID MUID,
+                                           uint32_t CurVersion) {
+  Type *MUIDTy = IntegerType::get(M.getContext(), 64);
+  Type *VersionTy = IntegerType::get(M.getContext(), 32);
+  Function *ReoptimizeFunc = M.getFunction("__orc_rt_reoptimize");
+  if (!ReoptimizeFunc) {
+    std::vector<Type *> ArgTys = {MUIDTy, VersionTy};
     FunctionType *FuncTy =
-        FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
-    DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
-                                    "__orc_rt_jit_dispatch", &M);
+        FunctionType::get(Type::getVoidTy(M.getContext()), ArgTys, false);
+    ReoptimizeFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
+                                      "__orc_rt_reoptimize", &M);
   }
-  size_t ArgBufferSizeConst =
-      SPSReoptimizeArgList::size(ReOptMaterializationUnitID{}, uint32_t{});
-  Constant *ArgBufferSize = ConstantInt::get(
-      IntegerType::get(M.getContext(), 64), ArgBufferSizeConst, false);
+  Constant *MUIDArg = ConstantInt::get(MUIDTy, MUID, false);
+  Constant *CurVersionArg = ConstantInt::get(VersionTy, CurVersion, false);
   IRBuilder<> IRB(&IP);
-  (void)IRB.CreateCall(DispatchFunc,
-                       {DispatchCtx, ReoptimizeTag, ArgBuffer, ArgBufferSize});
+  (void)IRB.CreateCall(ReoptimizeFunc, {MUIDArg, CurVersionArg});
 }
 
 ReOptimizeLayer::ReOptMaterializationUnitState &
diff --git a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
@@ -140,22 +140,13 @@ static Function *createRetFunction(Module *M, StringRef Name,
 TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
   MangleAndInterner Mangle(*ES, *DL);
 
-  auto &EPC = ES->getExecutorProcessControl();
-  EXPECT_THAT_ERROR(JD->define(absoluteSymbols(
-                        {{Mangle("__orc_rt_jit_dispatch"),
-                          {EPC.getJITDispatchInfo().JITDispatchFunction,
-                           JITSymbolFlags::Exported}},
-                         {Mangle("__orc_rt_jit_dispatch_ctx"),
-                          {EPC.getJITDispatchInfo().JITDispatchContext,
-                           JITSymbolFlags::Exported}},
-                         {Mangle("__orc_rt_reoptimize_tag"),
-                          {ExecutorAddr(), JITSymbolFlags::Exported}}})),
-                    Succeeded());
-
   auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer);
   EXPECT_THAT_ERROR(RM.takeError(), Succeeded());
 
   ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *DL, *CompileLayer, **RM);
+  if (auto Err = ROLayer->addOrcRTLiteSupport(*JD, *DL))
+    FAIL() << toString(std::move(Err));
+
   ROLayer->setReoptimizeFunc(
       [&](ReOptimizeLayer &Parent,
           ReOptimizeLayer::ReOptMaterializationUnitID MUID, unsigned CurVerison,
