[Mlir-commits] [mlir] [mlir][ExecutionEngine] Reserve one allocation per object on COFF (PR #233658)

[Phil Li]

https://github.com/Phil-amd updated https://github.com/llvm/llvm-project/pull/233658

>From 444d29b883a3dc4c4513c6de53e32a295218d085 Mon Sep 17 00:00:00 2001
From: Phil Li <haicli at amd.com>
Date: Tue, 20 Oct 2026 10:47:38 +0800
Subject: [PATCH] [mlir][ExecutionEngine] Reserve one allocation per object on
 COFF

RuntimeDyld resolves the IMAGE_REL_AMD64_ADDR32NB relocations in .pdata
against an object's lowest section, and aborts when another section lies
more than 4GB above it. SectionMemoryManager maps code and data
separately unless it reserves one block per object, so a process that
JIT-compiles many modules on x86-64 Windows hits this abort at random
(#55386, #65641).

Reserve the block on COFF, as #172833 does on AArch64. The new unit test
checks that code is never mapped on its own. It runs on AArch64 and COFF
hosts, so it runs on Windows too, where it fails without this change.

Assisted-by: Claude (Anthropic)
---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 11 +--
 mlir/unittests/ExecutionEngine/Invoke.cpp    | 76 +++++++++++++++++++-
 2 files changed, 80 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index d9dce6ef40f64..e721be4531803 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -319,9 +319,13 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        llvm::jitlink::JITLinkMemoryManager
                                            &IgnoredMemMgr) {
-    // Needed to respect AArch64 ABI requirements on the distance between
-    // TEXT and GOT sections.
-    bool reserveAlloc = llvmModule->getTargetTriple().isAArch64();
+    const llvm::Triple &targetTriple = llvmModule->getTargetTriple();
+    // Allocate each object's sections from one reserved block. AArch64 limits
+    // the distance between the TEXT and GOT sections. COFF's image-relative
+    // relocations, such as those in .pdata, need every section within 4GB
+    // above the lowest one, and RuntimeDyld aborts when they are not.
+    bool reserveAlloc =
+        targetTriple.isAArch64() || targetTriple.isOSBinFormatCOFF();
     auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
         session, [sectionMemoryMapper = options.sectionMemoryMapper,
                   reserveAlloc](const MemoryBuffer &) {
@@ -338,7 +342,6 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
     // COFF format binaries (Windows) need special handling to deal with
     // exported symbol visibility.
     // cf llvm/lib/ExecutionEngine/Orc/LLJIT.cpp LLJIT::createObjectLinkingLayer
-    const llvm::Triple &targetTriple = llvmModule->getTargetTriple();
     if (targetTriple.isOSBinFormatCOFF()) {
       objectLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
       objectLayer->setAutoClaimResponsibilityForObjectSymbols(true);
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 4061d16a83807..c457eebb8eda4 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -25,8 +25,13 @@
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/Support/Memory.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/TargetParser/Triple.h"
 
 #include "gmock/gmock.h"
 
@@ -39,9 +44,6 @@
 
 using namespace mlir;
 
-// The JIT isn't supported on Windows at that time
-#if !defined(_WIN32) && !defined(_AIX)
-
 static struct LLVMInitializer {
   LLVMInitializer() {
     llvm::InitializeNativeTarget();
@@ -60,6 +62,74 @@ static LogicalResult lowerToLLVMDialect(ModuleOp module) {
   return pm.run(module);
 }
 
+namespace {
+/// Forwards to sys::Memory and records what each mapping was requested for.
+class RecordingMemoryMapper : public llvm::SectionMemoryManager::MemoryMapper {
+public:
+  llvm::sys::MemoryBlock
+  allocateMappedMemory(llvm::SectionMemoryManager::AllocationPurpose purpose,
+                       size_t numBytes, const llvm::sys::MemoryBlock *nearBlock,
+                       unsigned flags, std::error_code &ec) override {
+    purposes.push_back(purpose);
+    return llvm::sys::Memory::allocateMappedMemory(numBytes, nearBlock, flags,
+                                                   ec);
+  }
+
+  std::error_code protectMappedMemory(const llvm::sys::MemoryBlock &block,
+                                      unsigned flags) override {
+    return llvm::sys::Memory::protectMappedMemory(block, flags);
+  }
+
+  std::error_code releaseMappedMemory(llvm::sys::MemoryBlock &block) override {
+    return llvm::sys::Memory::releaseMappedMemory(block);
+  }
+
+  std::vector<llvm::SectionMemoryManager::AllocationPurpose> purposes;
+};
+} // namespace
+
+// On AArch64 and COFF hosts the JIT carves each object's code and data from
+// one reserved mapping, because their relocations limit the distance between
+// sections. Code mapped on its own could land anywhere relative to the data.
+// Unlike the tests below, this one also runs on Windows.
+TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(CodeAndDataShareOneMapping)) {
+  llvm::Triple host(llvm::sys::getProcessTriple());
+  if (!host.isAArch64() && !host.isOSBinFormatCOFF())
+    GTEST_SKIP() << "only AArch64 and COFF hosts reserve one mapping";
+  std::string moduleStr = R"mlir(
+  func.func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
+    %res = arith.addi %arg0, %arg0 : i32
+    return %res : i32
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerBuiltinDialectTranslation(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  RecordingMemoryMapper mapper;
+  ExecutionEngineOptions options;
+  options.sectionMemoryMapper = &mapper;
+  auto jitOrError = ExecutionEngine::create(*module, options);
+  ASSERT_TRUE(!!jitOrError);
+  std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
+  int result = 0;
+  llvm::Error error =
+      jit->invoke("foo", 42, ExecutionEngine::Result<int>(result));
+  ASSERT_TRUE(!error);
+  ASSERT_EQ(result, 42 + 42);
+  ASSERT_FALSE(mapper.purposes.empty());
+  EXPECT_FALSE(llvm::is_contained(
+      mapper.purposes, llvm::SectionMemoryManager::AllocationPurpose::Code));
+}
+
+// The JIT isn't supported on Windows at that time
+#if !defined(_WIN32) && !defined(_AIX)
+
 TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(AddInteger)) {
 #ifdef __s390__
   std::string moduleStr = R"mlir(