[Mlir-commits] [mlir] [mlir][ExecutionEngine] Use JITLink for RISC-V (and AArch64) and fix … (PR #196023)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat May 30 08:57:17 PDT 2026
https://github.com/XYenChi updated https://github.com/llvm/llvm-project/pull/196023
>From fcb0c9d9f674846ce0f33e4076fd8e099ab18998 Mon Sep 17 00:00:00 2001
From: XYenChi <oriachiuan at gmail.com>
Date: Fri, 17 Apr 2026 18:47:31 +0800
Subject: [PATCH] [mlir][ExecutionEngine] Use JITLink by default and tune
JitRunner for RISC-V
The MLIR ExecutionEngine previously always built an
RTDyldObjectLinkingLayer. RuntimeDyld has limited/no support for
RISC-V relocations, which causes mlir-cpu-runner and other JIT-based
MLIR clients to fail on RISC-V hosts.
Switch the object linking layer selection to prefer JITLink
(ObjectLinkingLayer) by default and only fall back to
RTDyldObjectLinkingLayer for target/format combinations where JITLink
is not yet supported. The per-arch/format decision mirrors the logic
used by LLJIT::createObjectLinkingLayer (AArch64, Arm/Thumb on ELF,
PPC64 ELFv2, PPC64LE on ELF, and x86_64 on non-COFF use JITLink;
other targets stay on RTDyld).
Behavioral notes:
* The AArch64 SectionMemoryManager reserveAlloc workaround is no
longer needed on the JITLink path, because JITLink performs its
own allocation honoring the AArch64 TEXT/GOT distance requirement.
It is dropped from the RTDyld constructor as well, since AArch64
now goes through JITLink by default.
* GDB/perf JITEventListeners are only registered on the RTDyld path,
since the existing listener API is tied to RuntimeDyld. JITLink
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).
* Debug output in this block uses LDBG() instead of an undeclared
debug macro.
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 (e.g. PCREL_HI20/LO12) at link time.
Includes:
* Add llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h and
llvm/Support/DebugLog.h to ExecutionEngine.cpp.
* Add llvm/Support/CodeGen.h to JitRunner.cpp for the
Reloc/CodeModel enums.
Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
Assisted-by: Claude Opus 4.7
---
mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 154 +++++++++++++------
mlir/lib/ExecutionEngine/JitRunner.cpp | 10 ++
2 files changed, 113 insertions(+), 51 deletions(-)
diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 49a88e0000511..af518c8c0f570 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -10,23 +10,26 @@
// 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"
#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"
@@ -311,56 +314,105 @@ 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,
- llvm::jitlink::JITLinkMemoryManager
- &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);
-
- // 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);
- }
-
- // Resolve symbols from shared libraries.
- for (auto &libPath : jitDyLibPaths) {
- auto mb = llvm::MemoryBuffer::getFile(libPath);
- if (!mb) {
- errs() << "Failed to create MemoryBuffer for: " << libPath
- << "\nError: " << mb.getError().message() << "\n";
- continue;
- }
- auto &jd = session.createBareJITDylib(std::string(libPath));
- auto loaded = DynamicLibrarySearchGenerator::Load(
- libPath.str().c_str(), dataLayout.getGlobalPrefix());
- if (!loaded) {
- errs() << "Could not load " << libPath << ":\n " << loaded.takeError()
- << "\n";
- continue;
- }
- jd.addGenerator(std::move(*loaded));
- cantFail(objectLayer->add(jd, std::move(mb.get())));
- }
-
- return objectLayer;
- };
+ auto objectLinkingLayerCreator =
+ [&](ExecutionSession &session,
+ llvm::jitlink::JITLinkMemoryManager &ignoredMemMgr) {
+ // Needed to respect AArch64 ABI requirements on the distance between
+ // TEXT and GOT sections.
+
+ // 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, ignoredMemMgr);
+ } 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);
+
+ LDBG() << "mlir::ExecutionEngine initialized with "
+ "RTDyldObjectLinkingLayer engine";
+ 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
+ if (!useJITLink && targetTriple.isOSBinFormatCOFF()) {
+ if (auto *rtDyldLayer = cast<llvm::orc::RTDyldObjectLinkingLayer>(
+ objectLayer.get())) {
+ rtDyldLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
+ rtDyldLayer->setAutoClaimResponsibilityForObjectSymbols(true);
+ }
+ }
+
+ // Resolve symbols from shared libraries.
+ for (auto &libPath : jitDyLibPaths) {
+ auto mb = llvm::MemoryBuffer::getFile(libPath);
+ if (!mb) {
+ errs() << "Failed to create MemoryBuffer for: " << libPath
+ << "\nError: " << mb.getError().message() << "\n";
+ continue;
+ }
+ auto &jd = session.createBareJITDylib(std::string(libPath));
+ auto loaded = DynamicLibrarySearchGenerator::Load(
+ libPath.str().c_str(), dataLayout.getGlobalPrefix());
+ if (!loaded) {
+ errs() << "Could not load " << libPath << ":\n "
+ << loaded.takeError() << "\n";
+ continue;
+ }
+ jd.addGenerator(std::move(*loaded));
+ cantFail(objectLayer->add(jd, std::move(mb.get())));
+ }
+
+ return objectLayer;
+ };
// Callback to inspect the cache and recompile on demand. This follows Lang's
// LLJITWithObjectCache example.
diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index db0516533afcb..d805ee39009ca 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,15 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry ®istry,
tmBuilderOrError->getTargetTriple().setArchName(options.mArch);
}
+ // 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);
+ }
+
// Build TargetMachine
auto tmOrError = tmBuilderOrError->createTargetMachine();
More information about the Mlir-commits
mailing list