[Mlir-commits] [mlir] [MLIR] Enable dylib init/deinit in execution engine on AArch64 platform (PR #172833)

[Giacomo Castiglioni]

https://github.com/castigli updated https://github.com/llvm/llvm-project/pull/172833

>From deada8c3300163eb4e588454ff542f82aff8b214 Mon Sep 17 00:00:00 2001
From: Giacomo Castiglioni <giacastiglioni at gmail.com>
Date: Wed, 3 Dec 2025 09:46:38 +0100
Subject: [PATCH 1/2] Patch MLIR execution engine for aarch64.

---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 287c52a262c11..f8f8ee390a2d6 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -218,10 +218,7 @@ ExecutionEngine::ExecutionEngine(bool enableObjectDump,
 
 ExecutionEngine::~ExecutionEngine() {
   // Execute the global destructors from the module being processed.
-  // TODO: Allow JIT deinitialize for AArch64. Currently there's a bug causing a
-  // crash for AArch64 see related issue #71963.
-  if (jit && !jit->getTargetTriple().isAArch64())
-    llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
+  llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (LibraryDestroyFn destroy : destroyFns)
     destroy();
@@ -314,10 +311,14 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   // Callback to create the object layer with symbol resolution to current
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session) {
+    // Needed to respect AArch64 ABI requirements on the distance between
+    // TEXT and GOT sections.
+    bool reserveAlloc = llvmModule->getTargetTriple().isAArch64();
     auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
-        session, [sectionMemoryMapper =
-                      options.sectionMemoryMapper](const MemoryBuffer &) {
-          return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
+        session, [sectionMemoryMapper = options.sectionMemoryMapper,
+                  reserveAlloc](const MemoryBuffer &) {
+          return std::make_unique<SectionMemoryManager>(sectionMemoryMapper,
+                                                        reserveAlloc);
         });
 
     // Register JIT event listeners if they are enabled.
@@ -449,9 +450,6 @@ Error ExecutionEngine::invokePacked(StringRef name,
 void ExecutionEngine::initialize() {
   if (isInitialized)
     return;
-  // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a
-  // crash for AArch64 see related issue #71963.
-  if (!jit->getTargetTriple().isAArch64())
-    cantFail(jit->initialize(jit->getMainJITDylib()));
+  cantFail(jit->initialize(jit->getMainJITDylib()));
   isInitialized = true;
 }

>From e61af3e51fe566c27e9646378e651a7ed5d6f676 Mon Sep 17 00:00:00 2001
From: Giacomo Castiglioni <giacastiglioni at gmail.com>
Date: Thu, 18 Dec 2025 13:59:47 +0100
Subject: [PATCH 2/2] restore jit check on destructor

---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index f8f8ee390a2d6..9b16e09124aa3 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -218,7 +218,8 @@ ExecutionEngine::ExecutionEngine(bool enableObjectDump,
 
 ExecutionEngine::~ExecutionEngine() {
   // Execute the global destructors from the module being processed.
-  llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
+  if (jit)
+    llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (LibraryDestroyFn destroy : destroyFns)
     destroy();