From b7cb15feab961ac912e99fab5981972857265284 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sat, 6 Dec 2025 16:22:53 -0500 Subject: [PATCH 1/3] [MLIR][ExecutionEngine] Enable PIC option This PR enables the MLIR execution engine to dump object file as PIC code which is needed when the object file is bundled into a dynamic shared library. --- mlir/include/mlir-c/ExecutionEngine.h | 2 +- mlir/lib/Bindings/Python/ExecutionEngineModule.cpp | 10 +++++----- mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp | 8 ++++++-- mlir/test/CAPI/execution_engine.c | 4 ++-- mlir/test/CAPI/global_constructors.c | 2 +- mlir/test/python/execution_engine.py | 1 + 6 files changed, 16 insertions(+), 11 deletions(-) diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h index 1a58d68533f24..348a1642ee6a0 100644 --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -44,7 +44,7 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void); /// TODO: figure out other options. MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate( MlirModule op, int optLevel, int numPaths, - const MlirStringRef *sharedLibPaths, bool enableObjectDump); + const MlirStringRef *sharedLibPaths, bool enableObjectDump, bool enablePIC); /// Initialize the ExecutionEngine. Global constructors specified by /// `llvm.mlir.global_ctors` will be run. One common scenario is that kernel diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp index 8bb493ed7240c..be0785b126eaa 100644 --- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp +++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp @@ -75,13 +75,13 @@ NB_MODULE(_mlirExecutionEngine, m) { "__init__", [](PyExecutionEngine &self, MlirModule module, int optLevel, const std::vector &sharedLibPaths, - bool enableObjectDump) { + bool enableObjectDump, bool enablePIC) { llvm::SmallVector libPaths; for (const std::string &path : sharedLibPaths) libPaths.push_back({path.c_str(), path.length()}); - MlirExecutionEngine executionEngine = - mlirExecutionEngineCreate(module, optLevel, libPaths.size(), - libPaths.data(), enableObjectDump); + MlirExecutionEngine executionEngine = mlirExecutionEngineCreate( + module, optLevel, libPaths.size(), libPaths.data(), + enableObjectDump, enablePIC); if (mlirExecutionEngineIsNull(executionEngine)) throw std::runtime_error( "Failure while creating the ExecutionEngine."); @@ -89,7 +89,7 @@ NB_MODULE(_mlirExecutionEngine, m) { }, nb::arg("module"), nb::arg("opt_level") = 2, nb::arg("shared_libs") = nb::list(), - nb::arg("enable_object_dump") = true, + nb::arg("enable_object_dump") = true, nb::arg("enable_pic") = false, "Create a new ExecutionEngine instance for the given Module. The " "module must contain only dialects that can be translated to LLVM. " "Perform transformations and code generation at the optimization " diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp index 6c97499b28fd9..c3cb5e14f6160 100644 --- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp @@ -22,7 +22,7 @@ using namespace mlir; extern "C" MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, const MlirStringRef *sharedLibPaths, - bool enableObjectDump) { + bool enableObjectDump, bool enablePIC) { static bool initOnce = [] { llvm::InitializeNativeTarget(); llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm @@ -43,6 +43,9 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, consumeError(tmBuilderOrError.takeError()); return MlirExecutionEngine{nullptr}; } + if (enablePIC) { + tmBuilderOrError->setRelocationModel(llvm::Reloc::PIC_); + } auto tmOrError = tmBuilderOrError->createTargetMachine(); if (!tmOrError) { llvm::errs() << "Failed to create a TargetMachine for the host because: \n"; @@ -63,7 +66,8 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, jitOptions.jitCodeGenOptLevel = static_cast(optLevel); jitOptions.sharedLibPaths = libPaths; jitOptions.enableObjectDump = enableObjectDump; - auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions); + auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions, + std::move(tmOrError.get())); if (!jitOrError) { llvm::errs() << "Failed to create an ExecutionEngine because: \n"; consumeError(jitOrError.takeError()); diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c index 4751288c3ee4b..4df232f3caab3 100644 --- a/mlir/test/CAPI/execution_engine.c +++ b/mlir/test/CAPI/execution_engine.c @@ -69,7 +69,7 @@ void testSimpleExecution(void) { mlirRegisterAllLLVMTranslations(ctx); MlirExecutionEngine jit = mlirExecutionEngineCreate( module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL, - /*enableObjectDump=*/false); + /*enableObjectDump=*/false, /*enablePIC=*/false); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Execution engine creation failed"); exit(2); @@ -125,7 +125,7 @@ void testOmpCreation(void) { // against the OpenMP library. MlirExecutionEngine jit = mlirExecutionEngineCreate( module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL, - /*enableObjectDump=*/false); + /*enableObjectDump=*/false, /*enablePIC=*/false); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Engine creation failed with OpenMP"); exit(2); diff --git a/mlir/test/CAPI/global_constructors.c b/mlir/test/CAPI/global_constructors.c index bd2fe1416f0df..9aacaf2c513f3 100644 --- a/mlir/test/CAPI/global_constructors.c +++ b/mlir/test/CAPI/global_constructors.c @@ -79,7 +79,7 @@ void testGlobalCtorJitCallback(void) { // Create execution engine with initialization disabled MlirExecutionEngine jit = mlirExecutionEngineCreate( module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL, - /*enableObjectDump=*/false); + /*enableObjectDump=*/false, /*enablePIC=*/false); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Execution engine creation failed"); diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py index 005813d1788f5..b11340f2c19ce 100644 --- a/mlir/test/python/execution_engine.py +++ b/mlir/test/python/execution_engine.py @@ -807,6 +807,7 @@ def testDumpToObjectFile(): # because RTDyldObjectLinkingLayer::emit will try to resolve symbols before dumping # (see the jitLinkForORC call at the bottom there). shared_libs=[MLIR_C_RUNNER_UTILS], + enable_pic=True, ) # CHECK: Object file exists: True From 6c4987357dd9e1ea704feb4f04df21b27a8df789 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sun, 7 Dec 2025 07:55:33 -0500 Subject: [PATCH 2/3] Update mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp Co-authored-by: Mehdi Amini --- mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp index c3cb5e14f6160..81d86ad29af13 100644 --- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp @@ -43,9 +43,8 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, consumeError(tmBuilderOrError.takeError()); return MlirExecutionEngine{nullptr}; } - if (enablePIC) { + if (enablePIC) tmBuilderOrError->setRelocationModel(llvm::Reloc::PIC_); - } auto tmOrError = tmBuilderOrError->createTargetMachine(); if (!tmOrError) { llvm::errs() << "Failed to create a TargetMachine for the host because: \n"; From af0b61d49706436290b0dded2162e761bfc5b8e8 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 7 Dec 2025 09:57:47 -0500 Subject: [PATCH 3/3] Add docs per suggestion --- mlir/include/mlir-c/ExecutionEngine.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h index 348a1642ee6a0..2a81798bdbf37 100644 --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -41,6 +41,9 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void); /// generation. The number and array of paths corresponding to shared libraries /// that will be loaded are specified via `numPaths` and `sharedLibPaths` /// respectively. +/// The `enablePIC` arguments controls the relocation model, when true the +/// generated code is emitted as "position independent", making it possible to +/// save it and reload it as a shared object in another process. /// TODO: figure out other options. MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate( MlirModule op, int optLevel, int numPaths,