From 41bf73a983ca6176e795bc61d006fb56a64a74be Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 15:24:29 +0800 Subject: [PATCH 01/15] Allow defer initialization of jit --- mlir/include/mlir-c/ExecutionEngine.h | 10 +++-- .../mlir/ExecutionEngine/ExecutionEngine.h | 8 ++++ .../CAPI/ExecutionEngine/ExecutionEngine.cpp | 12 ++++-- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 24 +++++++---- mlir/test/CAPI/execution_engine.c | 4 +- mlir/unittests/ExecutionEngine/Invoke.cpp | 40 +++++++++++++++++++ 6 files changed, 81 insertions(+), 17 deletions(-) diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h index 99cddc5c2598d..eee56ee032c97 100644 --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -42,9 +42,13 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void); /// that will be loaded are specified via `numPaths` and `sharedLibPaths` /// respectively. /// TODO: figure out other options. -MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate( - MlirModule op, int optLevel, int numPaths, - const MlirStringRef *sharedLibPaths, bool enableObjectDump); +MLIR_CAPI_EXPORTED MlirExecutionEngine +mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, + const MlirStringRef *sharedLibPaths, + bool enableObjectDump, bool shouldInitialize); + +/// Execute the global constructors from the module. +MLIR_CAPI_EXPORTED void mlirExecutionEngineInitialize(MlirExecutionEngine jit); /// Destroy an ExecutionEngine instance. MLIR_CAPI_EXPORTED void mlirExecutionEngineDestroy(MlirExecutionEngine jit); diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index 96ccebcd5685e..247446b842e05 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -99,6 +99,10 @@ struct ExecutionEngineOptions { /// If `enablePerfNotificationListener` is set, the JIT compiler will notify /// the llvm's global Perf notification listener. bool enablePerfNotificationListener = true; + + /// Setting initialize=false to register symbols might be used in + /// initialization + bool shouldInitialize = true; }; /// JIT-backed execution engine for MLIR. Assumes the IR can be converted to @@ -227,6 +231,8 @@ class ExecutionEngine { llvm::function_ref symbolMap); + void initialize(); + private: /// Ordering of llvmContext and jit is important for destruction purposes: the /// jit must be destroyed before the context. @@ -250,6 +256,8 @@ class ExecutionEngine { /// Destroy functions in the libraries loaded by the ExecutionEngine that are /// called when this ExecutionEngine is destructed. SmallVector destroyFns; + + bool isInitialized = false; }; } // namespace mlir diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp index 306cebd236be9..9d445bc4210ca 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 shouldInitialize) { static bool initOnce = [] { llvm::InitializeNativeTarget(); llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm @@ -60,6 +60,7 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, jitOptions.jitCodeGenOptLevel = static_cast(optLevel); jitOptions.sharedLibPaths = libPaths; jitOptions.enableObjectDump = enableObjectDump; + jitOptions.shouldInitialize = shouldInitialize; auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions); if (!jitOrError) { consumeError(jitOrError.takeError()); @@ -68,6 +69,10 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, return wrap(jitOrError->release()); } +extern "C" void mlirExecutionEngineInitialize(MlirExecutionEngine jit) { + unwrap(jit)->initialize(); +} + extern "C" void mlirExecutionEngineDestroy(MlirExecutionEngine jit) { delete (unwrap(jit)); } @@ -106,9 +111,8 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, void *sym) { unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) { llvm::orc::SymbolMap symbolMap; - symbolMap[interner(unwrap(name))] = - { llvm::orc::ExecutorAddr::fromPtr(sym), - llvm::JITSymbolFlags::Exported }; + symbolMap[interner(unwrap(name))] = {llvm::orc::ExecutorAddr::fromPtr(sym), + llvm::JITSymbolFlags::Exported}; return symbolMap; }); } diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index f704fbfbe8fff..96fa60e15f425 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -106,7 +106,7 @@ void ExecutionEngine::dumpToObjectFile(StringRef filename) { } // Compilation is lazy and it doesn't populate object cache unless requested. // In case object dump is requested before cache is populated, we need to - // force compilation manually. + // force compilation manually. if (cache->isEmpty()) { for (std::string &functionName : functionNames) { auto result = lookupPacked(functionName); @@ -400,13 +400,9 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, return symbolMap; }; engine->registerSymbols(runtimeSymbolMap); - - // Execute the global constructors from the module being processed. - // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a - // crash for AArch64 see related issue #71963. - if (!engine->jit->getTargetTriple().isAArch64()) - cantFail(engine->jit->initialize(engine->jit->getMainJITDylib())); - + if (options.shouldInitialize) { + engine->initialize(); + } return std::move(engine); } @@ -451,3 +447,15 @@ Error ExecutionEngine::invokePacked(StringRef name, return Error::success(); } + +void ExecutionEngine::initialize() { + // Execute the global constructors from the module. + // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a + // crash for AArch64 see related issue #71963. + if (isInitialized) { + return; + } + if (!jit->getTargetTriple().isAArch64()) + cantFail(jit->initialize(jit->getMainJITDylib())); + isInitialized = true; +} diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c index 4751288c3ee4b..951558aef8c71 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, /*shouldInitialize=*/true); 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, /*shouldInitialize=*/true); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Engine creation failed with OpenMP"); exit(2); diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp index 887db227cfc4b..bd06f1ba15f25 100644 --- a/mlir/unittests/ExecutionEngine/Invoke.cpp +++ b/mlir/unittests/ExecutionEngine/Invoke.cpp @@ -316,4 +316,44 @@ TEST(NativeMemRefJit, MAYBE_JITCallback) { ASSERT_EQ(elt, coefficient * count++); } +static int initCnt = 0; +// A helper function that will be called during the JIT's initialization +static void initCallback() { initCnt += 1; } + +TEST(GlobalCtorJit, MAYBE_JITCallback) { + std::string moduleStr = R"mlir( + llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero] + llvm.func @ctor() { + func.call @init_callback() : () -> () + llvm.return + } + func.func private @init_callback() attributes { llvm.emit_c_interface } + )mlir"; + + DialectRegistry registry; + registerAllDialects(registry); + registerBuiltinDialectTranslation(registry); + registerLLVMDialectTranslation(registry); + MLIRContext context(registry); + auto module = parseSourceString(moduleStr, &context); + ASSERT_TRUE(!!module); + ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module))); + ExecutionEngineOptions jitOptions; + // defer the initialization to register symbols used in ctors + jitOptions.shouldInitialize = false; + auto jitOrError = ExecutionEngine::create(*module, jitOptions); + ASSERT_TRUE(!!jitOrError); + auto jit = std::move(jitOrError.get()); + // Define any extra symbols so they're available at initialization. + jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) { + llvm::orc::SymbolMap symbolMap; + symbolMap[interner("_mlir_ciface_init_callback")] = { + llvm::orc::ExecutorAddr::fromPtr(initCallback), + llvm::JITSymbolFlags::Exported}; + return symbolMap; + }); + jit->initialize(); + ASSERT_EQ(initCnt, 1); +} + #endif // _WIN32 From f48b3f4aea4363596330581928a08ffcb2a0d52e Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 15:25:05 +0800 Subject: [PATCH 02/15] wording --- mlir/unittests/ExecutionEngine/Invoke.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp index bd06f1ba15f25..f44fe76faff1b 100644 --- a/mlir/unittests/ExecutionEngine/Invoke.cpp +++ b/mlir/unittests/ExecutionEngine/Invoke.cpp @@ -317,7 +317,7 @@ TEST(NativeMemRefJit, MAYBE_JITCallback) { } static int initCnt = 0; -// A helper function that will be called during the JIT's initialization +// A helper function that will be called during the JIT's initialization. static void initCallback() { initCnt += 1; } TEST(GlobalCtorJit, MAYBE_JITCallback) { @@ -339,7 +339,7 @@ TEST(GlobalCtorJit, MAYBE_JITCallback) { ASSERT_TRUE(!!module); ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module))); ExecutionEngineOptions jitOptions; - // defer the initialization to register symbols used in ctors + // Defer the initialization to register symbols used in ctors. jitOptions.shouldInitialize = false; auto jitOrError = ExecutionEngine::create(*module, jitOptions); ASSERT_TRUE(!!jitOrError); From ee1a6dd468231bc5bdb40e548248383798cceea0 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 15:37:01 +0800 Subject: [PATCH 03/15] CAPI test --- mlir/test/CAPI/execution_engine.c | 55 +++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c index 951558aef8c71..c612545104e86 100644 --- a/mlir/test/CAPI/execution_engine.c +++ b/mlir/test/CAPI/execution_engine.c @@ -137,6 +137,60 @@ void testOmpCreation(void) { mlirContextDestroy(ctx); } +// Helper variable to track callback invocations +static int initCnt = 0; + +// Callback function that will be called during JIT initialization +static void initCallback(void) { initCnt += 1; } + +// CHECK-LABEL: Running test 'testGlobalCtorJitCallback' +void testGlobalCtorJitCallback(void) { + MlirContext ctx = mlirContextCreate(); + registerAllUpstreamDialects(ctx); + + // Create module with global constructor that calls our callback + MlirModule module = mlirModuleCreateParse( + ctx, mlirStringRefCreateFromCString( + // clang-format off +"module { \n" +" llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero] \n" +" llvm.func @ctor() { \n" +" func.call @init_callback() : () -> () \n" +" llvm.return \n" +" } \n" +" func.func private @init_callback() attributes { llvm.emit_c_interface } \n" +"} \n" + // clang-format on + )); + + lowerModuleToLLVM(ctx, module); + mlirRegisterAllLLVMTranslations(ctx); + + // Create execution engine with initialization disabled + MlirExecutionEngine jit = mlirExecutionEngineCreate( + module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL, + /*enableObjectDump=*/false, /*shouldInitialize=*/false); + + if (mlirExecutionEngineIsNull(jit)) { + fprintf(stderr, "Execution engine creation failed"); + exit(2); + } + + // Register callback symbol before initialization + mlirExecutionEngineRegisterSymbol( + jit, mlirStringRefCreateFromCString("_mlir_ciface_init_callback"), + (void *)initCallback); + + mlirExecutionEngineInitialize(jit); + + // CHECK: Init count: 1 + printf("Init count: %d\n", initCnt); + + mlirExecutionEngineDestroy(jit); + mlirModuleDestroy(module); + mlirContextDestroy(ctx); +} + int main(void) { #define _STRINGIFY(x) #x @@ -147,5 +201,6 @@ int main(void) { TEST(testSimpleExecution); TEST(testOmpCreation); + TEST(testGlobalCtorJitCallback); return 0; } From 5426bf4c24285882de1f3326278b85140e1b2a72 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 15:45:05 +0800 Subject: [PATCH 04/15] fix warning --- mlir/test/CAPI/execution_engine.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c index c612545104e86..1bf6d5d72be6e 100644 --- a/mlir/test/CAPI/execution_engine.c +++ b/mlir/test/CAPI/execution_engine.c @@ -179,7 +179,7 @@ void testGlobalCtorJitCallback(void) { // Register callback symbol before initialization mlirExecutionEngineRegisterSymbol( jit, mlirStringRefCreateFromCString("_mlir_ciface_init_callback"), - (void *)initCallback); + (void *)(uintptr_t)initCallback); mlirExecutionEngineInitialize(jit); From 3807a1775675276b0ba072e940a2cfa567bd2802 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 16:57:21 +0800 Subject: [PATCH 05/15] fmt --- mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp index 9d445bc4210ca..219e45a574574 100644 --- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp @@ -111,8 +111,9 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, void *sym) { unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) { llvm::orc::SymbolMap symbolMap; - symbolMap[interner(unwrap(name))] = {llvm::orc::ExecutorAddr::fromPtr(sym), - llvm::JITSymbolFlags::Exported}; + symbolMap[interner(unwrap(name))] = + { llvm::orc::ExecutorAddr::fromPtr(sym), + llvm::JITSymbolFlags::Exported }; return symbolMap; }); } From 88b93871b603f03c806df46443ca547d52d1f1f5 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 16:57:57 +0800 Subject: [PATCH 06/15] update docs --- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index 247446b842e05..49b5d11685f49 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -100,8 +100,7 @@ struct ExecutionEngineOptions { /// the llvm's global Perf notification listener. bool enablePerfNotificationListener = true; - /// Setting initialize=false to register symbols might be used in - /// initialization + /// Setting initialize=false to defer initialization bool shouldInitialize = true; }; From 7ce99afd9d20565bccd4a74ebf559e1c422703be Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 16:59:34 +0800 Subject: [PATCH 07/15] fmt --- mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp index 219e45a574574..624858c10eb7c 100644 --- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp @@ -112,8 +112,8 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) { llvm::orc::SymbolMap symbolMap; symbolMap[interner(unwrap(name))] = - { llvm::orc::ExecutorAddr::fromPtr(sym), - llvm::JITSymbolFlags::Exported }; + { llvm::orc::ExecutorAddr::fromPtr(sym), + llvm::JITSymbolFlags::Exported }; return symbolMap; }); } From bf7a0fde91996558863ecab82c448593c3a27555 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 17:01:27 +0800 Subject: [PATCH 08/15] refine comments --- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index 96fa60e15f425..7862cda3203c8 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -449,12 +449,11 @@ Error ExecutionEngine::invokePacked(StringRef name, } void ExecutionEngine::initialize() { - // Execute the global constructors from the module. - // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a - // crash for AArch64 see related issue #71963. 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())); isInitialized = true; From ca35dde48f0c08a72a0c1d2f807342128dd44f9c Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 17:47:13 +0800 Subject: [PATCH 09/15] fix python binding --- .../Bindings/Python/ExecutionEngineModule.cpp | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp index 81dada3553622..2157b152d210d 100644 --- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp +++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp @@ -7,8 +7,8 @@ //===----------------------------------------------------------------------===// #include "mlir-c/ExecutionEngine.h" -#include "mlir/Bindings/Python/NanobindAdaptors.h" #include "mlir/Bindings/Python/Nanobind.h" +#include "mlir/Bindings/Python/NanobindAdaptors.h" namespace nb = nanobind; using namespace mlir; @@ -75,13 +75,13 @@ NB_MODULE(_mlirExecutionEngine, m) { "__init__", [](PyExecutionEngine &self, MlirModule module, int optLevel, const std::vector &sharedLibPaths, - bool enableObjectDump) { + bool enableObjectDump, bool shouldInitialize) { 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, shouldInitialize); if (mlirExecutionEngineIsNull(executionEngine)) throw std::runtime_error( "Failure while creating the ExecutionEngine."); @@ -90,6 +90,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("should_initialize") = true, "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 " @@ -124,6 +125,12 @@ NB_MODULE(_mlirExecutionEngine, m) { }, nb::arg("name"), nb::arg("callback"), "Register `callback` as the runtime symbol `name`.") + .def( + "initialize", + [](PyExecutionEngine &executionEngine) { + mlirExecutionEngineInitialize(executionEngine.get()); + }, + "Initialize the ExecutionEngine.") .def( "dump_to_object_file", [](PyExecutionEngine &executionEngine, const std::string &fileName) { From fb0039ca05617b6821c545c90db6b23adc6bc5d0 Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Wed, 13 Aug 2025 18:08:24 +0800 Subject: [PATCH 10/15] Update mlir/lib/ExecutionEngine/ExecutionEngine.cpp Co-authored-by: Mehdi Amini --- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index 7862cda3203c8..fb4a1a5b37d83 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -449,9 +449,8 @@ Error ExecutionEngine::invokePacked(StringRef name, } void ExecutionEngine::initialize() { - if (isInitialized) { + 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()) From 1e3920c9a38be30f6159792749a0a7445879ac46 Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Wed, 13 Aug 2025 18:08:31 +0800 Subject: [PATCH 11/15] Update mlir/lib/ExecutionEngine/ExecutionEngine.cpp Co-authored-by: Mehdi Amini --- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index fb4a1a5b37d83..f4b8624ad3ad6 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -400,9 +400,8 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, return symbolMap; }; engine->registerSymbols(runtimeSymbolMap); - if (options.shouldInitialize) { + if (options.shouldInitialize) engine->initialize(); - } return std::move(engine); } From d90bd2933aaa04953d44cd7f34c723535c823aeb Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 18:18:00 +0800 Subject: [PATCH 12/15] more docs --- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index 49b5d11685f49..a3abfa60f5b62 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -230,6 +230,11 @@ class ExecutionEngine { llvm::function_ref symbolMap); + /// Initialize the ExecutionEngine. Global constructors specified by + /// `llvm.mlir.global_ctors` will be run. One common scenario is that kernel + /// binary compiled from `gpu.module` gets loaded during initialization. Make + /// sure all symbols are solvable before initialization by calling + /// `registerSymbols` or including shared libraries. void initialize(); private: From 28d752cbfde439932eea7c6e09c7b67b0e823621 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 18:21:26 +0800 Subject: [PATCH 13/15] fix typo --- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index a3abfa60f5b62..5efd87081bb29 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -233,7 +233,7 @@ class ExecutionEngine { /// Initialize the ExecutionEngine. Global constructors specified by /// `llvm.mlir.global_ctors` will be run. One common scenario is that kernel /// binary compiled from `gpu.module` gets loaded during initialization. Make - /// sure all symbols are solvable before initialization by calling + /// sure all symbols are resolvable before initialization by calling /// `registerSymbols` or including shared libraries. void initialize(); From 128bf5ddf700fcde4474b5b14a87ca8487fe15c1 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 21:01:07 +0800 Subject: [PATCH 14/15] address reviews --- mlir/include/mlir-c/ExecutionEngine.h | 7 +++---- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h | 3 --- mlir/lib/Bindings/Python/ExecutionEngineModule.cpp | 9 ++++----- mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp | 8 +++----- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 3 +-- mlir/lib/ExecutionEngine/JitRunner.cpp | 2 ++ mlir/test/CAPI/execution_engine.c | 6 +++--- mlir/unittests/ExecutionEngine/Invoke.cpp | 2 -- 8 files changed, 16 insertions(+), 24 deletions(-) diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h index eee56ee032c97..ac1eff7c94f99 100644 --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -42,10 +42,9 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void); /// that will be loaded are specified via `numPaths` and `sharedLibPaths` /// respectively. /// TODO: figure out other options. -MLIR_CAPI_EXPORTED MlirExecutionEngine -mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, - const MlirStringRef *sharedLibPaths, - bool enableObjectDump, bool shouldInitialize); +MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate( + MlirModule op, int optLevel, int numPaths, + const MlirStringRef *sharedLibPaths, bool enableObjectDump); /// Execute the global constructors from the module. MLIR_CAPI_EXPORTED void mlirExecutionEngineInitialize(MlirExecutionEngine jit); diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index 5efd87081bb29..5bd71d68d253a 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -99,9 +99,6 @@ struct ExecutionEngineOptions { /// If `enablePerfNotificationListener` is set, the JIT compiler will notify /// the llvm's global Perf notification listener. bool enablePerfNotificationListener = true; - - /// Setting initialize=false to defer initialization - bool shouldInitialize = true; }; /// JIT-backed execution engine for MLIR. Assumes the IR can be converted to diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp index 2157b152d210d..e93108cb11fc7 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 shouldInitialize) { + bool enableObjectDump) { 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, shouldInitialize); + MlirExecutionEngine executionEngine = + mlirExecutionEngineCreate(module, optLevel, libPaths.size(), + libPaths.data(), enableObjectDump); if (mlirExecutionEngineIsNull(executionEngine)) throw std::runtime_error( "Failure while creating the ExecutionEngine."); @@ -90,7 +90,6 @@ 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("should_initialize") = true, "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 624858c10eb7c..2dbb993b1640f 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 shouldInitialize) { + bool enableObjectDump) { static bool initOnce = [] { llvm::InitializeNativeTarget(); llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm @@ -60,7 +60,6 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, jitOptions.jitCodeGenOptLevel = static_cast(optLevel); jitOptions.sharedLibPaths = libPaths; jitOptions.enableObjectDump = enableObjectDump; - jitOptions.shouldInitialize = shouldInitialize; auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions); if (!jitOrError) { consumeError(jitOrError.takeError()); @@ -111,9 +110,8 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, void *sym) { unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) { llvm::orc::SymbolMap symbolMap; - symbolMap[interner(unwrap(name))] = - { llvm::orc::ExecutorAddr::fromPtr(sym), - llvm::JITSymbolFlags::Exported }; + symbolMap[interner(unwrap(name))] = {llvm::orc::ExecutorAddr::fromPtr(sym), + llvm::JITSymbolFlags::Exported}; return symbolMap; }); } diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index f4b8624ad3ad6..52162a43aeae3 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -400,8 +400,6 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, return symbolMap; }; engine->registerSymbols(runtimeSymbolMap); - if (options.shouldInitialize) - engine->initialize(); return std::move(engine); } @@ -437,6 +435,7 @@ Expected ExecutionEngine::lookup(StringRef name) const { Error ExecutionEngine::invokePacked(StringRef name, MutableArrayRef args) { + initialize(); auto expectedFPtr = lookupPacked(name); if (!expectedFPtr) return expectedFPtr.takeError(); diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp index 2107df37d1997..0ada4cc96570a 100644 --- a/mlir/lib/ExecutionEngine/JitRunner.cpp +++ b/mlir/lib/ExecutionEngine/JitRunner.cpp @@ -202,6 +202,8 @@ compileAndExecute(Options &options, Operation *module, StringRef entryPoint, auto engine = std::move(*expectedEngine); + engine->initialize(); + auto expectedFPtr = engine->lookupPacked(entryPoint); if (!expectedFPtr) return expectedFPtr.takeError(); diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c index 1bf6d5d72be6e..900588aeea2da 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, /*shouldInitialize=*/true); + /*enableObjectDump=*/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, /*shouldInitialize=*/true); + /*enableObjectDump=*/false); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Engine creation failed with OpenMP"); exit(2); @@ -169,7 +169,7 @@ void testGlobalCtorJitCallback(void) { // Create execution engine with initialization disabled MlirExecutionEngine jit = mlirExecutionEngineCreate( module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL, - /*enableObjectDump=*/false, /*shouldInitialize=*/false); + /*enableObjectDump=*/false); if (mlirExecutionEngineIsNull(jit)) { fprintf(stderr, "Execution engine creation failed"); diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp index f44fe76faff1b..e2f2fe0215b91 100644 --- a/mlir/unittests/ExecutionEngine/Invoke.cpp +++ b/mlir/unittests/ExecutionEngine/Invoke.cpp @@ -339,8 +339,6 @@ TEST(GlobalCtorJit, MAYBE_JITCallback) { ASSERT_TRUE(!!module); ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module))); ExecutionEngineOptions jitOptions; - // Defer the initialization to register symbols used in ctors. - jitOptions.shouldInitialize = false; auto jitOrError = ExecutionEngine::create(*module, jitOptions); ASSERT_TRUE(!!jitOrError); auto jit = std::move(jitOrError.get()); From 49a8ba0d406df0d28a905ea4dbf4ea2d626bc674 Mon Sep 17 00:00:00 2001 From: jackalcooper Date: Wed, 13 Aug 2025 21:06:44 +0800 Subject: [PATCH 15/15] more detailed docs for py and c --- mlir/include/mlir-c/ExecutionEngine.h | 6 +++++- mlir/lib/Bindings/Python/ExecutionEngineModule.cpp | 7 ++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h index ac1eff7c94f99..1a58d68533f24 100644 --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -46,7 +46,11 @@ MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate( MlirModule op, int optLevel, int numPaths, const MlirStringRef *sharedLibPaths, bool enableObjectDump); -/// Execute the global constructors from the module. +/// Initialize the ExecutionEngine. Global constructors specified by +/// `llvm.mlir.global_ctors` will be run. One common scenario is that kernel +/// binary compiled from `gpu.module` gets loaded during initialization. Make +/// sure all symbols are resolvable before initialization by calling +/// `mlirExecutionEngineRegisterSymbol` or including shared libraries. MLIR_CAPI_EXPORTED void mlirExecutionEngineInitialize(MlirExecutionEngine jit); /// Destroy an ExecutionEngine instance. diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp index e93108cb11fc7..4f7a4a628e246 100644 --- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp +++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp @@ -129,7 +129,12 @@ NB_MODULE(_mlirExecutionEngine, m) { [](PyExecutionEngine &executionEngine) { mlirExecutionEngineInitialize(executionEngine.get()); }, - "Initialize the ExecutionEngine.") + "Initialize the ExecutionEngine. Global constructors specified by " + "`llvm.mlir.global_ctors` will be run. One common scenario is that " + "kernel binary compiled from `gpu.module` gets loaded during " + "initialization. Make sure all symbols are resolvable before " + "initialization by calling `raw_register_runtime` or including " + "shared libraries.") .def( "dump_to_object_file", [](PyExecutionEngine &executionEngine, const std::string &fileName) {