[Mlir-commits] [mlir] [mlir][ExecutionEngine] Use JITLink for RISC-V (and AArch64) and fix … (PR #196023)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu May 7 21:02:05 PDT 2026


https://github.com/XYenChi updated https://github.com/llvm/llvm-project/pull/196023

>From 9d211de402c47c2481e31248c9d0de5520856b65 Mon Sep 17 00:00:00 2001
From: XYenChi <oriachiuan at gmail.com>
Date: Fri, 17 Apr 2026 18:47:31 +0800
Subject: [PATCH 1/3] [mlir][ExecutionEngine] Use JITLink for RISC-V (and
 AArch64) and fix JitRunner defaults

The MLIR ExecutionEngine currently always builds an
RTDyldObjectLinkingLayer. RuntimeDyld has limited/!no support for
RISC-V relocations, which causes mlir-cpu-runner and any JIT-based
MLIR client to fail on RISC-V hosts.

Switch the object linking layer selection to prefer JITLink
(ObjectLinkingLayer) on targets where it is the better-supported
backend (currently AArch64 and RISC-V), and keep
RTDyldObjectLinkingLayer for the remaining targets. This mirrors the
selection logic already used by LLJIT::createObjectLinkingLayer.

Behavioral notes:
  * AArch64 previously enabled SectionMemoryManager's reserveAlloc
    workaround under RTDyld; on the JITLink path this is unnecessary
    because JITLink performs its own allocation honoring the AArch64
    TEXT/GOT distance requirement.
  * GDB/perf JITEventListeners are only registered on the RTDyld
    path, since they are tied to RuntimeDyld. JITLink-based listener
    support can be wired up separately.
  * The COFF symbol-visibility workaround is now guarded so it only
    runs on the RTDyld path (it is an RTDyld-only API).

In addition, JitRunner now sets RelocModel=PIC_ and CodeModel=Medium
when the target triple is RISC-V, which matches the medany ABI
required by the JITLink RISC-V backend and avoids out-of-range
relocations at link time.

Includes:
  * Add llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h to
    ExecutionEngine.cpp.
  * Add llvm/Support/CodeGen.h to JitRunner.cpp for the
    Reloc/CodeModel enums.
---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 66 +++++++++++++++-----
 mlir/lib/ExecutionEngine/JitRunner.cpp       |  6 ++
 2 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 49a88e0000511..d099d27a92897 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -10,19 +10,21 @@
 // JIT engine.
 //
 //===----------------------------------------------------------------------===//
-#include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR/Export.h"
 
 #include "llvm/ExecutionEngine/JITEventListener.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -316,27 +318,57 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
                                            &IgnoredMemMgr) {
     // 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,
-                  reserveAlloc](const MemoryBuffer &) {
-          return std::make_unique<SectionMemoryManager>(sectionMemoryMapper,
-                                                        reserveAlloc);
-        });
-
-    // Register JIT event listeners if they are enabled.
-    if (engine->gdbListener)
-      objectLayer->registerJITEventListener(*engine->gdbListener);
-    if (engine->perfListener)
-      objectLayer->registerJITEventListener(*engine->perfListener);
+
+    // Check if we should use ObjectLinkingLayer (JITLink)
+    // JITLink supports modern architectures like RISC-V, AArch64
+    // RuntimeDyld is older and provides better compatibility with legacy
+    // platforms
+
+    // Decide which layer to use
+    bool useJITLink = llvmModule->getTargetTriple().isAArch64() ||
+                      llvmModule->getTargetTriple().isRISCV();
+
+    std::unique_ptr<llvm::orc::ObjectLayer> objectLayer;
+
+    if (useJITLink) {
+      // JITLink path
+      objectLayer = std::make_unique<llvm::orc::ObjectLinkingLayer>(session);
+
+      LLVM_DEBUG(llvm::dbgs() << "Using ObjectLinkingLayer (JITLink)\n");
+
+    } else {
+      // RuntimeDyld path
+      auto rtDyldLayer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(
+          session,
+          [sectionMemoryMapper =
+               options.sectionMemoryMapper](const llvm::MemoryBuffer &)
+              -> std::unique_ptr<llvm::RuntimeDyld::MemoryManager> {
+            return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
+          });
+
+      // Only RTDyld supports listener
+      if (engine->gdbListener)
+        rtDyldLayer->registerJITEventListener(*engine->gdbListener);
+
+      if (engine->perfListener)
+        rtDyldLayer->registerJITEventListener(*engine->perfListener);
+
+      LLVM_DEBUG(llvm::dbgs() << "Using RTDyldObjectLinkingLayer\n");
+
+      // Upcast
+      objectLayer = std::move(rtDyldLayer);
+    }
 
     // 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);
+    if (!useJITLink && targetTriple.isOSBinFormatCOFF()) {
+      if (auto *rtDyldLayer = dyn_cast<llvm::orc::RTDyldObjectLinkingLayer>(
+              objectLayer.get())) {
+        rtDyldLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
+        rtDyldLayer->setAutoClaimResponsibilityForObjectSymbols(true);
+      }
     }
 
     // Resolve symbols from shared libraries.
diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index db0516533afcb..4b52669df9a6e 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -31,6 +31,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
+#include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileUtilities.h"
@@ -361,6 +362,11 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
     tmBuilderOrError->getTargetTriple().setArchName(options.mArch);
   }
 
+  if (tmBuilderOrError->getTargetTriple().isRISCV()){
+    tmBuilderOrError->setRelocationModel(llvm::Reloc::PIC_);
+    tmBuilderOrError->setCodeModel(llvm::CodeModel::Medium);
+  }
+
   // Build TargetMachine
   auto tmOrError = tmBuilderOrError->createTargetMachine();
 

>From 97c28debc1bea31b42d5e1ded90d59a305ebb9eb Mon Sep 17 00:00:00 2001
From: XYenChi <oriachiuan at gmail.com>
Date: Wed, 6 May 2026 17:50:54 +0800
Subject: [PATCH 2/3] Update mlir/lib/ExecutionEngine/ExecutionEngine.cpp

use LDBG preferably for debug log.

Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
Assisted-by: Claude Opus 4.7
---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index d099d27a92897..895d3ad0cc6f1 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -332,10 +332,8 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
 
     if (useJITLink) {
       // JITLink path
+      LDBG() << "Using ObjectLinkingLayer (JITLink)";
       objectLayer = std::make_unique<llvm::orc::ObjectLinkingLayer>(session);
-
-      LLVM_DEBUG(llvm::dbgs() << "Using ObjectLinkingLayer (JITLink)\n");
-
     } else {
       // RuntimeDyld path
       auto rtDyldLayer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(
@@ -353,9 +351,7 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
       if (engine->perfListener)
         rtDyldLayer->registerJITEventListener(*engine->perfListener);
 
-      LLVM_DEBUG(llvm::dbgs() << "Using RTDyldObjectLinkingLayer\n");
-
-      // Upcast
+      LDBG() << "mlir::ExecutionEngine initialized with RTDyldObjectLinkingLayer engine";
       objectLayer = std::move(rtDyldLayer);
     }
 

>From 62139ee012ef4b02a2d8476b9001a805283f8fbf Mon Sep 17 00:00:00 2001
From: Yixuan Chen <chenyixuan at iscas.ac.cn>
Date: Fri, 8 May 2026 01:54:37 +0800
Subject: [PATCH 3/3] Updated the object layer selection to prefer JITLink by
 default

Replaced the old AArch64/RISCV-only condition with target-aware logic that starts from JITLink enabled and selectively disables it for unsupported cases.
Kept RuntimeDyld-specific listener registration and COFF-specific flag handling only on the RuntimeDyld path.
Fixed the JITLink construction to pass the memory manager into ObjectLinkingLayer.
Replaced undeclared debug macro usage with LLVM_DEBUG logging in this block.

Add comment for RISC-V.

Assisted-by: Claude Opus 4.7
---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 44 ++++++++++++++------
 mlir/lib/ExecutionEngine/JitRunner.cpp       |  6 ++-
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 895d3ad0cc6f1..83a54c4d71958 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -29,6 +29,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugLog.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/TargetParser/Host.h"
@@ -315,25 +316,45 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        llvm::jitlink::JITLinkMemoryManager
-                                           &IgnoredMemMgr) {
+                                           &ignoredMemMgr) {
     // Needed to respect AArch64 ABI requirements on the distance between
     // TEXT and GOT sections.
 
-    // Check if we should use ObjectLinkingLayer (JITLink)
-    // JITLink supports modern architectures like RISC-V, AArch64
-    // RuntimeDyld is older and provides better compatibility with legacy
-    // platforms
-
-    // Decide which layer to use
-    bool useJITLink = llvmModule->getTargetTriple().isAArch64() ||
-                      llvmModule->getTargetTriple().isRISCV();
+    // Prefer ObjectLinkingLayer (JITLink) by default and only fall back to
+    // RuntimeDyld for target / format combinations where JITLink is not
+    // currently supported.
+    const llvm::Triple &targetTriple = llvmModule->getTargetTriple();
+    bool useJITLink = true;
+    switch (targetTriple.getArch()) {
+    case Triple::aarch64:
+      useJITLink = !targetTriple.isOSBinFormatCOFF();
+      break;
+    case Triple::arm:
+    case Triple::armeb:
+    case Triple::thumb:
+    case Triple::thumbeb:
+      useJITLink = targetTriple.isOSBinFormatELF();
+      break;
+    case Triple::ppc64:
+      useJITLink = targetTriple.isPPC64ELFv2ABI();
+      break;
+    case Triple::ppc64le:
+      useJITLink = targetTriple.isOSBinFormatELF();
+      break;
+    case Triple::x86_64:
+      useJITLink = !targetTriple.isOSBinFormatCOFF();
+      break;
+    default:
+      break;
+    }
 
     std::unique_ptr<llvm::orc::ObjectLayer> objectLayer;
 
     if (useJITLink) {
       // JITLink path
       LDBG() << "Using ObjectLinkingLayer (JITLink)";
-      objectLayer = std::make_unique<llvm::orc::ObjectLinkingLayer>(session);
+      objectLayer = std::make_unique<llvm::orc::ObjectLinkingLayer>(session,
+                                                           ignoredMemMgr);
     } else {
       // RuntimeDyld path
       auto rtDyldLayer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(
@@ -358,9 +379,8 @@ 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 (!useJITLink && targetTriple.isOSBinFormatCOFF()) {
-      if (auto *rtDyldLayer = dyn_cast<llvm::orc::RTDyldObjectLinkingLayer>(
+      if (auto *rtDyldLayer = cast<llvm::orc::RTDyldObjectLinkingLayer>(
               objectLayer.get())) {
         rtDyldLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
         rtDyldLayer->setAutoClaimResponsibilityForObjectSymbols(true);
diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index 4b52669df9a6e..d805ee39009ca 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -362,7 +362,11 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
     tmBuilderOrError->getTargetTriple().setArchName(options.mArch);
   }
 
-  if (tmBuilderOrError->getTargetTriple().isRISCV()){
+  // On RISC-V, JITLink requires PIC relocations and the medium code model so
+  // that generated code and data fit within the +/-2GiB addressing range used
+  // by the default RISC-V relocations (e.g. PCREL_HI20/LO12). Without this,
+  // linking JITted objects can fail with out-of-range relocation errors.
+  if (tmBuilderOrError->getTargetTriple().isRISCV()) {
     tmBuilderOrError->setRelocationModel(llvm::Reloc::PIC_);
     tmBuilderOrError->setCodeModel(llvm::CodeModel::Medium);
   }



More information about the Mlir-commits mailing list